use taesd for all models

pipelines/controlnet.py

@@ -16,6 +16,7 @@ import psutil
 from config import Args
 from pydantic import BaseModel, Field
 from PIL import Image
+import math
 
 base_model = "SimianLuo/LCM_Dreamshaper_v7"
 taesd_model = "madebyollin/taesd"
@@ -68,13 +69,13 @@ class Pipeline:
             2159232, min=0, title="Seed", field="seed", hide=True, id="seed"
         )
         steps: int = Field(
-            4, min=2, max=15, title="Steps", field="range", hide=True, id="steps"
+            4, min=1, max=15, title="Steps", field="range", hide=True, id="steps"
         )
         width: int = Field(
-            512, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
+            768, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
         )
         height: int = Field(
-            512, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
+            768, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
         )
         guidance_scale: float = Field(
             0.2,
@@ -171,7 +172,7 @@ class Pipeline:
         if args.use_taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
-            )
+            ).to(device)
         self.canny_torch = SobelOperator(device=device)
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)
@@ -208,14 +209,18 @@ class Pipeline:
         control_image = self.canny_torch(
             params.image, params.canny_low_threshold, params.canny_high_threshold
         )
+        steps = params.steps
+        strength = params.strength
+        if int(steps * strength) < 1:
+            steps = math.ceil(1 / max(0.10, strength))
 
         results = self.pipe(
             image=params.image,
             control_image=control_image,
             prompt_embeds=prompt_embeds,
             generator=generator,
-            strength=params.strength,
-            num_inference_steps=params.steps,
+            strength=strength,
+            num_inference_steps=steps,
             guidance_scale=params.guidance_scale,
             width=params.width,
             height=params.height,

pipelines/controlnetLoraSD15.py

@@ -2,6 +2,7 @@ from diffusers import (
     StableDiffusionControlNetImg2ImgPipeline,
     ControlNetModel,
     LCMScheduler,
+    AutoencoderTiny,
 )
 from compel import Compel
 import torch
@@ -16,6 +17,7 @@ import psutil
 from config import Args
 from pydantic import BaseModel, Field
 from PIL import Image
+import math
 
 taesd_model = "madebyollin/taesd"
 controlnet_model = "lllyasviel/control_v11p_sd15_canny"
@@ -79,13 +81,13 @@ class Pipeline:
             2159232, min=0, title="Seed", field="seed", hide=True, id="seed"
         )
         steps: int = Field(
-            4, min=2, max=15, title="Steps", field="range", hide=True, id="steps"
+            4, min=1, max=15, title="Steps", field="range", hide=True, id="steps"
         )
         width: int = Field(
-            512, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
+            768, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
         )
         height: int = Field(
-            512, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
+            768, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
         )
         guidance_scale: float = Field(
             0.2,
@@ -200,6 +202,11 @@ class Pipeline:
             if psutil.virtual_memory().total < 64 * 1024**3:
                 pipe.enable_attention_slicing()
 
+            if args.use_taesd:
+                pipe.vae = AutoencoderTiny.from_pretrained(
+                    taesd_model, torch_dtype=torch_dtype, use_safetensors=True
+                ).to(device)
+
             # Load LCM LoRA
             pipe.load_lora_weights(lcm_lora_id, adapter_name="lcm")
             pipe.compel_proc = Compel(
@@ -222,7 +229,6 @@ class Pipeline:
 
     def predict(self, params: "Pipeline.InputParams") -> Image.Image:
         generator = torch.manual_seed(params.seed)
-        print(f"Using model: {params.base_model_id}")
         pipe = self.pipes[params.base_model_id]
 
         activation_token = base_models[params.base_model_id]
@@ -231,14 +237,18 @@ class Pipeline:
         control_image = self.canny_torch(
             params.image, params.canny_low_threshold, params.canny_high_threshold
         )
+        steps = params.steps
+        strength = params.strength
+        if int(steps * strength) < 1:
+            steps = math.ceil(1 / max(0.10, strength))
 
         results = pipe(
             image=params.image,
             control_image=control_image,
             prompt_embeds=prompt_embeds,
             generator=generator,
-            strength=params.strength,
-            num_inference_steps=params.steps,
+            strength=strength,
+            num_inference_steps=steps,
             guidance_scale=params.guidance_scale,
             width=params.width,
             height=params.height,

pipelines/controlnetLoraSDXL.py

@@ -3,6 +3,7 @@ from diffusers import (
     ControlNetModel,
     LCMScheduler,
     AutoencoderKL,
+    AutoencoderTiny,
 )
 from compel import Compel, ReturnedEmbeddingsType
 import torch
@@ -17,10 +18,12 @@ import psutil
 from config import Args
 from pydantic import BaseModel, Field
 from PIL import Image
+import math
 
 controlnet_model = "diffusers/controlnet-canny-sdxl-1.0"
 model_id = "stabilityai/stable-diffusion-xl-base-1.0"
 lcm_lora_id = "latent-consistency/lcm-lora-sdxl"
+taesd_model = "madebyollin/taesdxl"
 
 
 default_prompt = "Portrait of The Terminator with , glare pose, detailed, intricate, full of colour, cinematic lighting, trending on artstation, 8k, hyperrealistic, focused, extreme details, unreal engine 5 cinematic, masterpiece"
@@ -77,7 +80,7 @@ class Pipeline:
             2159232, min=0, title="Seed", field="seed", hide=True, id="seed"
         )
         steps: int = Field(
-            4, min=2, max=15, title="Steps", field="range", hide=True, id="steps"
+            2, min=1, max=15, title="Steps", field="range", hide=True, id="steps"
         )
         width: int = Field(
             1024, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
@@ -96,10 +99,10 @@ class Pipeline:
             id="guidance_scale",
         )
         strength: float = Field(
-            0.5,
+            1,
             min=0.25,
             max=1.0,
-            step=0.001,
+            step=0.0001,
             title="Strength",
             field="range",
             hide=True,
@@ -208,6 +211,10 @@ class Pipeline:
             returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
             requires_pooled=[False, True],
         )
+        if args.use_taesd:
+            self.pipe.vae = AutoencoderTiny.from_pretrained(
+                taesd_model, torch_dtype=torch_dtype, use_safetensors=True
+            ).to(device)
 
         if args.torch_compile:
             self.pipe.unet = torch.compile(
@@ -231,6 +238,10 @@ class Pipeline:
         control_image = self.canny_torch(
             params.image, params.canny_low_threshold, params.canny_high_threshold
         )
+        steps = params.steps
+        strength = params.strength
+        if int(steps * strength) < 1:
+            steps = math.ceil(1 / max(0.10, strength))
 
         results = self.pipe(
             image=params.image,
@@ -240,8 +251,8 @@ class Pipeline:
             negative_prompt_embeds=prompt_embeds[1:2],
             negative_pooled_prompt_embeds=pooled_prompt_embeds[1:2],
             generator=generator,
-            strength=params.strength,
-            num_inference_steps=params.steps,
+            strength=strength,
+            num_inference_steps=steps,
             guidance_scale=params.guidance_scale,
             width=params.width,
             height=params.height,

pipelines/controlnetSDXLTurbo.py

@@ -2,6 +2,7 @@ from diffusers import (
     StableDiffusionXLControlNetImg2ImgPipeline,
     ControlNetModel,
     AutoencoderKL,
+    AutoencoderTiny,
 )
 from compel import Compel, ReturnedEmbeddingsType
 import torch
@@ -20,6 +21,7 @@ import math
 
 controlnet_model = "diffusers/controlnet-canny-sdxl-1.0"
 model_id = "stabilityai/sdxl-turbo"
+taesd_model = "madebyollin/taesdxl"
 
 default_prompt = "Portrait of The Terminator with , glare pose, detailed, intricate, full of colour, cinematic lighting, trending on artstation, 8k, hyperrealistic, focused, extreme details, unreal engine 5 cinematic, masterpiece"
 default_negative_prompt = "blurry, low quality, render, 3D, oversaturated"
@@ -75,18 +77,18 @@ class Pipeline:
             2159232, min=0, title="Seed", field="seed", hide=True, id="seed"
         )
         steps: int = Field(
-            4, min=1, max=15, title="Steps", field="range", hide=True, id="steps"
+            2, min=1, max=15, title="Steps", field="range", hide=True, id="steps"
         )
         width: int = Field(
-            512, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
+            1024, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
         )
         height: int = Field(
-            512, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
+            1024, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
         )
         guidance_scale: float = Field(
             1.0,
             min=0,
-            max=20,
+            max=10,
             step=0.001,
             title="Guidance Scale",
             field="range",
@@ -197,6 +199,10 @@ class Pipeline:
             returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
             requires_pooled=[False, True],
         )
+        if args.use_taesd:
+            self.pipe.vae = AutoencoderTiny.from_pretrained(
+                taesd_model, torch_dtype=torch_dtype, use_safetensors=True
+            ).to(device)
 
         if args.torch_compile:
             self.pipe.unet = torch.compile(

pipelines/img2img.py

@@ -14,6 +14,7 @@ import psutil
 from config import Args
 from pydantic import BaseModel, Field
 from PIL import Image
+import math
 
 base_model = "SimianLuo/LCM_Dreamshaper_v7"
 taesd_model = "madebyollin/taesd"
@@ -64,13 +65,13 @@ class Pipeline:
             2159232, min=0, title="Seed", field="seed", hide=True, id="seed"
         )
         steps: int = Field(
-            4, min=2, max=15, title="Steps", field="range", hide=True, id="steps"
+            4, min=1, max=15, title="Steps", field="range", hide=True, id="steps"
         )
         width: int = Field(
-            512, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
+            768, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
         )
         height: int = Field(
-            512, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
+            768, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
         )
         guidance_scale: float = Field(
             0.2,
@@ -104,7 +105,7 @@ class Pipeline:
         if args.use_taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
-            )
+            ).to(device)
 
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)
@@ -138,12 +139,18 @@ class Pipeline:
     def predict(self, params: "Pipeline.InputParams") -> Image.Image:
         generator = torch.manual_seed(params.seed)
         prompt_embeds = self.compel_proc(params.prompt)
+
+        steps = params.steps
+        strength = params.strength
+        if int(steps * strength) < 1:
+            steps = math.ceil(1 / max(0.10, strength))
+
         results = self.pipe(
             image=params.image,
             prompt_embeds=prompt_embeds,
             generator=generator,
-            strength=params.strength,
-            num_inference_steps=params.steps,
+            strength=strength,
+            num_inference_steps=steps,
             guidance_scale=params.guidance_scale,
             width=params.width,
             height=params.height,

pipelines/img2imgSDXLTurbo.py

@@ -17,7 +17,7 @@ from PIL import Image
 import math
 
 base_model = "stabilityai/sdxl-turbo"
-taesd_model = "madebyollin/taesd"
+taesd_model = "madebyollin/taesdxl"
 
 default_prompt = "close-up photography of old man standing in the rain at night, in a street lit by lamps, leica 35mm summilux"
 default_negative_prompt = "blurry, low quality, render, 3D, oversaturated"
@@ -113,7 +113,7 @@ class Pipeline:
         if args.use_taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
-            )
+            ).to(device)
 
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)

