import gradio as gr
from diffusers import DiffusionPipeline
import torch

# Load the FLUX.1-dev model from Hugging Face
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

# Define a function to generate an image based on the prompt
def generate_image(prompt: str):
    image = pipe(prompt).images[0]
    return image

# Create the Gradio interface
iface = gr.Interface(fn=generate_image, 
                     inputs="text", 
                     outputs="image", 
                     title="Text-to-Image with FLUX.1-dev",
                     description="Enter a prompt to generate an image using the FLUX.1-dev model.")

# Launch the app
iface.launch()