import gradio as gr
import torch
from diffusers import FluxPipeline

# Load the model
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power

def generate_image(prompt):
    image = pipe(
        prompt,
        height=1024,
        width=1024,
        guidance_scale=3.5,
        num_inference_steps=50,
        max_sequence_length=512,
        generator=torch.Generator("cpu").manual_seed(0)
    ).images[0]
    return image

# Create the UI
ui = gr.Interface(
    generate_image,
    gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
    "image",
    title="FLUX.1-dev Image Generator",
    description="Generate images using the FLUX.1-dev model.",
)

# Launch the UI
ui.launch()