Update app.py

app.py

@@ -4,9 +4,9 @@ import io
 import random
 import os
 import time
-from typing import Optional
 from PIL import Image
 import json
+from typing import Optional
 
 # Project by Nymbo
 
@@ -15,76 +15,63 @@ API_TOKEN = os.getenv("HF_READ_TOKEN")
 headers = {"Authorization": f"Bearer {API_TOKEN}"}
 timeout = 100
 
-def query(
-    prompt: str,
-    negative_prompt: str = "",
+# Function to query the API and return the generated image
+def flux_krea_generate(
+    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,
+    cfg_scale: float = 7.0,
     sampler: str = "DPM++ 2M Karras",
     seed: int = -1,
     strength: float = 0.7,
     width: int = 1024,
-    height: int = 1024,
+    height: int = 1024
 ) -> Optional[Image.Image]:
-    """Generate an image from a text prompt using FLUX.1-Krea via the Hugging Face Inference API.
-
-    This function is exposed as an MCP tool via Gradio. It accepts common text-to-image
-    parameters and returns a single generated image.
-
+    """
+    Generate high-quality images using the FLUX.1-Krea-dev model from Hugging Face.
+    
+    This function creates detailed, ultra-high-quality images based on text prompts using 
+    the advanced FLUX.1-Krea-dev diffusion model. The generated images feature ultra detail,
+    ultra elaboration, and perfect quality.
+    
     Args:
-        prompt: The main text prompt describing the desired image. Required.
-        negative_prompt: Text describing what should be avoided in the image.
-        steps: Number of denoising steps (quality vs speed). Typical range 1–100.
-        cfg_scale: Classifier-free guidance scale (adherence to prompt). Typical range 1–20.
-        sampler: Sampling/scheduler method name. If unsupported by the backend, it may be ignored.
-        seed: Random seed. Use -1 for a random seed on each call.
-        strength: Image-to-image strength (0–1). Not used for pure text-to-image but included for compatibility.
-        width: Output image width in pixels. Usually multiples of 32.
-        height: Output image height in pixels. Usually multiples of 32.
-
+        prompt: Text description of the image to generate. Be detailed and descriptive for best results.
+        negative_prompt: Text describing what should NOT appear in the image. Helps avoid unwanted elements.
+        steps: Number of denoising steps (1-100). Higher values generally produce better quality but take longer.
+        cfg_scale: Classifier-free guidance scale (1-20). Higher values follow the prompt more closely.
+        sampler: Sampling method to use. Options: "DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM".
+        seed: Random seed for reproducible results. Use -1 for random seed.
+        strength: Strength of the generation process (0-1). Higher values give more creative freedom.
+        width: Width of generated image in pixels (64-1216, must be multiple of 32).
+        height: Height of generated image in pixels (64-1216, must be multiple of 32).
+    
     Returns:
-        A PIL.Image if generation succeeds; otherwise None.
-
-    Raises:
-        gr.Error: If the backend returns a non-200 HTTP status.
+        PIL Image object of the generated image, or None if generation failed.
     """
-    if not prompt:
+    if prompt == "" or prompt is None:
         return None
 
     key = random.randint(0, 999)
+    
+    API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN")])
+    headers = {"Authorization": f"Bearer {API_TOKEN}"}
 
-    api_token = os.getenv("HF_READ_TOKEN")
-    if not api_token:
-        raise gr.Error("Missing HF_READ_TOKEN environment variable.")
-
-    headers = {"Authorization": f"Bearer {api_token}"}
-
-    # Slightly augment the prompt for higher quality
-    prompt_aug = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
-    print(f'\033[1mGeneration {key}:\033[0m {prompt_aug}')
-
-    # Compute seed value
-    seed_value = seed if seed != -1 else random.randint(1, 1_000_000_000)
-
-    # Prepare payload aligned with HF Inference API conventions
-    parameters = {
-        "num_inference_steps": steps,
-        "guidance_scale": cfg_scale,
-        "width": width,
-        "height": height,
-        "seed": seed_value,
-        # Best-effort: some backends accept scheduler/sampler names; safe to send.
-        "scheduler": sampler,
-    }
-    if negative_prompt and negative_prompt.strip():
-        parameters["negative_prompt"] = negative_prompt
-    # Include strength only if meaningful
-    if 0 < strength <= 1:
-        parameters["strength"] = strength
-
+    # Add some extra flair to the prompt
+    enhanced_prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
+    print(f'\033[1mGeneration {key}:\033[0m {enhanced_prompt}')
+    
+    # Prepare the payload for the API call, including width and height
     payload = {
-        "inputs": prompt_aug,
-        "parameters": parameters,
+        "inputs": enhanced_prompt,
+        "is_negative": False,
+        "steps": steps,
+        "cfg_scale": cfg_scale,
+        "seed": seed if seed != -1 else random.randint(1, 1000000000),
+        "strength": strength,
+        "parameters": {
+            "width": width,  # Pass the width to the API
+            "height": height  # Pass the height to the API
+        }
     }
 
     # Send the request to the API and handle the response
@@ -95,12 +82,12 @@ def query(
         if response.status_code == 503:
             raise gr.Error(f"{response.status_code} : The model is being loaded")
         raise gr.Error(f"{response.status_code}")
-
+    
     try:
         # Convert the response content into an image
         image_bytes = response.content
         image = Image.open(io.BytesIO(image_bytes))
-        print(f'\033[1mGeneration {key} completed!\033[0m ({prompt_aug})')
+        print(f'\033[1mGeneration {key} completed!\033[0m ({enhanced_prompt})')
         return image
     except Exception as e:
         print(f"Error when trying to open the image: {e}")
@@ -118,7 +105,8 @@ css = """
 # 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><h1>FLUX.1-Krea-dev MCP Server</h1></center>")
+    gr.HTML("<center><p>High-quality image generation via Model Context Protocol</p></center>")
     
     # Container for all the UI elements
     with gr.Column(elem_id="app-container"):
@@ -149,22 +137,8 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', css=css) as app:
         with gr.Row():
             image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
         
-        # Bind the button to the query function with the added width and height inputs
-        text_button.click(
-            query,
-            inputs=[
-                text_prompt,
-                negative_prompt,
-                steps,
-                cfg,
-                method,
-                seed,
-                strength,
-                width,
-                height,
-            ],
-            outputs=image_output,
-        )
+        # Bind the button to the flux_krea_generate function with the added width and height inputs
+        text_button.click(flux_krea_generate, inputs=[text_prompt, negative_prompt, steps, cfg, method, seed, strength, width, height], outputs=image_output)
 
-# Launch the Gradio app
+# Launch the Gradio app with MCP server enabled
 app.launch(show_api=True, share=False, mcp_server=True)