kandisnky 2.2

dkackman · Jul 27, 2023 · 8b6679d · 8b6679d
1 parent e6dd833
commit 8b6679d
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 86 deletions.
diff --git a/swarm/diffusion/diffusion_func.py b/swarm/diffusion/diffusion_func.py
@@ -19,7 +19,7 @@ def diffusion_callback(device_identifier, model_name, **kwargs):
     lora = kwargs.pop("lora", None)
     cross_attention_scale = kwargs.pop("cross_attention_scale", 1.0)
 
-    # set output_type if already there or upscale is selected (we use the latent upscaler)
+    # set output_type if already there or upscale is/ selected (we use the latent upscaler)
     output_type = kwargs.pop("output_type", "latent" if upscale else None)
     if output_type is not None:
         kwargs["output_type"] = output_type
@@ -43,14 +43,22 @@ def diffusion_callback(device_identifier, model_name, **kwargs):
                 "preprocessed_input", [kwargs.get("control_image")]
             )
 
-    pipeline = pipeline_type.from_pretrained(
-        model_name,
-        revision=kwargs.pop("revision", "main"),
-        variant=kwargs.pop("variant", None),
-        torch_dtype=torch.float16,
-        controlnet=controlnet if "controlnet" in locals() else None,
-        use_safe_tensors=use_safe_tensors,
-    )
+        pipeline = pipeline_type.from_pretrained(
+            model_name,
+            revision=kwargs.pop("revision", "main"),
+            variant=kwargs.pop("variant", None),
+            torch_dtype=torch.float16,
+            controlnet=controlnet if "controlnet" in locals() else None,
+            use_safe_tensors=use_safe_tensors,
+        )
+    else:
+        pipeline = pipeline_type.from_pretrained(
+            model_name,
+            revision=kwargs.pop("revision", "main"),
+            variant=kwargs.pop("variant", None),
+            torch_dtype=torch.float16,
+            use_safe_tensors=use_safe_tensors,
+        )
 
     if textual_inversion is not None:
         try:

diff --git a/swarm/diffusion/kandinsky.py b/swarm/diffusion/kandinsky.py
diff --git a/swarm/job_arguments.py b/swarm/job_arguments.py
@@ -6,7 +6,6 @@
 from .audio.audioldm import txt2audio_diffusion_callback
 from .audio.bark import bark_diffusion_callback
 from .diffusion.diffusion_func_if import diffusion_if_callback
-from .diffusion.kandinsky import kandinsky_callback
 from .type_helpers import get_type
 from .pre_processors.controlnet import scale_to_size
 from .external_resources import get_image, get_control_image, max_size, download_images
@@ -37,9 +36,6 @@ async def format_args(job):
     if args["model_name"].startswith("DeepFloyd/"):
         return diffusion_if_callback, args
 
-    if args["model_name"].startswith("kandinsky-"):
-        return await format_kandinsky_args(args)
-
     return await format_stable_diffusion_args(args, workflow)