fix(misc): fix failing tests, pipeline init in rare circumstances

kiri-art · Sep 4, 2023 · 9338648 · 9338648
1 parent 3f1f980
commit 9338648
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 8 deletions.
diff --git a/api/app.py b/api/app.py
@@ -16,7 +16,12 @@
 import skimage
 import skimage.measure
 from getScheduler import getScheduler, SCHEDULERS
-from getPipeline import getPipelineForModel, listAvailablePipelines, clearPipelines
+from getPipeline import (
+    getPipelineClass,
+    getPipelineForModel,
+    listAvailablePipelines,
+    clearPipelines,
+)
 import re
 import requests
 from download import download_model, normalize_model_id
@@ -228,7 +233,7 @@ def sendStatus():
         model_dir = os.path.join(MODELS_DIR, normalized_model_id)
         pipeline_name = call_inputs.get("PIPELINE", None)
         if pipeline_name:
-            pipeline_class = getattr(diffusers_pipelines, pipeline_name)
+            pipeline_class = getPipelineClass(pipeline_name)
         if last_model_id != normalized_model_id:
             # if not downloaded_models.get(normalized_model_id, None):
             if not os.path.isdir(model_dir):
@@ -250,7 +255,7 @@ def sendStatus():
                     hf_model_id=hf_model_id,
                     model_precision=model_precision,
                     send_opts=send_opts,
-                    pipeline_class=pipeline_class,
+                    pipeline_class=pipeline_class if pipeline_name else None,
                 )
                 # downloaded_models.update({normalized_model_id: True})
             clearPipelines()
@@ -267,7 +272,7 @@ def sendStatus():
                 precision=model_precision,
                 revision=model_revision,
                 send_opts=send_opts,
-                pipeline_class=pipeline_class,
+                pipeline_class=pipeline_class if pipeline_name else None,
             )
             await send(
                 "loadModel", "done", {"startRequestId": startRequestId}, send_opts

diff --git a/api/getPipeline.py b/api/getPipeline.py
@@ -48,6 +48,13 @@ def clearPipelines():
     _pipelines = {}
 
 
+def getPipelineClass(pipeline_name: str):
+    if hasattr(diffusers_pipelines, pipeline_name):
+        return getattr(diffusers_pipelines, pipeline_name)
+    elif pipeline_name in availableCommunityPipelines():
+        return DiffusionPipeline
+
+
 def getPipelineForModel(
     pipeline_name: str, model, model_id, model_revision, model_precision
 ):

diff --git a/api/loadModel.py b/api/loadModel.py
@@ -31,7 +31,7 @@ def loadModel(
     precision=None,
     revision=None,
     send_opts={},
-    pipeline_class=AutoPipelineForText2Image,
+    pipeline_class=None,
 ):
     torch_dtype = torch_dtype_from_precision(precision)
     if revision == "":
@@ -44,18 +44,23 @@ def loadModel(
             "load": load,
             "precision": precision,
             "revision": revision,
+            "pipeline_class": pipeline_class,
         },
     )
+
+    if not pipeline_class:
+        pipeline_class = AutoPipelineForText2Image
+
+    pipeline = pipeline_class if PIPELINE == "ALL" else getattr(_pipelines, PIPELINE)
+    print("pipeline", pipeline_class)
+
     print(
         ("Loading" if load else "Downloading")
         + " model: "
         + model_id
         + (f" ({revision})" if revision else "")
     )
 
-    pipeline = pipeline_class if PIPELINE == "ALL" else getattr(_pipelines, PIPELINE)
-    print("pipeline", pipeline_class)
-
     scheduler = getScheduler(model_id, DEFAULT_SCHEDULER, not load)
 
     model_dir = os.path.join(MODELS_DIR, model_id)