Adding HF inference image generation to MCP tools (Flux Krea by default)

app.py

@@ -9,6 +9,8 @@ from __future__ import annotations
 import re
 import json
 import sys
+import os
+import random
 from io import StringIO
 from typing import List, Dict, Tuple, Annotated
 
@@ -18,6 +20,8 @@ from bs4 import BeautifulSoup
 from readability import Document
 from urllib.parse import urljoin, urldefrag, urlparse
 from duckduckgo_search import DDGS
+from PIL import Image
+from huggingface_hub import InferenceClient
 
 # Optional imports for Kokoro TTS (loaded lazily)
 import numpy as np
@@ -626,7 +630,7 @@ CSS_STYLES = """
     }
     /* Default: add subtitle under titles */
     .gradio-container h1::after {
-        content: "Fetch Webpage | Search DuckDuckGo | Code Interpreter | Kokoro TTS";
+    content: "Fetch Webpage | Search DuckDuckGo | Code Interpreter | Kokoro TTS | Flux Krea (Serverless)";
         display: block;
         font-size: 1rem;
         font-weight: 500;
@@ -660,14 +664,136 @@ kokoro_interface = gr.Interface(
     allow_flagging="never",
 )
 
-# Build tabbed app including Kokoro
+# ==========================
+# Flux Krea (Serverless)
+# ==========================
+
+HF_API_TOKEN = os.getenv("HF_READ_TOKEN")
+
+
+def Flux_Krea_Generate(  # <-- MCP tool #5 (Flux Krea - Serverless)
+    prompt: Annotated[str, "Text description of the image to generate."],
+    negative_prompt: Annotated[str, "What should NOT appear in the image." ] = (
+        "(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: Annotated[int, "Number of denoising steps (1–100). Higher = slower, potentially higher quality."] = 35,
+    cfg_scale: Annotated[float, "Classifier-free guidance scale (1–20). Higher = follow the prompt more closely."] = 7.0,
+    sampler: Annotated[str, "Sampling method label (UI only). Common options: 'DPM++ 2M Karras', 'DPM++ SDE Karras', 'Euler', 'Euler a', 'Heun', 'DDIM'."] = "DPM++ 2M Karras",
+    seed: Annotated[int, "Random seed for reproducibility. Use -1 for a random seed per call."] = -1,
+    width: Annotated[int, "Output width in pixels (64–1216, multiple of 32 recommended)."] = 1024,
+    height: Annotated[int, "Output height in pixels (64–1216, multiple of 32 recommended)."] = 1024,
+) -> Image.Image:
+    """
+    Generate a single image from a text prompt using the black-forest-labs/FLUX.1-Krea-dev
+    model via Hugging Face Inference (serverless). Returns a PIL image.
+
+    Notes (MCP):
+    - Per the latest Gradio MCP docs, images returned from tools are handled by the server and
+      converted to file URLs automatically for MCP clients. Ensure type hints and this docstring
+      "Args:" block are present so the tool schema is accurate.
+
+    Args:
+        prompt (str): Text description of the image to generate.
+        negative_prompt (str): What should NOT appear in the image.
+        steps (int): Number of denoising steps (1–100). Higher can improve quality.
+        cfg_scale (float): Guidance scale (1–20). Higher = follow the prompt more closely.
+        sampler (str): Sampling method label for UI; not all providers expose this control.
+        seed (int): Random seed. Use -1 to randomize on each call.
+        width (int): Output width in pixels (64–1216; multiples of 32 recommended).
+        height (int): Output height in pixels (64–1216; multiples of 32 recommended).
+
+    Returns:
+        PIL.Image.Image: The generated image.
+
+    Error modes:
+        - Raises gr.Error with a user-friendly message on auth/model/load errors.
+    """
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please provide a non-empty prompt.")
+
+    # Slightly enhance prompt for quality (kept consistent with Serverless space)
+    enhanced_prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
+
+    # Try multiple providers for resilience
+    providers = ["auto", "replicate", "fal-ai"]
+    last_error: Exception | None = None
+
+    for provider in providers:
+        try:
+            client = InferenceClient(api_key=HF_API_TOKEN, provider=provider)
+            image = client.text_to_image(
+                prompt=enhanced_prompt,
+                negative_prompt=negative_prompt,
+                model="black-forest-labs/FLUX.1-Krea-dev",
+                width=width,
+                height=height,
+                num_inference_steps=steps,
+                guidance_scale=cfg_scale,
+                seed=seed if seed != -1 else random.randint(1, 1_000_000_000),
+            )
+            return image
+        except Exception as e:  # try next provider, transform last one to friendly error
+            last_error = e
+            continue
+
+    # If we reach here, all providers failed
+    msg = str(last_error) if last_error else "Unknown error"
+    if "404" in msg:
+        raise gr.Error("Model not found. Ensure 'FLUX.1-Krea-dev' is accessible with your HF token.")
+    if "503" in msg:
+        raise gr.Error("The model is warming up. Please try again shortly.")
+    if "401" in msg or "403" in msg:
+        raise gr.Error("Authentication failed. Set HF_READ_TOKEN environment variable with access to the model.")
+    raise gr.Error(f"Image generation failed: {msg}")
+
+
+flux_krea_interface = gr.Interface(
+    fn=Flux_Krea_Generate,
+    inputs=[
+        gr.Textbox(label="Prompt", placeholder="Enter a prompt", lines=2),
+        gr.Textbox(
+            label="Negative Prompt",
+            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=2,
+        ),
+        gr.Slider(minimum=1, maximum=100, value=35, step=1, label="Steps"),
+        gr.Slider(minimum=1.0, maximum=20.0, value=7.0, step=0.1, label="CFG Scale"),
+        gr.Radio(label="Sampler", value="DPM++ 2M Karras", choices=[
+            "DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"
+        ]),
+        gr.Slider(minimum=-1, maximum=1_000_000_000, value=-1, step=1, label="Seed (-1 = random)"),
+        gr.Slider(minimum=64, maximum=1216, value=1024, step=32, label="Width"),
+        gr.Slider(minimum=64, maximum=1216, value=1024, step=32, label="Height"),
+    ],
+    outputs=gr.Image(label="Generated Image"),
+    title="Flux Krea (Serverless)",
+    description=(
+        "<div style=\"text-align:center\">Generate images with FLUX.1-Krea-dev via Hugging Face Inference. "
+        "No input image required. 'Strength' is omitted. Set HF_READ_TOKEN as needed.</div>"
+    ),
+    api_description=(
+        "Generate a single image from a text prompt using black-forest-labs/FLUX.1-Krea-dev on Hugging Face Inference. "
+        "Parameters: prompt (str), negative_prompt (str), steps (int, 1–100), cfg_scale (float, 1–20), "
+        "sampler (str, label only), seed (int, -1=random), width/height (int, 64–1216). Returns a PIL.Image."
+    ),
+    allow_flagging="never",
+)
+
+# Build tabbed app including Flux Krea
 demo = gr.TabbedInterface(
-    interface_list=[fetch_interface, concise_interface, code_interface, kokoro_interface],
+    interface_list=[fetch_interface, concise_interface, code_interface, kokoro_interface, flux_krea_interface],
     tab_names=[
         "Fetch Webpage",
         "DuckDuckGo Search",
         "Python Code Executor",
         "Kokoro TTS",
+        "Flux Krea (Serverless)",
     ],
     title="Tools MCP",
     theme="Nymbo/Nymbo_Theme",