import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

# Load the Stable Diffusion v1.5 model (on CPU with float32)
model_id = "sd-legacy/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
pipe = pipe.to("cpu")

# Image generation function
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Gradio UI setup
demo = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(lines=2, placeholder="Describe your image...", label="Prompt"),
    outputs=gr.Image(type="pil", label="Generated Image"),
    title="🎨 Stable Diffusion v1.5 (CPU)",
    description="Enter a prompt and generate an image using Stable Diffusion v1.5 on CPU.",
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch()