import gradio as gr
import torch
from diffusers import FluxPipeline

# Load the model directly into your Space
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.to("cuda")

def generate_image(prompt):
    image = pipe(
        prompt, 
        height=1024, 
        width=1024, 
        guidance_scale=3.5, 
        num_inference_steps=28
    ).images[0]
    return image

# Build a quick Gradio Interface
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter your prompt"),
    outputs=gr.Image(label="Generated Image"),
    title="FLUX.1-dev Generator"
)

interface.launch()