Update app.py

app.py

@@ -208,36 +208,35 @@ def infer(images_b64_json, prompt, seed, randomize_seed, guidance_scale, steps,
         torch.cuda.empty_cache()
 
 
-# ─────────────────────────────────────────────
-#  Clean API endpoint — usable via Gradio API
-# ─────────────────────────────────────────────
+# ─────────────────────────────────────────
+#  API endpoint — /api/predict or /api/edit
+# ─────────────────────────────────────────
 @spaces.GPU
 def api_edit(
-    image: Image.Image,
+    image: "Image.Image",
     prompt: str,
     seed: int = 0,
     randomize_seed: bool = True,
     guidance_scale: float = 1.0,
     steps: int = 4,
-    extra_image: Image.Image = None,
+    extra_image: "Image.Image" = None,
 ) -> tuple:
     """
-    Edit an image using the FireRed / Qwen image editing pipeline.
+    Edit an image with a text prompt.
 
     Parameters
     ----------
-    image           : The main input image (PIL Image).
-    prompt          : Edit instruction, e.g. 'Convert to anime style'.
-    seed            : Random seed (ignored when randomize_seed=True).
-    randomize_seed  : If True a random seed is used each call.
-    guidance_scale  : CFG scale (1.0–10.0). Keep at 1.0 for the rapid model.
-    steps           : Inference steps (default 4 for the rapid/distilled model).
-    extra_image     : Optional second reference image (e.g. clothing reference).
+    image          : Input PIL image.
+    prompt         : Edit instruction e.g. 'Convert to anime style'.
+    seed           : Seed (ignored when randomize_seed=True).
+    randomize_seed : Use a random seed each call.
+    guidance_scale : CFG scale 1.0–10.0 (keep 1.0 for rapid model).
+    steps          : Inference steps (default 4).
+    extra_image    : Optional second reference image.
 
     Returns
     -------
-    result_image    : Edited PIL Image.
-    seed_used       : The seed that was actually used.
+    (result_image, seed_used)
     """
     import base64, json
     from io import BytesIO
@@ -247,19 +246,14 @@ def api_edit(
     if not prompt or not prompt.strip():
         raise gr.Error("prompt is required.")
 
-    # Pack images into the b64 JSON format that infer() expects
-    pil_list = [image]
-    if extra_image is not None:
-        pil_list.append(extra_image)
-
+    pil_list = [image] + ([extra_image] if extra_image is not None else [])
     b64_list = []
     for img in pil_list:
         buf = BytesIO()
         img.save(buf, format="PNG")
         b64_list.append("data:image/png;base64," + base64.b64encode(buf.getvalue()).decode())
 
-    images_b64_json = json.dumps(b64_list)
-    return infer(images_b64_json, prompt, seed, randomize_seed, guidance_scale, steps)
+    return infer(json.dumps(b64_list), prompt, seed, randomize_seed, guidance_scale, steps)
 
 
 css = r"""
@@ -1202,43 +1196,31 @@ with gr.Blocks() as demo:
         queue=False,
     )
 
-# ── Dedicated API tab (separate from the custom HTML UI) ──
-with gr.Blocks() as api_demo:
-    gr.Markdown("## Image Edit API")
-    gr.Markdown(
-        "Send requests to `/api/edit` programmatically. "
-        "See the [API docs](/api) for the full schema."
-    )
-    with gr.Row():
-        with gr.Column():
-            api_image_in   = gr.Image(type="pil", label="Input Image")
-            api_extra_in   = gr.Image(type="pil", label="Reference Image (optional)")
-            api_prompt_in  = gr.Textbox(label="Prompt", placeholder="Convert to anime style")
-            api_seed_in    = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
-            api_rand_in    = gr.Checkbox(value=True, label="Randomize seed")
-            api_cfg_in     = gr.Slider(1.0, 10.0, value=1.0, step=0.1, label="Guidance scale")
-            api_steps_in   = gr.Slider(1, 50, value=4, step=1, label="Steps")
-            api_run_btn    = gr.Button("Edit", variant="primary")
-        with gr.Column():
-            api_image_out  = gr.Image(type="pil", label="Result")
-            api_seed_out   = gr.Number(label="Seed used")
-
-    api_run_btn.click(
+# Wire a named API route directly onto the existing demo block
+# This exposes /api/edit without touching the UI at all
+with demo:
+    _api_img_in    = gr.Image(type="pil", visible=False, label="image")
+    _api_extra_in  = gr.Image(type="pil", visible=False, label="extra_image")
+    _api_prompt_in = gr.Textbox(visible=False, label="prompt")
+    _api_seed_in   = gr.Slider(0, MAX_SEED, value=0, step=1, visible=False, label="seed")
+    _api_rand_in   = gr.Checkbox(value=True, visible=False, label="randomize_seed")
+    _api_cfg_in    = gr.Slider(1.0, 10.0, value=1.0, step=0.1, visible=False, label="guidance_scale")
+    _api_steps_in  = gr.Slider(1, 50, value=4, step=1, visible=False, label="steps")
+    _api_img_out   = gr.Image(type="pil", visible=False, label="result_image")
+    _api_seed_out  = gr.Number(visible=False, label="seed_used")
+
+    _api_btn = gr.Button(visible=False)
+    _api_btn.click(
         fn=api_edit,
-        inputs=[api_image_in, api_prompt_in, api_seed_in, api_rand_in,
-                api_cfg_in, api_steps_in, api_extra_in],
-        outputs=[api_image_out, api_seed_out],
-        api_name="edit",          # → callable as client.predict(..., api_name="/edit")
+        inputs=[_api_img_in, _api_prompt_in, _api_seed_in, _api_rand_in,
+                _api_cfg_in, _api_steps_in, _api_extra_in],
+        outputs=[_api_img_out, _api_seed_out],
+        api_name="edit",
     )
 
-combined = gr.TabbedInterface(
-    [demo, api_demo],
-    ["UI", "API"],
-    title="FireRed Image Edit",
-)
-
 if __name__ == "__main__":
-    combined.queue(max_size=30).launch(
+    demo.queue(max_size=30).launch(
+        css=css,
         mcp_server=True,
         ssr_mode=False,
         show_error=True,