Update app.py

app.py

@@ -1,23 +1,49 @@
 import gradio as gr
+import os, shutil
+from typing import Any
 
-# Load the model
+# ----------------------------------
+#  Clean out stale Gradio cache
+# ----------------------------------
+if os.path.isdir(".gradio/cached_examples"):
+    shutil.rmtree(".gradio/cached_examples")
+
+# ----------------------------------
+#  Load the model
+# ----------------------------------
 model = gr.load("models/jbigs/stellathedog")
 
-# Define custom examples
 examples = [
     ["Stella the dog running in a park"],
     ["Stella the dog on a beach at sunset"],
 ]
 
-# Function to call model and return only the image part
-def generate_output(input_text):
-    result = model(input_text)
-    return result[0]  # Extract just the image
+# ----------------------------------
+#  Utility to pull an image‑like object
+# ----------------------------------
+def _extract_image(x: Any):
+    if isinstance(x, (tuple, list)):        # first element from tuple/list
+        return _extract_image(x[0])
+    if isinstance(x, dict):                 # common keys or first value
+        for k in ("image", "img", "picture", "output", "result"):
+            if k in x:
+                return _extract_image(x[k])
+        return _extract_image(next(iter(x.values())))
+    return x                                # assume already Gradio‑compatible
+
+# ----------------------------------
+#  Interface callback
+# ----------------------------------
+def generate_output(prompt: str):
+    raw = model(prompt)
+    return _extract_image(raw)
 
-# Create the interface
-gr.Interface(
+demo = gr.Interface(
     fn=generate_output,
     inputs=gr.Textbox(label="Input"),
     outputs=gr.Image(label="Output"),
     examples=examples,
-).launch()
+)
+
+if __name__ == "__main__":
+    demo.launch()