Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import random
 import os
 from PIL import Image
-from typing import Optional
+from typing import Optional, Dict, Tuple
 from huggingface_hub import InferenceClient
 
 # Project by Nymbo
@@ -10,9 +10,8 @@ from huggingface_hub import InferenceClient
 API_TOKEN = os.getenv("HF_READ_TOKEN")
 timeout = 100
 
-# Function to query the API and return the generated image
 def flux_krea_generate(
-    prompt: str, 
+    prompt: str,
     negative_prompt: str = "(deformed, distorted, disfigured), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, misspellings, typos",
     steps: int = 35,
     cfg_scale: float = 7.0,
@@ -20,27 +19,28 @@ def flux_krea_generate(
     seed: int = -1,
     strength: float = 0.7,
     width: int = 1024,
-    height: int = 1024
+    height: int = 1024,
 ) -> Optional[Image.Image]:
-    """
-    Text-to-image generation with FLUX.1-Krea-dev (no input image required).
+    """Generate a single image from a text prompt using FLUX.1-Krea-dev.
 
-    This tool generates a single image from a text prompt using the
-    black-forest-labs/FLUX.1-Krea-dev model on Hugging Face Inference.
+    Contract (for UI):
+    - Inputs: prompt (required), optional tuning params below. No input image used.
+    - Output: a PIL.Image for Gradio image component wiring.
+    - Errors: raises gr.Error on auth/model/service issues.
 
     Args:
-        prompt: Text description of the image to generate.
-        negative_prompt: What should NOT appear in the image.
-        steps: Number of denoising steps (1-100). Higher is slower but can improve quality.
-        cfg_scale: Classifier-free guidance scale (1-20). Higher = follow the prompt more closely.
-        sampler: Sampling method to use. One of: "DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM".
-        seed: Random seed for reproducible results. Use -1 for a random seed per call.
-        strength: Generation strength (0-1). Kept for parity; not an input image strength.
-        width: Output width in pixels (64-1216, multiple of 32 recommended).
-        height: Output height in pixels (64-1216, multiple of 32 recommended).
+        prompt: Required. Describes what to create. Keep it specific and concise.
+        negative_prompt: Phrases/objects to avoid in the output.
+        steps: Number of diffusion steps (1-100). Higher may improve detail but is slower.
+        cfg_scale: Classifier-free guidance scale (1-20). Higher forces closer adherence to prompt.
+        sampler: Sampler algorithm label (UI only; provider may ignore or auto-select).
+        seed: Set a deterministic seed (>=0). Use -1 to randomize per call.
+        strength: Kept for parity; has no effect without an input image (0-1).
+        width: Output width in pixels (64-1216). Prefer multiples of 32.
+        height: Output height in pixels (64-1216). Prefer multiples of 32.
 
     Returns:
-        A PIL.Image of the generated result. No input image is expected or required.
+        PIL.Image or None if prompt is empty.
     """
     if prompt == "" or prompt is None:
         return None
@@ -96,6 +96,120 @@ def flux_krea_generate(
             raise gr.Error(f"Image generation failed: {str(e)}")
         return None
 
+
+def _space_base_url() -> Optional[str]:
+    """Return the public base URL of this Space if available.
+
+    Looks up SPACE_ID (e.g. "username/space-name") and converts it to
+    the public subdomain "https://username-space-name.hf.space".
+    If SPACE_ID is not set, optionally respects HF_SPACE_BASE_URL.
+    """
+    # Explicit override if provided
+    explicit = os.getenv("HF_SPACE_BASE_URL")
+    if explicit:
+        return explicit.rstrip("/")
+
+    space_id = os.getenv("SPACE_ID")
+    if not space_id:
+        return None
+    sub = space_id.replace("/", "-")
+    return f"https://{sub}.hf.space"
+
+
+def _save_image_and_url(image: Image.Image, key: int) -> Tuple[str, Optional[str]]:
+    """Save the image to a temporary path and construct a public URL if running on Spaces.
+
+    Returns (local_path, public_url_or_None).
+    """
+    # Ensure POSIX-like temp dir (Spaces is Linux). Still works locally.
+    out_dir = os.path.join("/tmp", "flux-krea-outputs")
+    os.makedirs(out_dir, exist_ok=True)
+    file_path = os.path.join(out_dir, f"flux_{key}.png")
+    image.save(file_path)
+
+    base_url = _space_base_url()
+    public_url = None
+    if base_url:
+        # Gradio serves local files via /file=<absolute-path>
+        # Normalize backslashes to forward slashes in case of local dev on Windows
+        posix_path = file_path.replace("\\", "/")
+        public_url = f"{base_url}/file={posix_path}"
+
+    return file_path, public_url
+
+
+def flux_krea_generate_mcp(
+    prompt: str,
+    negative_prompt: str = "(deformed, distorted, disfigured), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, misspellings, typos",
+    steps: int = 35,
+    cfg_scale: float = 7,
+    sampler: str = "DPM++ 2M Karras",
+    seed: int = -1,
+    strength: float = 0.7,
+    width: int = 1024,
+    height: int = 1024,
+) -> Dict[str, object]:
+    """Generate an image (MCP tool) and return a JSON payload with a public URL.
+
+    This endpoint is tailored for Model Context Protocol (MCP) clients per the
+    latest Hugging Face MCP Space guidance (see hf-docs-search: "Spaces as MCP servers").
+
+    Inputs:
+      - prompt (str, required): Description of the desired image.
+      - negative_prompt (str): Items to avoid in the generation.
+      - steps (int, 1-100): Denoising steps. Higher is slower and may add detail.
+      - cfg_scale (float, 1-20): Guidance strength. Higher adheres more to the prompt.
+      - sampler (str): Sampler label (informational; provider may auto-select).
+      - seed (int): -1 for random per call; otherwise a deterministic seed >= 0.
+      - strength (float, 0-1): No-op in pure text-to-image; kept for cross-app parity.
+      - width (int, 64-1216): Output width (prefer multiples of 32).
+      - height (int, 64-1216): Output height (prefer multiples of 32).
+
+    Returns (JSON):
+      - image_path (str): Absolute path to the saved image on the Space VM.
+      - image_url (str|None): Publicly accessible URL to the image on the Space
+        (present when running on Spaces; None when running locally).
+      - seed (int): The seed used for this run (randomized if input was -1).
+      - width (int), height (int): Echo of output dimensions.
+      - sampler (str): Echo of requested sampler.
+
+    Error Modes:
+      - Raises a Gradio-friendly error with a concise message for common HTTP
+        failure codes (401/403 auth; 404 model; 503 warmup).
+    """
+    if not prompt:
+        raise gr.Error("'prompt' is required and cannot be empty.")
+
+    # Reuse core generator for image creation
+    image = flux_krea_generate(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        steps=steps,
+        cfg_scale=cfg_scale,
+        sampler=sampler,
+        seed=seed,
+        strength=strength,
+        width=width,
+        height=height,
+    )
+
+    if image is None:
+        raise gr.Error("No image generated.")
+
+    # Save and build URLs
+    key = random.randint(0, 999)
+    file_path, public_url = _save_image_and_url(image, key)
+
+    # Expose URL explicitly for MCP clients (LLMs need a resolvable URL)
+    return {
+        "image_path": file_path,
+        "image_url": public_url,
+        "seed": seed,
+        "width": width,
+        "height": height,
+        "sampler": sampler,
+    }
+
 # CSS to style the app
 css = """
 #app-container {
@@ -152,13 +266,14 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
         )
 
     # Expose a dedicated MCP/API endpoint with a clear schema (text-to-image only)
-    # This avoids clients misinterpreting the UI event as requiring an input image.
+    # This wrapper returns both a local file path and a fully-qualified public URL
+    # when running on Spaces so LLMs can access the finished image.
     gr.api(
-        flux_krea_generate,
+        flux_krea_generate_mcp,
         api_name="generate_image",
         api_description=(
             "Generate an image from a text prompt using FLUX.1-Krea-dev. "
-            "Inputs are text and numeric parameters only; no input image is required."
+            "Returns JSON with image_path and image_url (public URL when on Spaces)."
         ),
     )