import torch
from diffusers import StableDiffusion3Pipeline
import gradio as gr

# Load model (auto CPU/GPU)
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3-medium-diffusers",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
)
pipe.to(device)

def generate(prompt):
    image = pipe(
        prompt,
        negative_prompt="",
        num_inference_steps=28,
        guidance_scale=7.0,
    ).images[0]
    return image

ui = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(label="Enter your prompt"),
    outputs=gr.Image(label="Generated Image"),
    title="Stable Diffusion 3 Medium - HF Space",
)

ui.launch()