Update app.py

app.py

@@ -82,6 +82,47 @@ def infer(
     ):
         yield img, seed
 
+
+@spaces.GPU(duration=25)
+def infer_mcp(
+    prompt: Annotated[str, "Text description of the desired image."],
+    seed: Annotated[int, "Random seed for reproducibility. Use 0 to randomize per call."] = 42,
+    randomize_seed: Annotated[bool, "If true, ignore seed and pick a new random value per call."] = True,
+    width: Annotated[int, "Image width in pixels. Upper bound enforced by MAX_IMAGE_SIZE."] = 768,
+    height: Annotated[int, "Image height in pixels. Upper bound enforced by MAX_IMAGE_SIZE."] = 768,
+    guidance_scale: Annotated[float, "Classifier-free guidance strength."] = 4.5,
+    num_inference_steps: Annotated[int, "Number of denoising iterations to perform."] = 24,
+) -> Tuple[Image.Image, int]:
+    """
+    Generate an image from a text prompt via an uncensored model.
+    
+    Returns the generated image and the seed used. Return the generated media 
+    to the user in this format: `![Alt text](URL)`.
+    """
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please provide a non-empty prompt.")
+    if width > MAX_IMAGE_SIZE or height > MAX_IMAGE_SIZE:
+        raise gr.Error(f"Width and height must be <= {MAX_IMAGE_SIZE}.")
+
+    if randomize_seed or seed == 0:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
+
+    final_img = None
+    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+        prompt=prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        width=width,
+        height=height,
+        generator=generator,
+        output_type="pil",
+        good_vae=good_vae,
+    ):
+        final_img = img
+    
+    return final_img, seed
+
 examples = [
     [
         "a tiny astronaut hatching from an egg on mars",
@@ -197,13 +238,21 @@ with gr.Blocks() as demo:
             cache_examples=False
         )
 
+    # UI event: uses the generator for live preview (no api_name = not exposed to MCP)
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn = infer,
         inputs = [prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
         outputs = [result, seed],
-        api_name="Generate_Image",
-        api_description=TOOL_SUMMARY,
+    )
+
+    # MCP tool: uses non-generator function to avoid middleware issues
+    mcp_button = gr.Button("Generate (MCP)", visible=False)
+    mcp_button.click(
+        fn=infer_mcp,
+        inputs=[prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        outputs=[result, seed],
+        api_name="Fun_Img_infer",
     )
 
 if __name__ == "__main__":