feat(pipelines): allow calling of ALL PIPELINES (official+community)

gadicc · gadicc · commit 1ccbaad1f405 · 2022-12-08T13:50:09.000Z
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,6 +6,12 @@
     all pipelines.  Using `init_image` now shoes a deprecation warning and
     will be removed in future.
 
+  * **ALL THE PIPELINES**.  We no longer load a list of hard-coded pipelines
+    in `init()`.  Instead, we init and cache each on first use (for faster
+    first calls on cold boots), and, *all* pipelines, both official diffusers
+    and community pipelines, are available.
+    [Full details](https://banana-forums.dev/t/all-your-pipelines-are-belong-to-us/83)
+
   * **Changed `sd-base` to `diffusers-api` as the default tag / name used
     in the README examples and optional [./build][build script].
 
diff --git a/README.md b/README.md
@@ -7,8 +7,11 @@ Please give credit and link back to this repo if you use it in a public project.
 
 ## Features
 
-* Pipelines: txt2img, img2img and inpainting in a single container
 * Models: stable-diffusion, waifu-diffusion, and easy to add others (e.g. jp-sd)
+* Pipelines: txt2img, img2img and inpainting in a single container
+
+  (
+  [all diffusers official and community pipelines](https://banana-forums.dev/t/all-your-pipelines-are-belong-to-us/83) are wrapped, but untested)
 * All model inputs supported, including setting nsfw filter per request
 * *Permute* base config to multiple forks based on yaml config with vars
 * Optionally send signed event logs / performance data to a REST endpoint
diff --git a/app.py b/app.py
@@ -2,14 +2,7 @@
 import torch
 
 from torch import autocast
-from diffusers import (
-    pipelines as _pipelines,
-    LMSDiscreteScheduler,
-    DDIMScheduler,
-    PNDMScheduler,
-    DiffusionPipeline,
-    __version__,
-)
+from diffusers import __version__
 import base64
 from io import BytesIO
 import PIL
@@ -22,10 +15,10 @@
 import skimage.measure
 from PyPatchMatch import patch_match
 from getScheduler import getScheduler, SCHEDULERS
+from getPipeline import getPipelineForModel, listAvailablePipelines, clearPipelines
 import re
 import requests
 from download import download_model
-from precision import revision, torch_dtype
 
 RUNTIME_DOWNLOADS = os.getenv("RUNTIME_DOWNLOADS") == "1"
 USE_DREAMBOOTH = os.getenv("USE_DREAMBOOTH") == "1"
@@ -36,53 +29,9 @@
 PIPELINE = os.environ.get("PIPELINE")
 HF_AUTH_TOKEN = os.getenv("HF_AUTH_TOKEN")
 
-PIPELINES = [
-    "StableDiffusionPipeline",
-    "StableDiffusionImg2ImgPipeline",
-    "StableDiffusionInpaintPipeline",
-    "StableDiffusionInpaintPipelineLegacy",
-]
-
-COMMUNITY_PIPELINES = [
-    "lpw_stable_diffusion",
-]
-
 torch.set_grad_enabled(False)
 
 
-def createPipelinesFromModel(model, model_id):
-    pipelines = dict()
-
-    for pipeline in PIPELINES:
-        if hasattr(_pipelines, pipeline):
-            if hasattr(model, "components"):
-                pipelines[pipeline] = getattr(_pipelines, pipeline)(**model.components)
-            else:
-                pipelines[pipeline] = getattr(_pipelines, pipeline)(
-                    vae=model.vae,
-                    text_encoder=model.text_encoder,
-                    tokenizer=model.tokenizer,
-                    unet=model.unet,
-                    scheduler=model.scheduler,
-                    safety_checker=model.safety_checker,
-                    feature_extractor=model.feature_extractor,
-                )
-        else:
-            print(f'Skipping non-existent pipeline "{PIPELINE}"')
-
-    for pipeline in COMMUNITY_PIPELINES:
-        pipelines[pipeline] = DiffusionPipeline.from_pretrained(
-            model_id,
-            revision=revision,
-            torch_dtype=torch_dtype,
-            custom_pipeline="./diffusers/examples/community/" + pipeline + ".py",
-            local_files_only=True,
-            **model.components,
-        )
-
-    return pipelines
-
-
 class DummySafetyChecker:
     @staticmethod
     def __call__(images, clip_input):
@@ -93,8 +42,6 @@ def __call__(images, clip_input):
 # Load your model to GPU as a global variable here using the variable name "model"
 def init():
     global model  # needed for bananna optimizations
-    global pipelines
-    global schedulers
     global dummy_safety_checker
     global initTime
 
@@ -120,9 +67,6 @@ def init():
     if not RUNTIME_DOWNLOADS:
         model = loadModel(MODEL_ID)
 
-        if PIPELINE == "ALL":
-            pipelines = createPipelinesFromModel(model, MODEL_ID)
-
     send("init", "done")
     initTime = get_now() - initStart
 
@@ -200,13 +144,13 @@ def inference(all_inputs: dict) -> dict:
                 downloaded_models.update({model_id: True})
             model = loadModel(model_id)
             if PIPELINE == "ALL":
-                pipelines = createPipelinesFromModel(model, model_id)
+                clearPipelines()
             last_model_id = model_id
 
     if MODEL_ID == "ALL":
         if last_model_id != model_id:
             model = loadModel(model_id)
-            pipelines = createPipelinesFromModel(model, model_id)
+            clearPipelines()
             last_model_id = model_id
     else:
         if model_id != MODEL_ID and not RUNTIME_DOWNLOADS:
@@ -220,7 +164,17 @@ def inference(all_inputs: dict) -> dict:
             }
 
     if PIPELINE == "ALL":
-        pipeline = pipelines.get(call_inputs.get("PIPELINE"))
+        pipeline_name = call_inputs.get("PIPELINE")
+        pipeline = getPipelineForModel(pipeline_name, model, model_id)
+        if not pipeline:
+            return {
+                "$error": {
+                    "code": "NO_SUCH_PIPELINE",
+                    "message": f'"{pipeline_name}" is not an official nor community Diffusers pipelines',
+                    "requested": pipeline_name,
+                    "available": listAvailablePipelines(),
+                }
+            }
     else:
         pipeline = model
 
diff --git a/getPipeline.py b/getPipeline.py
@@ -0,0 +1,90 @@
+import time
+import os, fnmatch
+from diffusers import (
+    DiffusionPipeline,
+    pipelines as diffusers_pipelines,
+)
+from precision import revision, torch_dtype
+
+_pipelines = {}
+_availableCommunityPipelines = None
+
+
+def listAvailablePipelines():
+    return (
+        list(
+            filter(
+                lambda key: key.endswith("Pipeline"),
+                list(diffusers_pipelines.__dict__.keys()),
+            )
+        )
+        + availableCommunityPipelines()
+    )
+
+
+def availableCommunityPipelines():
+    global _availableCommunityPipelines
+    if not _availableCommunityPipelines:
+        _availableCommunityPipelines = list(
+            map(
+                lambda s: s[0:-3],
+                fnmatch.filter(os.listdir("diffusers/examples/community"), "*.py"),
+            )
+        )
+
+    return _availableCommunityPipelines
+
+
+def clearPipelines():
+    """
+    Clears the pipeline cache.  Important to call this when changing the
+    loaded model, as pipelines include references to the model and would
+    therefore prevent memory being reclaimed after unloading the previous
+    model.
+    """
+    pipelines = {}
+
+
+def getPipelineForModel(pipeline_name: str, model, model_id):
+    """
+    Inits a new pipeline, re-using components from a previously loaded
+    model.  The pipeline is cached and future calls with the same
+    arguments will return the previously initted instance.  Be sure
+    to call `clearPipelines()` if loading a new model, to allow the
+    previous model to be garbage collected.
+    """
+    pipeline = _pipelines.get(pipeline_name)
+    if pipeline:
+        return pipeline
+
+    start = time.time()
+
+    if hasattr(diffusers_pipelines, pipeline_name):
+        if hasattr(model, "components"):
+            pipeline = getattr(diffusers_pipelines, pipeline_name)(**model.components)
+        else:
+            pipeline = getattr(diffusers_pipelines, pipeline_name)(
+                vae=model.vae,
+                text_encoder=model.text_encoder,
+                tokenizer=model.tokenizer,
+                unet=model.unet,
+                scheduler=model.scheduler,
+                safety_checker=model.safety_checker,
+                feature_extractor=model.feature_extractor,
+            )
+
+    elif pipeline_name in availableCommunityPipelines():
+        pipeline = DiffusionPipeline.from_pretrained(
+            model_id,
+            revision=revision,
+            torch_dtype=torch_dtype,
+            custom_pipeline="./diffusers/examples/community/" + pipeline_name + ".py",
+            local_files_only=True,
+            **model.components,
+        )
+
+    if pipeline:
+        _pipelines.update({pipeline_name: pipeline})
+        diff = round((time.time() - start) * 1000)
+        print(f"Initialized {pipeline_name} for {model_id} in {diff}ms")
+        return pipeline
