import gradio as gr
import replicate
import os

def generate_image(prompt, api_key):
    # Set the API key for the current session
    os.environ["REPLICATE_API_TOKEN"] = api_key
    
    # Prepare the input for the model
    inputs = {
        "prompt": prompt,
        "prompt_upsampling": True
    }
    
    # Run the model and get the output URL
    output_url = replicate.run(
        "black-forest-labs/flux-1.1-pro",
        input=inputs
    )
    
    # Return the generated image URL
    return output_url[0]

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(
            lines=2,
            placeholder="Enter your prompt here...",
            label="Prompt"
        ),
        gr.Textbox(
            lines=1,
            placeholder="Enter your Replicate API key...",
            type="password",
            label="Replicate API Key"
        )
    ],
    outputs=gr.Image(type="auto"),
    title="FLUX 1.1 Pro Text-to-Image Generator",
    description="Generate images from text prompts using the FLUX 1.1 Pro model."
)

# Launch the Gradio app
iface.launch()