more updates.

app.py

@@ -56,8 +56,11 @@ 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")
+    "Kandinsky 2.2 (T2I)": (
+        AutoPipelineForText2Image,
+        "kandinsky-community/kandinsky-2-2-decoder",
+    ),
+    "Würstchen (T2I)": (AutoPipelineForText2Image, "warp-ai/wuerstchen"),
 }
 
 
@@ -67,6 +70,7 @@ def load_pipeline(
     do_torch_compile: bool = False,
 ):
     # Get pipeline details.
+    print(f"Loading pipeline: {pipeline_to_benchmark}")
     pipeline_details = pipeline_mapping[pipeline_to_benchmark]
     pipeline_cls = pipeline_details[0]
     pipeline_ckpt = pipeline_details[1]
@@ -96,14 +100,15 @@ 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:
+        print("Setting memory layout.")
         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)": 
+        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)":
@@ -116,16 +121,23 @@ def load_pipeline(
 
     # Optional torch compilation.
     if do_torch_compile:
+        print("Compiling pipeline.")
         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)
-        
+            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
@@ -135,6 +147,7 @@ def load_pipeline(
                 pipeline.adapter, mode="reduce-overhead", fullgraph=True
             )
 
+    print("Pipeline loaded.")
     return pipeline
 
 
@@ -144,6 +157,11 @@ def generate(
     use_channels_last: bool = False,
     do_torch_compile: bool = False,
 ):
+    if isinstance(pipeline_to_benchmark, list):
+        # It can only happen when we don't select a pipeline to benchmark.
+        raise ValueError(
+            "pipeline_to_benchmark cannot be None. Please select a pipeline to benchmark."
+        )
     print("Start...")
     print("Torch version", torch.__version__)
     print("Torch CUDA version", torch.version.cuda)
@@ -189,7 +207,7 @@ with gr.Blocks() as demo:
     use_channels_last = gr.Checkbox(label="Use `channels_last` memory layout?")
     pipeline_to_benchmark = gr.Dropdown(
         list(pipeline_mapping.keys()),
-        value=["Stable Diffusion V1.5"],
+        value=None,
         multiselect=False,
         label="Pipeline to benchmark",
     )