Unify UI to single interface with optional image input; add generate defaults so examples work

app.py

@@ -1,7 +1,8 @@
 """Mage-Flow: Efficient Native-Resolution Foundation Model for Image Generation and Editing.
 
-Gradio Space demo supporting text-to-image generation and instruction-based image editing
-using Microsoft's Mage-Flow model (Turbo variants for fast inference).
+Gradio Space demo with a single unified interface: enter a prompt and optionally
+upload an image. If an image is provided, the request is routed to the
+Mage-Flow-Edit-Turbo model; otherwise it is routed to the Mage-Flow-Turbo model.
 """
 import os
 
@@ -25,27 +26,53 @@ pipe_edit = MageFlowPipeline.from_pretrained(EDIT_MODEL, device="cuda")
 @spaces.GPU(duration=60)
 def generate(
     prompt: str,
-    negative_prompt: str,
-    steps: int,
-    cfg: float,
-    height: int,
-    width: int,
-    seed: int,
+    image=None,
+    negative_prompt: str = " ",
+    steps: int = 4,
+    cfg: float = 1.0,
+    height: int = 1024,
+    width: int = 1024,
+    max_size: int = 1024,
+    seed: int = 42,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Generate an image from a text prompt using Mage-Flow-Turbo.
+    """Generate or edit an image with Mage-Flow.
+
+    If ``image`` is provided, the reference image is edited with
+    Mage-Flow-Edit-Turbo. Otherwise a new image is generated from the text
+    prompt with Mage-Flow-Turbo.
 
     Args:
-        prompt: Text description of the image to generate.
-        negative_prompt: What to avoid in the generation.
+        prompt: Text description (generation) or edit instruction (editing).
+        image: Optional reference image. When given, routes to the edit model.
+        negative_prompt: What to avoid in the result.
         steps: Number of denoising steps (Turbo uses 4).
         cfg: Classifier-free guidance scale (Turbo uses 1.0).
-        height: Output image height (multiple of 16).
-        width: Output image width (multiple of 16).
+        height: Output image height for text-to-image (multiple of 16).
+        width: Output image width for text-to-image (multiple of 16).
+        max_size: Longest side of edited output (0 = keep source resolution).
         seed: Random seed for reproducibility.
     """
     if not (prompt or "").strip():
         raise gr.Error("Prompt is empty.")
+
+    if image is not None:
+        # Route to the edit model when an image is provided.
+        if isinstance(image, str):
+            image = Image.open(image)
+        refs = [image.convert("RGB")]
+        out = pipe_edit.edit(
+            [prompt],
+            [refs],
+            neg_prompts=[negative_prompt or " "],
+            seeds=[int(seed)],
+            steps=int(steps),
+            cfg=float(cfg),
+            max_size=int(max_size) if max_size else None,
+        )[0]
+        return out
+
+    # No image: route to the text-to-image model.
     img = pipe_t2i.generate(
         [prompt],
         neg_prompts=[negative_prompt or " "],
@@ -58,152 +85,98 @@ def generate(
     return img
 
 
-@spaces.GPU(duration=60)
-def edit(
-    prompt: str,
-    negative_prompt: str,
-    ref_img,
-    steps: int,
-    cfg: float,
-    max_size: int,
-    seed: int,
-    progress=gr.Progress(track_tqdm=True),
-):
-    """Edit a reference image using an instruction prompt with Mage-Flow-Edit-Turbo.
-
-    Args:
-        prompt: Edit instruction describing how to modify the image.
-        negative_prompt: What to avoid in the edited result.
-        ref_img: Reference image to edit.
-        steps: Number of denoising steps (Turbo uses 4).
-        cfg: Classifier-free guidance scale.
-        max_size: Longest side of output (0 = keep source resolution).
-        seed: Random seed for reproducibility.
-    """
-    if not (prompt or "").strip():
-        raise gr.Error("Edit instruction is empty.")
-    if ref_img is None:
-        raise gr.Error("Upload a reference image to edit.")
-    if isinstance(ref_img, str):
-        ref_img = Image.open(ref_img)
-    refs = [ref_img.convert("RGB")]
-    out = pipe_edit.edit(
-        [prompt],
-        [refs],
-        neg_prompts=[negative_prompt or " "],
-        seeds=[int(seed)],
-        steps=int(steps),
-        cfg=float(cfg),
-        max_size=int(max_size) if max_size else None,
-    )[0]
-    return out
-
-
 ASSETS_DIR = os.path.join(os.path.dirname(__file__), "mage_flow", "assets")
 
 CSS = """
-#col-container { max-width: 1100px; margin: 0 auto; }
+#col-container { margin: 0 auto; max-width: 1100px; }
 .dark .gradio-container { color: var(--body-text-color); }
 """
 
 with gr.Blocks(css=CSS) as demo:
-    gr.Markdown(
-        "# Mage-Flow\n"
-        "Efficient Native-Resolution Foundation Model for Image Generation and Editing.\n\n"
-        "Model: [microsoft/Mage-Flow-Turbo](https://huggingface.co/microsoft/Mage-Flow-Turbo) | "
-        "[Paper](https://huggingface.co/papers/2607.19064) | "
-        "[GitHub](https://github.com/microsoft/Mage)"
-    )
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            "# Mage-Flow\n"
+            "Efficient Native-Resolution Foundation Model for Image Generation and Editing. "
+            "Enter a prompt to generate an image, or upload an image to edit it.\n\n"
+            "Models: [Mage-Flow-Turbo](https://huggingface.co/microsoft/Mage-Flow-Turbo) & "
+            "[Mage-Flow-Edit-Turbo](https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo) | "
+            "[Paper](https://huggingface.co/papers/2607.19064) | "
+            "[GitHub](https://github.com/microsoft/Mage)"
+        )
 
-    with gr.Tab("Text → Image"):
         with gr.Row():
             with gr.Column(scale=1):
-                t_prompt = gr.Textbox(
-                    label="Prompt",
-                    lines=3,
-                    value="A close-up portrait of an elderly African man with deep wrinkles, "
-                          "wearing a traditional hat, soft natural lighting, ultra realistic.",
-                )
-                t_neg = gr.Textbox(label="Negative prompt", value=" ", lines=1)
-                with gr.Row():
-                    t_steps = gr.Slider(1, 50, value=4, step=1, label="Steps")
-                    t_cfg = gr.Slider(1.0, 10.0, value=1.0, step=0.5, label="CFG")
                 with gr.Row():
-                    t_h = gr.Slider(256, 1536, value=1024, step=16, label="Height")
-                    t_w = gr.Slider(256, 1536, value=1024, step=16, label="Width")
-                t_seed = gr.Number(value=42, precision=0, label="Seed")
-                t_btn = gr.Button("Generate", variant="primary")
+                    prompt = gr.Textbox(
+                        label="Prompt",
+                        show_label=False,
+                        max_lines=3,
+                        placeholder="Describe an image to generate, or an edit instruction for an uploaded image",
+                        container=False,
+                        scale=4,
+                    )
+                    run_btn = gr.Button("Run", variant="primary", scale=1)
+
+                with gr.Accordion("Input image (optional — enables editing)", open=True):
+                    image = gr.Image(
+                        type="pil",
+                        label="Input image",
+                        show_label=False,
+                        height=300,
+                    )
+
+                with gr.Accordion("Advanced Settings", open=False):
+                    negative_prompt = gr.Textbox(label="Negative prompt", value=" ", lines=1)
+                    with gr.Row():
+                        steps = gr.Slider(1, 50, value=4, step=1, label="Steps")
+                        cfg = gr.Slider(1.0, 10.0, value=1.0, step=0.5, label="CFG")
+                    with gr.Row():
+                        height = gr.Slider(256, 1536, value=1024, step=16, label="Height (text→image)")
+                        width = gr.Slider(256, 1536, value=1024, step=16, label="Width (text→image)")
+                    max_size = gr.Slider(
+                        0, 1536, value=1024, step=16,
+                        label="Max output side for editing (0 = keep source size)",
+                    )
+                    seed = gr.Number(value=42, precision=0, label="Seed")
+
             with gr.Column(scale=1):
-                t_out = gr.Image(type="pil", label="Output", height=560)
+                result = gr.Image(type="pil", label="Output", height=560)
 
+        gr.Markdown("### Text → Image examples")
         gr.Examples(
             examples=[
                 ["A close-up portrait of an elderly African man with deep wrinkles, wearing a traditional hat, soft natural lighting, ultra realistic."],
                 ["A serene mountain landscape at sunset, with snow-capped peaks reflecting golden light, photorealistic."],
                 ["A cute robot playing a guitar in a neon-lit cyberpunk city, digital art style."],
             ],
-            inputs=[t_prompt],
-            outputs=t_out,
+            inputs=[prompt],
+            outputs=result,
             fn=generate,
             cache_examples=True,
             cache_mode="lazy",
         )
 
-        t_btn.click(
-            lambda: None, None, t_out
-        ).then(
-            generate,
-            [t_prompt, t_neg, t_steps, t_cfg, t_h, t_w, t_seed],
-            t_out,
-            api_name="generate",
-        )
-
-    with gr.Tab("Image Edit"):
-        with gr.Row():
-            with gr.Column(scale=1):
-                e_prompt = gr.Textbox(
-                    label="Edit instruction",
-                    lines=2,
-                    value="change the background to a city street",
-                )
-                e_neg = gr.Textbox(label="Negative prompt", value=" ", lines=1)
-                e_ref = gr.Image(
-                    type="pil", label="Reference image", height=280,
-                    value=os.path.join(ASSETS_DIR, "dog.jpg"),
-                )
-                with gr.Row():
-                    e_steps = gr.Slider(1, 50, value=4, step=1, label="Steps")
-                    e_cfg = gr.Slider(1.0, 10.0, value=1.0, step=0.5, label="CFG")
-                e_max = gr.Slider(
-                    0, 1536, value=1024, step=16,
-                    label="Max output side (0 = keep source size)",
-                )
-                e_seed = gr.Number(value=42, precision=0, label="Seed")
-                e_btn = gr.Button("Edit", variant="primary")
-            with gr.Column(scale=1):
-                e_out = gr.Image(type="pil", label="Output", height=560)
-
+        gr.Markdown("### Image editing examples")
         gr.Examples(
             examples=[
                 ["change the background to a city street", os.path.join(ASSETS_DIR, "dog.jpg")],
                 ["make it look like a painting", os.path.join(ASSETS_DIR, "cuisine.jpg")],
                 ["add a hat to the person", os.path.join(ASSETS_DIR, "portrait.jpg")],
             ],
-            inputs=[e_prompt, e_ref],
-            outputs=e_out,
-            fn=edit,
+            inputs=[prompt, image],
+            outputs=result,
+            fn=generate,
             cache_examples=True,
             cache_mode="lazy",
         )
 
-        e_btn.click(
-            lambda: None, None, e_out
-        ).then(
-            edit,
-            [e_prompt, e_neg, e_ref, e_steps, e_cfg, e_max, e_seed],
-            e_out,
-            api_name="edit",
-        )
+    inputs = [prompt, image, negative_prompt, steps, cfg, height, width, max_size, seed]
+    run_btn.click(lambda: None, None, result).then(
+        generate, inputs, result, api_name="generate",
+    )
+    prompt.submit(lambda: None, None, result).then(
+        generate, inputs, result, api_name=False,
+    )
 
 if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Citrus(), mcp_server=True, show_error=True)
+    demo.launch(theme=gr.themes.Citrus(), mcp_server=True, show_error=True)