pipelines/txt2img.py

@@ -62,10 +62,10 @@ class Pipeline:
             4, min=2, max=15, title="Steps", field="range", hide=True, id="steps"
         )
         width: int = Field(
-            512, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
+            768, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
         )
         height: int = Field(
-            512, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
+            768, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
         )
         guidance_scale: float = Field(
             8.0,
@@ -88,7 +88,7 @@ class Pipeline:
         if args.use_taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
-            )
+            ).to(device)
 
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)

pipelines/txt2imgLora.py

@@ -95,7 +95,7 @@ class Pipeline:
         if args.use_taesd:
             self.pipe.vae = AutoencoderTiny.from_pretrained(
                 taesd_model, torch_dtype=torch_dtype, use_safetensors=True
-            )
+            ).to(device)
         self.pipe.scheduler = LCMScheduler.from_config(self.pipe.scheduler.config)
         self.pipe.set_progress_bar_config(disable=True)
         self.pipe.to(device=device, dtype=torch_dtype)

pipelines/txt2imgLoraSDXL.py

@@ -1,8 +1,4 @@
-from diffusers import (
-    DiffusionPipeline,
-    LCMScheduler,
-    AutoencoderKL,
-)
+from diffusers import DiffusionPipeline, LCMScheduler, AutoencoderKL, AutoencoderTiny
 from compel import Compel, ReturnedEmbeddingsType
 import torch
 
