# app.py

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

# Load model on CPU only
@st.cache_resource
def load_model():
    pipe = StableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        torch_dtype=torch.float32,             # CPU needs float32
        use_auth_token=True                    # If needed for private models
    )
    pipe.to("cpu")                             # Ensure it's forced to CPU
    return pipe

# UI
st.title("🎨 AI Image Generator (CPU Compatible)")
st.markdown("No GPU? No problem. This app runs Stable Diffusion entirely on your CPU.")

prompt = st.text_area("Enter your prompt:", 
                      "A surreal multi-dimensional alien forest with glowing trees and floating rocks, 8K")

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

if st.button("Generate Image"):
    with st.spinner("Generating image on CPU... this may take 2-5 minutes."):
        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")