From 45898e37276193c6ab8b74aa4ccff8cb3af444f3 Mon Sep 17 00:00:00 2001
From: Rick Staa <rick.staa@outlook.com>
Date: Fri, 31 May 2024 14:30:21 +0200
Subject: [PATCH] feat(runner): disable DEEPCACHE for lightning/turbo models

This commit ensures that people can not use the deepcache optimization
with lightning and turbo models. As explained in
https://github.com/livepeer/ai-worker/issues/82#issuecomment-2141983903
this optimization does not offer any speedup for these models while it
does reduce image quality.
---
 runner/app/pipelines/image_to_image.py | 20 +++++++++++++++++---
 runner/app/pipelines/text_to_image.py  | 12 ++++++++++--
 runner/app/pipelines/util.py           | 22 ++++++++++++++++++++++
 3 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/runner/app/pipelines/image_to_image.py b/runner/app/pipelines/image_to_image.py
index 9ddf687a..a2ec720f 100644
--- a/runner/app/pipelines/image_to_image.py
+++ b/runner/app/pipelines/image_to_image.py
@@ -1,5 +1,11 @@
 from app.pipelines.base import Pipeline
-from app.pipelines.util import get_torch_device, get_model_dir, SafetyChecker
+from app.pipelines.util import (
+    get_torch_device,
+    get_model_dir,
+    SafetyChecker,
+    is_lightning_model,
+    is_turbo_model,
+)
 
 from diffusers import (
     AutoPipelineForImage2Image,
@@ -119,14 +125,22 @@ def __init__(self, model_id: str):
                     "call may be slow if 'SFAST' is enabled."
                 )
 
-        if deepcache_enabled:
+        if deepcache_enabled and not (
+            is_lightning_model(model_id) or is_turbo_model(model_id)
+        ):
             logger.info(
-                "TextToImagePipeline will be optimized with DeepCache for %s",
+                "ImageToImagePipeline will be optimized with DeepCache for %s",
                 model_id,
             )
             from app.pipelines.optim.deepcache import enable_deepcache
 
             self.ldm = enable_deepcache(self.ldm)
+        elif deepcache_enabled:
+            logger.warning(
+                "DeepCache is not supported for Lightning or Turbo models. "
+                "ImageToImagePipeline will NOT be optimized with DeepCache for %s",
+                model_id,
+            )
 
         safety_checker_device = os.getenv("SAFETY_CHECKER_DEVICE", "cuda").lower()
         self._safety_checker = SafetyChecker(device=safety_checker_device)
diff --git a/runner/app/pipelines/text_to_image.py b/runner/app/pipelines/text_to_image.py
index db3ac6e8..eedf4132 100644
--- a/runner/app/pipelines/text_to_image.py
+++ b/runner/app/pipelines/text_to_image.py
@@ -14,7 +14,7 @@
 from safetensors.torch import load_file
 
 from app.pipelines.base import Pipeline
-from app.pipelines.util import get_model_dir, get_torch_device, SafetyChecker
+from app.pipelines.util import get_model_dir, get_torch_device, SafetyChecker, is_lightning_model, is_turbo_model
 
 logger = logging.getLogger(__name__)
 
@@ -139,7 +139,9 @@ def __init__(self, model_id: str):
                     "call may be slow if 'SFAST' is enabled."
                 )
 
-        if deepcache_enabled:
+        if deepcache_enabled and not (
+            is_lightning_model(model_id) or is_turbo_model(model_id)
+        ):
             logger.info(
                 "TextToImagePipeline will be optimized with DeepCache for %s",
                 model_id,
@@ -147,6 +149,12 @@ def __init__(self, model_id: str):
             from app.pipelines.optim.deepcache import enable_deepcache
 
             self.ldm = enable_deepcache(self.ldm)
+        elif deepcache_enabled:
+            logger.warning(
+                "DeepCache is not supported for Lightning or Turbo models. "
+                "TextToImagePipeline will NOT be optimized with DeepCache for %s",
+                model_id,
+            )
 
         safety_checker_device = os.getenv("SAFETY_CHECKER_DEVICE", "cuda").lower()
         self._safety_checker = SafetyChecker(device=safety_checker_device)
diff --git a/runner/app/pipelines/util.py b/runner/app/pipelines/util.py
index 2a73b798..e2f1c938 100644
--- a/runner/app/pipelines/util.py
+++ b/runner/app/pipelines/util.py
@@ -51,6 +51,28 @@ def validate_torch_device(device_name: str) -> bool:
         return False
 
 
+def is_lightning_model(model_id: str) -> bool:
+    """Checks if the model is a Lightning model.
+
+    Args:
+        model_id: Model ID.
+
+    Returns:
+        True if the model is a Lightning model, False otherwise.
+    """
+    return "-lightning" in model_id.lower()
+
+def is_turbo_model(model_id: str) -> bool:
+    """Checks if the model is a Turbo model.
+
+    Args:
+        model_id: Model ID.
+
+    Returns:
+        True if the model is a Turbo model, False otherwise.
+    """
+    return "-turbo" in model_id.lower()
+
 class SafetyChecker:
     """Checks images for unsafe or inappropriate content using a pretrained model.