Update app.py

app.py

@@ -78,6 +78,15 @@ DEFAULT_CFG = {
 }
 
 def image_to_data_uri(img):
+    """
+    Convert a PIL Image to a base64 data URI.
+
+    Args:
+        img: The PIL Image to convert.
+
+    Returns:
+        str: A data URI string containing the base64-encoded PNG image.
+    """
     buffered = io.BytesIO()
     img.save(buffered, format="PNG")
     img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
@@ -85,6 +94,16 @@ def image_to_data_uri(img):
 
 
 def upsample_prompt_logic(prompt, image_list):
+    """
+    Enhance a text prompt using a Vision-Language Model.
+
+    Args:
+        prompt (str): The original text prompt to enhance.
+        image_list: Optional list of PIL Images for context-aware enhancement.
+
+    Returns:
+        str: The enhanced prompt, or the original prompt if enhancement fails.
+    """
     try:
         if image_list and len(image_list) > 0:
             # Image + Text Editing Mode
@@ -125,8 +144,18 @@ def upsample_prompt_logic(prompt, image_list):
 
 
 def update_dimensions_from_image(image_list):
-    """Update width/height sliders based on uploaded image aspect ratio.
-    Keeps one side at 1024 and scales the other proportionally, with both sides as multiples of 8."""
+    """
+    Update width/height based on uploaded image aspect ratio.
+    
+    Keeps one side at 1024 and scales the other proportionally,
+    with both sides as multiples of 8.
+
+    Args:
+        image_list: Gallery list of tuples (image, caption) from Gradio.
+
+    Returns:
+        tuple: A tuple of (width, height) integers, both multiples of 8.
+    """
     if image_list is None or len(image_list) == 0:
         return 1024, 1024  # Default dimensions
     
@@ -155,12 +184,69 @@ def update_dimensions_from_image(image_list):
 
 
 def update_steps_from_mode(mode_choice):
-    """Update the number of inference steps based on the selected mode."""
+    """
+    Update inference steps and guidance scale based on the selected mode.
+
+    Args:
+        mode_choice (str): The selected mode, either "Distilled (4 steps)" or "Base (50 steps)".
+
+    Returns:
+        tuple: A tuple of (num_inference_steps, guidance_scale).
+    """
     return DEFAULT_STEPS[mode_choice], DEFAULT_CFG[mode_choice]
 
 
 @spaces.GPU(duration=85)
-def infer(prompt, input_images=None, mode_choice="Distilled (4 steps)", seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, guidance_scale=4.0, prompt_upsampling=False, progress=gr.Progress(track_tqdm=True)):
+def infer(
+    prompt: str,
+    input_images=None,
+    mode_choice: str = "Distilled (4 steps)",
+    seed: int = 42,
+    randomize_seed: bool = False,
+    width: int = 1024,
+    height: int = 1024,
+    num_inference_steps: int = 4,
+    guidance_scale: float = 4.0,
+    prompt_upsampling: bool = False,
+    progress=gr.Progress(track_tqdm=True)
+):
+    """
+    Generate or edit images using FLUX.2 Klein 9B model.
+    
+    This tool can generate images from text prompts, or edit/combine existing images
+    based on text instructions. Use the distilled mode for fast 4-step generation,
+    or base mode for higher quality 50-step generation.
+
+    Args:
+        prompt (str): Text description of the image to generate, or editing instructions when input images are provided.
+        input_images: Optional list of input images for editing or combining. Provide image URLs.
+        mode_choice (str): Model mode - "Distilled (4 steps)" for fast generation or "Base (50 steps)" for higher quality.
+        seed (str): Random seed for reproducible generation. Use "0" with randomize_seed=True for random results.
+        randomize_seed (str): Set to "true" to use a random seed, "false" to use the specified seed.
+        width (str): Output image width in pixels (256-1024, must be multiple of 8).
+        height (str): Output image height in pixels (256-1024, must be multiple of 8).
+        num_inference_steps (str): Number of denoising steps. Use "4" for distilled mode, "50" for base mode.
+        guidance_scale (str): How closely to follow the prompt. Use "1.0" for distilled, "4.0" for base mode.
+        prompt_upsampling (str): Set to "true" to automatically enhance the prompt using a VLM.
+
+    Returns:
+        tuple: A tuple containing the generated PIL Image and the seed used.
+    """
+    # Convert string inputs to proper types for MCP compatibility
+    if isinstance(seed, str):
+        seed = int(seed)
+    if isinstance(randomize_seed, str):
+        randomize_seed = randomize_seed.lower() == "true"
+    if isinstance(width, str):
+        width = int(width)
+    if isinstance(height, str):
+        height = int(height)
+    if isinstance(num_inference_steps, str):
+        num_inference_steps = int(num_inference_steps)
+    if isinstance(guidance_scale, str):
+        guidance_scale = float(guidance_scale)
+    if isinstance(prompt_upsampling, str):
+        prompt_upsampling = prompt_upsampling.lower() == "true"
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -356,7 +442,9 @@ FLUX.2 [Klein] is a distilled model capable of generating, editing and combining
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[prompt, input_images, mode_choice, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling],
-        outputs=[result, seed]
+        outputs=[result, seed],
+        api_name="generate"  # Explicit API name for MCP tool
     )
 
-demo.launch()
+# Launch with MCP server enabled
+demo.launch(mcp_server=True)