Update app.py

app.py

@@ -5,24 +5,13 @@
 
 import gradio as gr
 import numpy as np
+import spaces
 import torch
 import random
 from PIL import Image
 
 from diffusers import FluxKontextPipeline
- 
-try:
-    import spaces  # Hugging Face Spaces runtime (GPU decorators)
-    GPU_DECORATOR = spaces.GPU
-except Exception:
-    # Fallback: no-op decorator when not running on Spaces
-    class _NoOpGPU:
-        def __call__(self, *args, **kwargs):
-            def _wrap(fn):
-                return fn
-            return _wrap
-
-    GPU_DECORATOR = _NoOpGPU()
+from diffusers.utils import load_image
 
 # -----------------------------
 # Constants & model bootstrap
@@ -33,18 +22,16 @@ MAX_SEED = np.iinfo(np.int32).max  # <-- (layman's) the biggest safe random seed
 
 # Load the FLUX.1 Kontext-dev pipeline once and keep it on GPU for speed
 # (layman's) this downloads the model and prepares it to run on your graphics card
-_device = "cuda" if torch.cuda.is_available() else "cpu"
-_dtype = torch.bfloat16 if _device == "cuda" else torch.float32
 pipe = FluxKontextPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-Kontext-dev",
-    torch_dtype=_dtype,
-).to(_device)
+    torch_dtype=torch.bfloat16
+).to("cuda")
 
 # ---------------------------------------------------------
 # Core editing function (works WITH or WITHOUT input image)
 # ---------------------------------------------------------
 
-@GPU_DECORATOR
+@spaces.GPU
 def infer(
     input_image: Image.Image | None,
     prompt: str,
@@ -103,7 +90,7 @@ def infer(
 # NEW: Dedicated text-to-image function (separate MCP tool)
 # ------------------------------------------------------------
 
-@GPU_DECORATOR  # (layman's) make sure we run on the GPU so it's fast
+@spaces.GPU  # (layman's) make sure we run on the GPU so it's fast
 def text_to_image(
     prompt: str,
     seed: int = 42,
@@ -153,7 +140,7 @@ def text_to_image(
 # Lightweight helper for the Examples
 # -------------------------------------
 
-@GPU_DECORATOR(duration=25)
+@spaces.GPU(duration=25)
 def infer_example(input_image: Image.Image | None, prompt: str) -> tuple[Image.Image, int]:
     # (layman's) small wrapper used by the clickable examples
     image, seed, _ = infer(input_image, prompt)
@@ -259,18 +246,38 @@ Image editing and manipulation model guidance-distilled from FLUX.1 Kontext [pro
     # (Optional) If you want a 1-click "reuse image" flow in the UI later:
     # reuse_button.click(fn=lambda image: image, inputs=[result], outputs=[input_image])
 
-    # ------------------------------------------------------------------
-    # Register a dedicated MCP tool that does text-to-image only (no UI).
-    # Placing gr.api inside the Blocks context registers it with this demo.
-    # ------------------------------------------------------------------
-    gr.api(
-        text_to_image,
-        api_name="text_to_image",  # MCP tool route
-        api_description=(
-            "Generate a brand-new image from text (no input image required) "
-            "using FLUX.1 Kontext-dev. Returns the image and the seed used."
-        ),
-    )
+# ------------------------------------------------------------------
+# NEW: Create a dedicated Interface for text-to-image MCP tool
+# This ensures better compatibility with MCP clients
+# ------------------------------------------------------------------
+
+# Create a separate interface for the text-to-image tool
+text_to_image_interface = gr.Interface(
+    fn=text_to_image,
+    inputs=[
+        gr.Text(label="Prompt", placeholder="Describe the image you want to generate", value=""),
+        gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=42, step=1),
+        gr.Checkbox(label="Randomize seed", value=True),
+        gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, value=2.5, step=0.1),
+        gr.Slider(label="Steps", minimum=1, maximum=30, value=20, step=1),
+        gr.Slider(label="Width", minimum=256, maximum=2048, value=1024, step=64),
+        gr.Slider(label="Height", minimum=256, maximum=2048, value=1024, step=64),
+    ],
+    outputs=[
+        gr.Image(label="Generated Image"),
+        gr.Number(label="Seed Used")
+    ],
+    title="FLUX.1 Text-to-Image Generator",
+    description="Generate high-quality images from text descriptions using FLUX.1 Kontext-dev",
+    api_name="text_to_image"
+)
+
+# Mount both interfaces using gr.TabbedInterface for better organization
+combined_demo = gr.TabbedInterface(
+    [demo, text_to_image_interface], 
+    ["Image Editor", "Text-to-Image Generator"],
+    title="FLUX.1 Kontext Tools"
+)
 
 # (layman's) start the app with MCP enabled so tools show up to agents (e.g., Claude/Cursor)
-demo.launch(mcp_server=True)
+combined_demo.launch(mcp_server=True)