import gradio as gr
import replicate
import os
import requests
from PIL import Image
from io import BytesIO

examples = [
    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
    "An astronaut riding a green horse",
    "A delicious ceviche cheesecake slice",
]

css="""
#col-container {
    margin: 0 auto;
    max-width: 640px;
}
"""

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
    )
    
    # Fetch the image from the URL
    response = requests.get(output_url)
    image = Image.open(BytesIO(response.content))
    
    # Return the image
    return image


with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""
        # FLUX 1.1 Pro Text-to-Image Generator
        """)
        
        with gr.Row():
            with gr.Column():
                api_key = gr.Text(
                    label="Replicate API Key",
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your Replicate API key...",
                    container=False,
                    type="password",
                )
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    container=False,
                )
            
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Image(label="Result", show_label=False)

        
        gr.Examples(
            examples = examples,
            inputs = [prompt]
        )
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn = generate_image,
        inputs = [prompt, api_key],
        outputs = [result,]
    )

demo.queue().launch()