:tada: init

app.py

@@ -306,16 +306,40 @@ def inference(image: Image.Image, prompt: str, progress=gr.Progress(track_tqdm=T
 
 
 if __name__ == '__main__':
-    ca_output = [gr.Image(type="pil", label="Attention Map") for _ in range(16)]
+    unet_mapping = [
+        "0: Down 64",
+        "1: Down 64",
+        "2: Down 32",
+        "3: Down 32",
+        "4: Down 16",
+        "5: Down 16",
+        "6: Mid 8",
+        "7: Up 16",
+        "8: Up 16",
+        "9: Up 16",
+        "10: Up 32",
+        "11: Up 32",
+        "12: Up 32",
+        "13: Up 64",
+        "14: Up 64",
+        "15: Up 64",
+    ]
+    ca_output = [gr.Image(type="pil", label=unet_mapping[i]) for i in range(16)]
 
     iface = gr.Interface(
         title="Stable Diffusion Attention Visualizer",
-        description="",
+        description="This is a visualizer for the attention maps of the Stable Diffusion model. ",
         fn=inference,
         inputs=[
-            gr.Image(type="pil", label="Input Image", width=512, height=512),
-            gr.Textbox(label="Prompt", placeholder="Enter a prompt here..."),
+            gr.Image(type="pil", label="Input", width=512, height=512),
+            gr.Textbox(label="Prompt", placeholder="e.g.) A photo of dog...")
         ],
         outputs=ca_output,
+        cache_examples=True,
+        examples=[
+            [Image.open("assets/aeroplane.png"), "plane, background"],
+            [Image.open("assets/dogcat.png"), "a photo of dog and cat"],
+        ]
     )
+
     iface.launch()