Update app.py

app.py

@@ -2,16 +2,20 @@ import gradio as gr
 import random
 import os
 from PIL import Image
-from typing import Optional, Dict, Tuple
+from typing import Optional
 from huggingface_hub import InferenceClient
+import tempfile
+import json
+import uuid
 
-# Project by Nymbo
+# Project by Nymbo - Converted to MCP Server
 
+# Configuration
 API_TOKEN = os.getenv("HF_READ_TOKEN")
 timeout = 100
 
 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,
@@ -19,41 +23,45 @@ def flux_krea_generate(
     seed: int = -1,
     strength: float = 0.7,
     width: int = 1024,
-    height: int = 1024,
-) -> Optional[Image.Image]:
-    """Generate a single image from a text prompt using FLUX.1-Krea-dev.
-
-    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.
+    height: int = 1024
+) -> str:
+    """
+    Generate high-quality professional images using FLUX.1-Krea-dev via HuggingFace Serverless Inference.
+    
+    This MCP tool creates realistic, professional-quality images from text prompts using the 
+    FLUX.1-Krea-dev model through HuggingFace's serverless inference infrastructure. The model 
+    excels at generating natural-looking images without typical AI artifacts, making it ideal 
+    for product photography, e-commerce visuals, concept art, fashion photography, and stock images.
 
     Args:
-        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.
+        prompt: Detailed text description of the image to generate. Best results with approximately 60-70 words max. Include technical photography terms, specific lighting, and material textures for photorealistic results.
+        negative_prompt: Elements to avoid in the generated image (deformities, poor anatomy, etc.).
+        steps: Number of denoising steps (1-100). Higher values generally improve quality but increase generation time.
+        cfg_scale: Classifier-free guidance scale (1-20). Higher values make the model follow the prompt more closely but may reduce creativity.
+        sampler: Sampling method for diffusion process. Options: "DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM".
+        seed: Random seed for reproducible results. Use -1 for random generation each time.
+        strength: Generation strength parameter (0-1). Controls the influence of the generation process.
+        width: Output image width in pixels (64-1216). Recommend multiples of 32 for best results.
+        height: Output image height in pixels (64-1216). Recommend multiples of 32 for best results.
 
     Returns:
-        PIL.Image or None if prompt is empty.
+        JSON string containing the image URL and generation metadata that can be accessed by LLMs and MCP clients.
     """
-    if prompt == "" or prompt is None:
-        return None
+    if not prompt or prompt.strip() == "":
+        return json.dumps({
+            "success": False,
+            "error": "Prompt is required and cannot be empty",
+            "image_url": None
+        })
 
-    key = random.randint(0, 999)
+    generation_id = random.randint(0, 999)
     
-    # Add some extra flair to the prompt
+    # Enhance prompt for better quality
     enhanced_prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
-    print(f'\033[1mGeneration {key}:\033[0m {enhanced_prompt}')
+    print(f'\033[1mMCP Generation {generation_id}:\033[0m {enhanced_prompt}')
     
     try:
-        # Initialize the Hugging Face Inference Client
-        # Try different providers in order of preference
+        # Initialize the Hugging Face Inference Client with fallback providers
         providers = ["auto", "replicate", "fal-ai"]
         
         for provider in providers:
@@ -63,7 +71,7 @@ def flux_krea_generate(
                     provider=provider
                 )
                 
-                # Generate the image using the proper client
+                # Generate the image using serverless inference
                 image = client.text_to_image(
                     prompt=enhanced_prompt,
                     negative_prompt=negative_prompt,
@@ -75,207 +83,252 @@ def flux_krea_generate(
                     seed=seed if seed != -1 else random.randint(1, 1000000000)
                 )
                 
-                print(f'\033[1mGeneration {key} completed with {provider}!\033[0m ({enhanced_prompt})')
-                return image
+                # Save the image to a temporary file that Gradio can serve
+                # This creates a local file that Gradio can serve with a public URL
+                if image:
+                    # Generate a unique filename
+                    unique_id = str(uuid.uuid4())[:8]
+                    temp_filename = f"flux_krea_{generation_id}_{unique_id}.png"
+                    
+                    # Save to temp directory
+                    temp_file = tempfile.NamedTemporaryFile(
+                        delete=False, 
+                        suffix=".png", 
+                        prefix=f"flux_krea_{generation_id}_"
+                    )
+                    image.save(temp_file.name)
+                    temp_file.close()
+                    
+                    # For MCP clients, we need to provide a URL that they can access
+                    # Gradio will serve temp files via the /file/ endpoint
+                    file_url = temp_file.name
+                    
+                    print(f'\033[1mMCP Generation {generation_id} completed with {provider}!\033[0m')
+                    
+                    # Return JSON with both file path and success metadata
+                    result = {
+                        "success": True,
+                        "image_url": file_url,
+                        "image_path": file_url,
+                        "generation_id": generation_id,
+                        "provider": provider,
+                        "model": "black-forest-labs/FLUX.1-Krea-dev",
+                        "prompt": enhanced_prompt,
+                        "parameters": {
+                            "width": width,
+                            "height": height,
+                            "steps": steps,
+                            "cfg_scale": cfg_scale,
+                            "seed": seed if seed != -1 else "random",
+                            "sampler": sampler
+                        },
+                        "metadata": {
+                            "tool": "flux_krea_generate",
+                            "timestamp": str(generation_id),
+                            "mcp_compatible": True
+                        }
+                    }
+                    
+                    return json.dumps(result)
                 
             except Exception as provider_error:
                 print(f"Provider {provider} failed: {provider_error}")
-                if provider == providers[-1]:  # Last provider
+                if provider == providers[-1]:  # Last provider failed
                     raise provider_error
                 continue
         
     except Exception as e:
-        print(f"Error during image generation: {e}")
+        print(f"Error during MCP image generation: {e}")
+        error_message = "Image generation failed due to an unknown error."
+        
         if "404" in str(e):
-            raise gr.Error("Model not found. Please ensure the FLUX.1-Krea-dev model is accessible with your API token.")
+            error_message = "Model not found. Please ensure the FLUX.1-Krea-dev model is accessible with your API token."
         elif "503" in str(e):
-            raise gr.Error("The model is currently being loaded. Please try again in a moment.")
+            error_message = "The model is currently being loaded. Please try again in a moment."
         elif "401" in str(e) or "403" in str(e):
-            raise gr.Error("Authentication failed. Please check your HF_READ_TOKEN environment variable.")
-        else:
-            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)
+            error_message = "Authentication failed. Please check your HF_READ_TOKEN environment variable."
+        
+        return json.dumps({
+            "success": False,
+            "error": error_message,
+            "image_url": None,
+            "generation_id": generation_id
+        })
 
-    # 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,
-    }
+# For UI compatibility - this function returns a PIL Image for the Gradio interface
+def flux_krea_generate_ui(*args) -> Optional[Image.Image]:
+    """UI wrapper that returns PIL Image for Gradio interface"""
+    result_json = flux_krea_generate(*args)
+    try:
+        result = json.loads(result_json)
+        if result.get("success") and result.get("image_path"):
+            return Image.open(result["image_path"])
+    except:
+        pass
+    return None
 
-# CSS to style the app
+# CSS for improved styling
 css = """
 #app-container {
-    max-width: 800px;
+    max-width: 900px;
     margin-left: auto;
     margin-right: auto;
 }
+.mcp-badge {
+    background: linear-gradient(45deg, #ff6b6b, #4ecdc4);
+    color: white;
+    padding: 5px 10px;
+    border-radius: 15px;
+    font-size: 12px;
+    font-weight: bold;
+}
 """
 
-# Build the Gradio UI with Blocks
-with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
-    # Add a title to the app
-    gr.HTML("<center><h1>FLUX.1-Krea-dev</h1></center>")
-    gr.HTML("<center><p>High-quality image generation via Model Context Protocol</p></center>")
+# Build the Gradio app
+with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css, title="FLUX.1-Krea MCP Server") as app:
+    # Header
+    gr.HTML("""
+    <center>
+        <h1>🚀 FLUX.1-Krea-dev <span class="mcp-badge">MCP SERVER</span></h1>
+        <p>High-quality serverless image generation via Model Context Protocol</p>
+        <p><em>Professional-grade images • No AI artifacts • MCP-compatible</em></p>
+    </center>
+    """)
     
-    # Container for all the UI elements
+    # Main interface
     with gr.Column(elem_id="app-container"):
-        # Add a text input for the main prompt
+        # Prompt input
         with gr.Row():
-            with gr.Column(elem_id="prompt-container"):
+            text_prompt = gr.Textbox(
+                label="Image Prompt", 
+                placeholder="Describe the image you want to generate (60-70 words recommended)", 
+                lines=3, 
+                elem_id="prompt-text-input"
+            )
+        
+        # Advanced settings accordion
+        with gr.Row():
+            with gr.Accordion("🔧 Advanced Generation Settings", open=False):
+                negative_prompt = gr.Textbox(
+                    label="Negative Prompt", 
+                    placeholder="What should NOT be in the image", 
+                    value="(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",
+                    lines=3, 
+                    elem_id="negative-prompt-text-input"
+                )
+                
+                with gr.Row():
+                    width = gr.Slider(
+                        label="Width", 
+                        value=1024, 
+                        minimum=64, 
+                        maximum=1216, 
+                        step=32,
+                        info="Output width in pixels"
+                    )
+                    height = gr.Slider(
+                        label="Height", 
+                        value=1024, 
+                        minimum=64, 
+                        maximum=1216, 
+                        step=32,
+                        info="Output height in pixels"
+                    )
+                
                 with gr.Row():
-                    text_prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here", lines=2, elem_id="prompt-text-input")
+                    steps = gr.Slider(
+                        label="Sampling Steps", 
+                        value=35, 
+                        minimum=1, 
+                        maximum=100, 
+                        step=1,
+                        info="More steps = higher quality, longer generation time"
+                    )
+                    cfg = gr.Slider(
+                        label="CFG Scale", 
+                        value=7, 
+                        minimum=1, 
+                        maximum=20, 
+                        step=1,
+                        info="How closely to follow the prompt"
+                    )
                 
-                # Accordion for advanced settings
                 with gr.Row():
-                    with gr.Accordion("Advanced Settings", open=False):
-                        negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="What should not be in the image", value="(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", lines=3, elem_id="negative-prompt-text-input")
-                        with gr.Row():
-                            width = gr.Slider(label="Width", value=1024, minimum=64, maximum=1216, step=32)
-                            height = gr.Slider(label="Height", value=1024, minimum=64, maximum=1216, step=32)
-                        steps = gr.Slider(label="Sampling steps", value=35, minimum=1, maximum=100, step=1)
-                        cfg = gr.Slider(label="CFG Scale", value=7, minimum=1, maximum=20, step=1)
-                        strength = gr.Slider(label="Strength", value=0.7, minimum=0, maximum=1, step=0.001)
-                        seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=1000000000, step=1) # Setting the seed to -1 will make it random
-                        method = gr.Radio(label="Sampling method", value="DPM++ 2M Karras", choices=["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"])
+                    strength = gr.Slider(
+                        label="Strength", 
+                        value=0.7, 
+                        minimum=0, 
+                        maximum=1, 
+                        step=0.01,
+                        info="Generation strength parameter"
+                    )
+                    seed = gr.Slider(
+                        label="Seed", 
+                        value=-1, 
+                        minimum=-1, 
+                        maximum=1000000000, 
+                        step=1,
+                        info="Use -1 for random seed"
+                    )
+                
+                sampler = gr.Radio(
+                    label="Sampling Method", 
+                    value="DPM++ 2M Karras", 
+                    choices=["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"],
+                    info="Algorithm used for image generation"
+                )
 
-        # Add a button to trigger the image generation
+        # Generation button
         with gr.Row():
-            text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
+            generate_button = gr.Button("🎨 Generate Image", variant='primary', elem_id="gen-button", scale=1)
         
-        # Image output area to display the generated image
+        # Output area
         with gr.Row():
-            # Output component only; no input image is required by the tool
-            image_output = gr.Image(label="Image Output", elem_id="gallery")
+            image_output = gr.Image(
+                label="Generated Image", 
+                elem_id="gallery",
+                show_share_button=True,
+                show_download_button=True
+            )
         
-        # Bind the button to the flux_krea_generate function for the UI only
-        # Hide this event as an MCP tool to avoid schema confusion (UI wires image output)
-        text_button.click(
-            flux_krea_generate,
-            inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, width, height],
+        # MCP Information
+        with gr.Row():
+            gr.HTML("""
+            <div style="background: #f0f0f0; padding: 15px; border-radius: 10px; margin-top: 20px;">
+                <h3>📡 MCP Server Information</h3>
+                <p><strong>Server Endpoint:</strong> <code>/gradio_api/mcp/sse</code></p>
+                <p><strong>Tool Name:</strong> <code>flux_krea_generate</code></p>
+                <p>This server exposes the image generation function as an MCP tool that returns JSON with image URLs for LLM integration.</p>
+            </div>
+            """)
+        
+        # Wire up the UI event (this is separate from the MCP tool)
+        generate_button.click(
+            flux_krea_generate_ui,
+            inputs=[text_prompt, negative_prompt, steps, cfg, sampler, seed, strength, width, height],
             outputs=image_output,
-            show_api=False,
-            api_description=False,
+            show_api=False  # Hide from API docs since we have dedicated MCP tool
         )
 
-    # Expose a dedicated MCP/API endpoint with a clear schema (text-to-image only)
-    # 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.
+    # Expose the MCP tool with clear API documentation
     gr.api(
-        flux_krea_generate_mcp,
-        api_name="generate_image",
+        flux_krea_generate,
+        api_name="flux_krea_generate",
         api_description=(
-            "Generate an image from a text prompt using FLUX.1-Krea-dev. "
-            "Returns JSON with image_path and image_url (public URL when on Spaces)."
-        ),
+            "MCP Tool: Generate professional-quality images using FLUX.1-Krea-dev via serverless inference. "
+            "Returns JSON with image URL and metadata for MCP clients and LLMs. "
+            "Optimized for natural, realistic images without AI artifacts."
+        )
     )
 
-# Launch the Gradio app with MCP server enabled
-app.launch(show_api=True, share=False, mcp_server=True)
+# Launch with MCP server enabled
+if __name__ == "__main__":
+    # Enable MCP server functionality
+    app.launch(
+        mcp_server=True,
+        show_api=True,
+        share=False,
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True
+    )