feat(sdxl,compel): Support. AutoPipeline default, safety_check fix

Author: gadicc

api/app.py

@@ -28,7 +28,10 @@
 from threading import Timer
 import extras
 
+from diffusers import StableDiffusionXLPipeline
+
 from lib.textual_inversions import handle_textual_inversions
+from lib.prompts import prepare_prompts
 from lib.vars import (
     RUNTIME_DOWNLOADS,
     USE_DREAMBOOTH,
@@ -290,7 +293,7 @@ def sendStatus():
     if PIPELINE == "ALL":
         pipeline_name = call_inputs.get("PIPELINE", None)
         if not pipeline_name:
-            pipeline_name = "StableDiffusionPipeline"
+            pipeline_name = "AutoPipelineForText2Image"
             result["$meta"].update({"PIPELINE": pipeline_name})
 
         pipeline = getPipelineForModel(
@@ -329,7 +332,11 @@ def sendStatus():
         }
 
     safety_checker = call_inputs.get("safety_checker", True)
-    pipeline.safety_checker = model.safety_checker if safety_checker else None
+    pipeline.safety_checker = (
+        model.safety_checker
+        if safety_checker and hasattr(model, "safety_checker")
+        else None
+    )
     is_url = call_inputs.get("is_url", False)
     image_decoder = getFromUrl if is_url else decodeBase64Image
 
@@ -399,6 +406,8 @@ def sendStatus():
                         else 1
                     )
                     cross_attention_kwargs.update({"scale": storage_query_scale})
+                    # https://github.com/damian0815/compel/issues/42#issuecomment-1656989385
+                    pipeline._lora_scale = storage_query_scale
                     if storage_query_fname:
                         fname = storage_query_fname[0]
                     else:
@@ -569,8 +578,22 @@ def callback(step: int, timestep: int, latents: torch.FloatTensor):
                 "inference", step / model_inputs.get("num_inference_steps", 50)
             )
 
+    is_sdxl = isinstance(model, StableDiffusionXLPipeline)
+    print("is_sdxl", is_sdxl)
+
     with torch.inference_mode():
         custom_pipeline_method = call_inputs.get("custom_pipeline_method", None)
+        print(
+            pipeline,
+            {
+                "cross_attention_kwargs": cross_attention_kwargs,
+                "callback": callback,
+                "**model_inputs": model_inputs,
+            },
+        )
+
+        if call_inputs.get("compel_prompts", False):
+            prepare_prompts(pipeline, model_inputs, is_sdxl)
 
         try:
             async_pipeline = asyncio.to_thread(
@@ -581,13 +604,13 @@ def callback(step: int, timestep: int, latents: torch.FloatTensor):
                 callback=callback,
                 **model_inputs,
             )
