Merge branch 'main' into text2igm

app-controlnet.py

@@ -35,6 +35,7 @@ import psutil
 MAX_QUEUE_SIZE = int(os.environ.get("MAX_QUEUE_SIZE", 0))
 TIMEOUT = float(os.environ.get("TIMEOUT", 0))
 SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None)
+TORCH_COMPILE = os.environ.get("TORCH_COMPILE", None)   
 WIDTH = 512
 HEIGHT = 512
 # disable tiny autoencoder for better quality speed tradeoff
@@ -100,15 +101,20 @@ pipe.unet.to(memory_format=torch.channels_last)
 if psutil.virtual_memory().total < 64 * 1024**3:
     pipe.enable_attention_slicing()
 
-if not mps_available and not xpu_available:
-    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-    pipe(prompt="warmup", image=[Image.new("RGB", (768, 768))], control_image=[Image.new("RGB", (768, 768))])
-
 compel_proc = Compel(
     tokenizer=pipe.tokenizer,
     text_encoder=pipe.text_encoder,
     truncate_long_prompts=False,
 )
+if TORCH_COMPILE:
+    pipe.text_encoder = torch.compile(pipe.text_encoder, mode="max-autotune", fullgraph=False)
+    pipe.tokenizer = torch.compile(pipe.tokenizer, mode="max-autotune", fullgraph=False)
+    pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=False)
+    pipe.vae = torch.compile(pipe.vae, mode="max-autotune", fullgraph=False)
+
+    pipe(prompt="warmup", image=[Image.new("RGB", (768, 768))], control_image=[Image.new("RGB", (768, 768))])
+
+
 user_queue_map = {}
 
 

app-img2img.py

@@ -29,6 +29,8 @@ import psutil
 MAX_QUEUE_SIZE = int(os.environ.get("MAX_QUEUE_SIZE", 0))
 TIMEOUT = float(os.environ.get("TIMEOUT", 0))
 SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None)
+TORCH_COMPILE = os.environ.get("TORCH_COMPILE", None)   
+
 WIDTH = 512
 HEIGHT = 512
 # disable tiny autoencoder for better quality speed tradeoff
@@ -76,8 +78,12 @@ pipe.unet.to(memory_format=torch.channels_last)
 if psutil.virtual_memory().total < 64 * 1024**3:
     pipe.enable_attention_slicing()
 
-if not mps_available and not xpu_available:
-    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+if TORCH_COMPILE:
+    pipe.text_encoder = torch.compile(pipe.text_encoder, mode="max-autotune", fullgraph=False)
+    pipe.tokenizer = torch.compile(pipe.tokenizer, mode="max-autotune", fullgraph=False)
+    pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=False)
+    pipe.vae = torch.compile(pipe.vae, mode="max-autotune", fullgraph=False)
+
     pipe(prompt="warmup", image=[Image.new("RGB", (512, 512))])
 
 compel_proc = Compel(

app-txt2img.py

@@ -30,6 +30,8 @@ import psutil
 MAX_QUEUE_SIZE = int(os.environ.get("MAX_QUEUE_SIZE", 0))
 TIMEOUT = float(os.environ.get("TIMEOUT", 0))
 SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None)
+TORCH_COMPILE = os.environ.get("TORCH_COMPILE", None)   
+
 WIDTH = 768
 HEIGHT = 768
 # disable tiny autoencoder for better quality speed tradeoff
@@ -76,8 +78,12 @@ pipe.unet.to(memory_format=torch.channels_last)
 if psutil.virtual_memory().total < 64 * 1024**3:
     pipe.enable_attention_slicing()
 
-if not mps_available and not xpu_available:
-    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+if TORCH_COMPILE:
+    pipe.text_encoder = torch.compile(pipe.text_encoder, mode="max-autotune", fullgraph=False)
+    pipe.tokenizer = torch.compile(pipe.tokenizer, mode="max-autotune", fullgraph=False)
+    pipe.unet = torch.compile(pipe.unet, mode="max-autotune", fullgraph=False)
+    pipe.vae = torch.compile(pipe.vae, mode="max-autotune", fullgraph=False)
+
     pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)
 
 compel_proc = Compel(

latent_consistency_controlnet.py

@@ -25,7 +25,6 @@ from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
 
 from diffusers import (
     AutoencoderKL,
-    AutoencoderTiny,
     ConfigMixin,
     DiffusionPipeline,
     SchedulerMixin,
@@ -50,6 +49,17 @@ import PIL.Image
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(encoder_output, generator):
+    if hasattr(encoder_output, "latent_dist"):
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
 class LatentConsistencyModelPipeline_controlnet(DiffusionPipeline):
     _optional_components = ["scheduler"]
 
@@ -276,22 +286,17 @@ class LatentConsistencyModelPipeline_controlnet(DiffusionPipeline):
                 )
 
             elif isinstance(generator, list):
-                if isinstance(self.vae, AutoencoderTiny):
-                    init_latents = [
-                        self.vae.encode(image[i : i + 1]).latents
-                        for i in range(batch_size)
-                    ]
-                else:
-                    init_latents = [
-                        self.vae.encode(image[i : i + 1]).latent_dist.sample(generator[i])
-                        for i in range(batch_size)
-                    ]
+                init_latents = [
+                    retrieve_latents(
+                        self.vae.encode(image[i : i + 1]), generator=generator[i]
+                    )
+                    for i in range(batch_size)
+                ]
                 init_latents = torch.cat(init_latents, dim=0)
             else:
-                if isinstance(self.vae, AutoencoderTiny):
-                    init_latents = self.vae.encode(image).latents
-                else:
-                    init_latents = self.vae.encode(image).latent_dist.sample(generator)
+                init_latents = retrieve_latents(
+                    self.vae.encode(image), generator=generator
+                )
 
             init_latents = self.vae.config.scaling_factor * init_latents
 

requirements.txt

@@ -1,4 +1,4 @@
-diffusers==0.22.1
+diffusers==0.22.2
 transformers==4.34.1
 gradio==3.50.2
 --extra-index-url https://download.pytorch.org/whl/cu121