import gradio as gr
from gradio_client import Client

# Connect to the existing FLUX.1-dev space
flux_client = Client("black-forest-labs/FLUX.1-dev")

def flux_dev_generate(
    prompt: str, 
    seed: int = 0, 
    randomize_seed: bool = True,
    width: int = 1024, 
    height: int = 1024,
    guidance_scale: float = 3.5, 
    num_inference_steps: int = 28
) -> str:
    """Generate high-quality images using FLUX.1-dev model.
    
    Args:
        prompt: Text description of the image to generate
        seed: Random seed for reproducibility
        randomize_seed: Whether to randomize the seed
        width: Image width in pixels
        height: Image height in pixels  
        guidance_scale: How closely to follow the prompt (1.0-20.0)
        num_inference_steps: Quality vs speed tradeoff (1-50)
    
    Returns:
        Generated image file
    """
    result = flux_client.predict(
        prompt=prompt,
        seed=seed,
        randomize_seed=randomize_seed,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        api_name="/infer"
    )
    return result

demo = gr.Interface(
    fn=flux_dev_generate,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Describe your image..."),
        gr.Number(label="Seed", value=0),
        gr.Checkbox(label="Randomize Seed", value=True),
        gr.Slider(256, 2048, value=1024, label="Width"),
        gr.Slider(256, 2048, value=1024, label="Height"),
        gr.Slider(1.0, 20.0, value=3.5, label="Guidance Scale"),
        gr.Slider(1, 50, value=28, label="Inference Steps")
    ],
    outputs="image",
    title="FLUX.1-dev MCP Wrapper"
)

# This makes it accessible to MCP clients!
demo.launch(mcp_server=True)