Spaces:
Running
Running
| import gradio as gr | |
| import os, shutil | |
| from typing import Any | |
| # ---------------------------------- | |
| # 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") | |
| examples = [ | |
| ["Stella the dog running in a park"], | |
| ["Stella the dog on a beach at sunset"], | |
| ] | |
| # ---------------------------------- | |
| # 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) | |
| demo = gr.Interface( | |
| fn=generate_output, | |
| inputs=gr.Textbox(label="Input"), | |
| outputs=gr.Image(label="Output"), | |
| examples=examples, | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |