import os
import torch
from diffusers import StableDiffusionPipeline
import gradio as gr

# Store Hugging Face cache in /tmp so it resets each run (saves storage space)
os.environ["HF_HOME"] = "/tmp"

# Use a lighter model (stable-diffusion-turbo) for faster + smaller downloads
model_id = "stabilityai/sd-turbo"

# Load model
pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16
)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

def generate(prompt):
    image = pipe(prompt).images[0]
    return image

# Gradio UI
iface = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(label="Enter your prompt"),
    outputs=gr.Image(type="pil"),
    title="Stable Diffusion Demo",
    description="Enter a text prompt and generate an AI image!"
)

iface.launch()