Update app.py

app.py

@@ -5,13 +5,24 @@
 
 import gradio as gr
 import numpy as np
-import spaces
 import torch
 import random
 from PIL import Image
 
 from diffusers import FluxKontextPipeline
-from diffusers.utils import load_image
+ 
+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()
 
 # -----------------------------
 # Constants & model bootstrap
@@ -22,16 +33,18 @@ 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=torch.bfloat16
-).to("cuda")
+    torch_dtype=_dtype,
+).to(_device)
 
 # ---------------------------------------------------------
 # Core editing function (works WITH or WITHOUT input image)
 # ---------------------------------------------------------
 
-@spaces.GPU
+@GPU_DECORATOR
 def infer(
     input_image: Image.Image | None,
     prompt: str,
@@ -90,7 +103,7 @@ def infer(
 # NEW: Dedicated text-to-image function (separate MCP tool)
 # ------------------------------------------------------------
 
-@spaces.GPU  # (layman's) make sure we run on the GPU so it's fast
+@GPU_DECORATOR  # (layman's) make sure we run on the GPU so it's fast
 def text_to_image(
     prompt: str,
     seed: int = 42,
@@ -140,7 +153,7 @@ def text_to_image(
 # Lightweight helper for the Examples
 # -------------------------------------
 
-@spaces.GPU(duration=25)
+@GPU_DECORATOR(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)
@@ -246,18 +259,18 @@ 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])
 
-# ------------------------------------------------------------------
-# NEW: Register a dedicated MCP tool that does text-to-image only.
-#      This does not create any extra UI — it's a clean API endpoint.
-# ------------------------------------------------------------------
-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."
-    ),
-)
+    # ------------------------------------------------------------------
+    # 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."
+        ),
+    )
 
 # (layman's) start the app with MCP enabled so tools show up to agents (e.g., Claude/Cursor)
-demo.launch(mcp_server=True)
+demo.launch(mcp_server=True)