app.py

import gradio as gr
import random
import os
from PIL import Image
from typing import Optional
from huggingface_hub import InferenceClient

# Project by Nymbo

API_TOKEN = os.getenv("HF_READ_TOKEN")
timeout = 100

# 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.0,
    sampler: str = "DPM++ 2M Karras",
    seed: int = -1,
    strength: float = 0.7,
    width: int = 1024,
    height: int = 1024
) -> Optional[Image.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: 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:
        PIL Image object of the generated image, or None if generation failed.
    """
    if prompt == "" or prompt is None:
        return None

    key = random.randint(0, 999)
    
    # 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}')
    
    try:
        # Initialize the Hugging Face Inference Client
        # Try different providers in order of preference
        providers = ["auto", "replicate", "fal-ai"]
        
        for provider in providers:
            try:
                client = InferenceClient(
                    api_key=API_TOKEN,
                    provider=provider
                )
                
                # Generate the image using the proper client
                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, 1000000000)
                )
                
                print(f'\033[1mGeneration {key} completed with {provider}!\033[0m ({enhanced_prompt})')
                return image
                
            except Exception as provider_error:
                print(f"Provider {provider} failed: {provider_error}")
                if provider == providers[-1]:  # Last provider
                    raise provider_error
                continue
        
    except Exception as e:
        print(f"Error during image generation: {e}")
        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.")
        elif "503" in str(e):
            raise gr.Error("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

# CSS to style the app
css = """
#app-container {
    max-width: 800px;
    margin-left: auto;
    margin-right: auto;
}
"""

# 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 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"):
        # Add a text input for the main prompt
        with gr.Row():
            with gr.Column(elem_id="prompt-container"):
                with gr.Row():
                    text_prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here", lines=2, elem_id="prompt-text-input")
                
                # 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"])

        # Add a button to trigger the image generation
        with gr.Row():
            text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
        
        # Image output area to display the generated image
        with gr.Row():
            image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
        
        # 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 with MCP server enabled
app.launch(show_api=True, share=False, mcp_server=True)