# app.py

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

# Load model on CPU
@st.cache_resource
def load_model():
    pipe = StableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",  # You can change to another model if needed
        torch_dtype=torch.float32          # Use float32 on CPU
    )
    return pipe.to("cpu")

# Streamlit UI
st.title("🧠 AI Image Generator (CPU-Compatible)")
st.markdown("Generate original, multidimensional images using Stable Diffusion — no GPU required!")

# Prompt and settings
prompt = st.text_area("Enter your creative prompt:", 
                      "A multi-dimensional alien city with glowing fractals, floating geometry, cosmic lighting, 8K resolution")

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

# Generate button
if st.button("Generate Image"):
    with st.spinner("Generating image. This may take a few minutes on CPU..."):
        pipe = load_model()
        image = pipe(prompt, guidance_scale=guidance).images[0]

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

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