add support for other pipelines too.

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import torch
 from diffusers import (
+    AutoPipelineForText2Image,
     StableDiffusionXLControlNetPipeline,
     DiffusionPipeline,
     StableDiffusionImg2ImgPipeline,
@@ -55,6 +56,8 @@ pipeline_mapping = {
         "stabilityai/stable-diffusion-xl-base-1.0",
         "TencentARC/t2i-adapter-canny-sdxl-1.0",
     ),
+    "Kandinsky 2.2 (T2I)": (AutoPipelineForText2Image, "kandinsky-community/kandinsky-2-2-decoder"),
+    "Würstchen (T2I)": (AutoPipelineForText2Image, "warp-ai/wuerstchen")
 }
 
 
@@ -93,22 +96,36 @@ def load_pipeline(
         pipeline = pipeline_cls.from_pretrained(pipeline_ckpt, controlnet=controlnet)
     elif "Adapters" in pipeline_to_benchmark:
         pipeline = pipeline_cls.from_pretrained(pipeline_ckpt, adapter=adapter)
+    
     pipeline.to(device)
 
     # Optionally set memory layout.
     if use_channels_last:
-        pipeline.unet.to(memory_format=torch.channels_last)
+        if pipeline_to_benchmark not in ["Würstchen (T2I)", "Kandinsky 2.2 (T2I)"]:
+            pipeline.unet.to(memory_format=torch.channels_last)
+        elif pipeline_to_benchmark == "Würstchen (T2I)": 
+            pipeline.prior.to(memory_format=torch.channels_last)
+            pipeline.decoder.to(memory_format=torch.channels_last)
+        elif pipeline_to_benchmark == "Kandinsky 2.2 (T2I)":
+            pipeline.unet.to(memory_format=torch.channels_last)
 
-    if hasattr(pipeline, "controlnet"):
-        pipeline.controlnet.to(memory_format=torch.channels_last)
-    elif hasattr(pipeline, "adapter"):
-        pipeline.adapter.to(memory_format=torch.channels_last)
+        if hasattr(pipeline, "controlnet"):
+            pipeline.controlnet.to(memory_format=torch.channels_last)
+        elif hasattr(pipeline, "adapter"):
+            pipeline.adapter.to(memory_format=torch.channels_last)
 
     # Optional torch compilation.
     if do_torch_compile:
-        pipeline.unet = torch.compile(
-            pipeline.unet, mode="reduce-overhead", fullgraph=True
-        )
+        if pipeline_to_benchmark not in ["Würstchen (T2I)", "Kandinsky 2.2 (T2I)"]:
+            pipeline.unet = torch.compile(
+                pipeline.unet, mode="reduce-overhead", fullgraph=True
+            )
+        elif pipeline_to_benchmark == "Würstchen (T2I)":
+            pipeline.prior = torch.compile(pipeline.prior, mode="reduce-overhead", fullgraph=True)
+            pipeline.decoder = torch.compile(pipeline.decoder, mode="reduce-overhead", fullgraph=True)
+        elif  pipeline_to_benchmark == "Kandinsky 2.2 (T2I)":
+            pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
+        
         if hasattr(pipeline, "controlnet"):
             pipeline.controlnet = torch.compile(
                 pipeline.controlnet, mode="reduce-overhead", fullgraph=True

utils.py

@@ -5,8 +5,7 @@ def get_image_for_img_to_img(pipeline_to_benchmark):
     if pipeline_to_benchmark == "SD I2I":
         url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
         init_image = load_image(url).convert("RGB")
-        size = (768, 512)
-        init_image = init_image.resize(size)
+        init_image = init_image.resize((512, 512))
     elif pipeline_to_benchmark == "SDXL I2I":
         url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl-img2img.png"
         init_image = load_image(url).convert("RGB")