import streamlit as st
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image
import io
import os

# Force CPU usage on Spaces
os.environ["CUDA_VISIBLE_DEVICES"] = ""

@st.cache_resource
def load_model():
    pipe = StableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        torch_dtype=torch.float32,
    )
    pipe.to("cpu")  # Force CPU usage
    return pipe

# Streamlit UI
st.title("🎨 AI Image Generator (CPU - Hugging Face Spaces)")

prompt = st.text_input("Enter your prompt:", 
                       "A multi-dimensional futuristic city with glowing lights, fractals, 8K")

guidance = st.slider("Creativity (Guidance Scale)", 1.0, 20.0, 7.5)

if st.button("Generate"):
    with st.spinner("Generating image on CPU... please wait ⏳"):
        pipe = load_model()
        image = pipe(prompt, guidance_scale=guidance).images[0]

        st.image(image, caption="Generated Image", use_column_width=True)

        buf = io.BytesIO()
        image.save(buf, format="PNG")
        st.download_button("Download Image", buf.getvalue(), "generated.png", "image/png")