@@ -16,9 +12,9 @@ from config import Args
 from pydantic import BaseModel, Field
 from PIL import Image
 
-controlnet_model = "diffusers/controlnet-canny-sdxl-1.0"
 model_id = "stabilityai/stable-diffusion-xl-base-1.0"
 lcm_lora_id = "latent-consistency/lcm-lora-sdxl"
+taesd_model = "madebyollin/taesdxl"
 
 
 default_prompt = "close-up photography of old man standing in the rain at night, in a street lit by lamps, leica 35mm summilux"
@@ -76,7 +72,7 @@ class Pipeline:
             2159232, min=0, title="Seed", field="seed", hide=True, id="seed"
         )
         steps: int = Field(
-            4, min=2, max=15, title="Steps", field="range", hide=True, id="steps"
+            4, min=1, max=15, title="Steps", field="range", hide=True, id="steps"
         )
         width: int = Field(
             1024, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
@@ -127,6 +123,10 @@ class Pipeline:
             returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
             requires_pooled=[False, True],
         )
+        if args.use_taesd:
+            self.pipe.vae = AutoencoderTiny.from_pretrained(
+                taesd_model, torch_dtype=torch_dtype, use_safetensors=True
+            ).to(device)
 
         if args.torch_compile:
             self.pipe.unet = torch.compile(