import streamlit as st
from diffusers import StableDiffusionPipeline
import torch
import os

# Set cache to writable directory
os.environ["HF_HOME"] = "/tmp"
os.environ["TRANSFORMERS_CACHE"] = "/tmp"
os.environ["HF_HUB_CACHE"] = "/tmp"

st.set_page_config(page_title="AI Design Generator", layout="centered")

@st.cache_resource
def load_model():
    model = StableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5", 
        torch_dtype=torch.float32
    )
    model.to("cpu")  # CPU unless GPU enabled in Space settings
    return model

pipe = load_model()

st.title("🧠 AI Design Generator")
prompt = st.text_input("Enter your design prompt", "a modern packaging design with bright colors")

if st.button("Generate Image"):
    with st.spinner("Generating..."):
        image = pipe(prompt).images[0]
        st.image(image, caption="Generated Design", use_column_width=True)