-            if call_inputs.get("PIPELINE") != "StableDiffusionPipeline":
-                # autocast im2img and inpaint which are broken in 0.4.0, 0.4.1
-                # still broken in 0.5.1
-                with autocast(device_id):
-                    images = (await async_pipeline).images
-            else:
-                images = (await async_pipeline).images
+            # if call_inputs.get("PIPELINE") != "StableDiffusionPipeline":
+            #    # autocast im2img and inpaint which are broken in 0.4.0, 0.4.1
+            #    # still broken in 0.5.1
+            #    with autocast(device_id):
+            #        images = (await async_pipeline).images
+            # else:
+            images = (await async_pipeline).images
 
         except Exception as err:
             return {

api/getPipeline.py

@@ -65,8 +65,11 @@ def getPipelineForModel(
     start = time.time()
 
     if hasattr(diffusers_pipelines, pipeline_name):
-        if hasattr(model, "components"):
-            pipeline = getattr(diffusers_pipelines, pipeline_name)(**model.components)
+        pipeline_class = getattr(diffusers_pipelines, pipeline_name)
+        if hasattr(pipeline_class, "from_pipe"):
+            pipeline = pipeline_class.from_pipe(model)
+        elif hasattr(model, "components"):
+            pipeline = pipeline_class(**model.components)
         else:
             pipeline = getattr(diffusers_pipelines, pipeline_name)(
                 vae=model.vae,

api/lib/prompts.py

@@ -0,0 +1,60 @@
+from compel import Compel, DiffusersTextualInversionManager, ReturnedEmbeddingsType
+
+
+def prepare_prompts(pipeline, model_inputs, is_sdxl):
+    textual_inversion_manager = DiffusersTextualInversionManager(pipeline)
+    if is_sdxl:
+        compel = Compel(
+            tokenizer=[pipeline.tokenizer, pipeline.tokenizer_2],
+            text_encoder=[pipeline.text_encoder, pipeline.text_encoder_2],
+            # diffusers has no ti in sdxl yet
+            # https://github.com/huggingface/diffusers/issues/4376#issuecomment-1659016141
+            # textual_inversion_manager=textual_inversion_manager,
+            truncate_long_prompts=False,
+            returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+            requires_pooled=[False, True],
+        )
+        conditioning, pooled = compel(model_inputs.get("prompt"))
+        negative_conditioning, negative_pooled = compel(
+            model_inputs.get("negative_prompt")
+        )
+        [
+            conditioning,
+            negative_conditioning,
+        ] = compel.pad_conditioning_tensors_to_same_length(
+            [conditioning, negative_conditioning]
+        )
+        model_inputs.update(
+            {
+                "prompt": None,
+                "negative_prompt": None,
+                "prompt_embeds": conditioning,
+                "negative_prompt_embeds": negative_conditioning,
+                "pooled_prompt_embeds": pooled,
+                "negative_pooled_prompt_embeds": negative_pooled,
+            }
+        )
+
+    else:
+        compel = Compel(
+            tokenizer=pipeline.tokenizer,
+            text_encoder=pipeline.text_encoder,
+            textual_inversion_manager=textual_inversion_manager,
+            truncate_long_prompts=False,
+        )
+        conditioning = compel(model_inputs.get("prompt"))
+        negative_conditioning = compel(model_inputs.get("negative_prompt"))
+        [
+            conditioning,
+            negative_conditioning,
+        ] = compel.pad_conditioning_tensors_to_same_length(
+            [conditioning, negative_conditioning]
+        )
+        model_inputs.update(
+            {
+                "prompt": None,
+                "negative_prompt": None,
+                "prompt_embeds": conditioning,
+                "negative_prompt_embeds": negative_conditioning,
+            }
+        )

api/loadModel.py

@@ -1,6 +1,6 @@
 import torch
 import os
-from diffusers import pipelines as _pipelines, StableDiffusionPipeline
+from diffusers import pipelines as _pipelines, AutoPipelineForText2Image
 from getScheduler import getScheduler, DEFAULT_SCHEDULER
 from precision import torch_dtype_from_precision
 from device import device
@@ -25,7 +25,14 @@
 ]
 
 
-def loadModel(model_id: str, load=True, precision=None, revision=None, send_opts={}):
+def loadModel(
+    model_id: str,
+    load=True,
+    precision=None,
+    revision=None,
+    send_opts={},
+    pipeline_class=AutoPipelineForText2Image,
+):
     torch_dtype = torch_dtype_from_precision(precision)
     if revision == "":
         revision = None
@@ -46,9 +53,7 @@ def loadModel(model_id: str, load=True, precision=None, revision=None, send_opts
         + (f" ({revision})" if revision else "")
     )
 
-    pipeline = (
-        StableDiffusionPipeline if PIPELINE == "ALL" else getattr(_pipelines, PIPELINE)
-    )
+    pipeline = pipeline_class if PIPELINE == "ALL" else getattr(_pipelines, PIPELINE)
 
     scheduler = getScheduler(model_id, DEFAULT_SCHEDULER, not load)
 

requirements.txt

@@ -62,3 +62,4 @@ xtarfile[zstd]==0.1.0
 bitsandbytes==0.40.2 # released 2023-07-17
 
 invisible-watermark==0.2.0 # released 2023-07-06
+compel==2.0.1 # released 2023-07-29