import streamlit as st
import torch
from diffusers import StableDiffusionXLPipeline
from PIL import Image

# Load the model (CPU version)
@st.cache_resource()
def load_model():
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "cagliostrolab/animagine-xl-4.0",
        torch_dtype=torch.float32,  # Change to float32 for CPU
        use_safetensors=True,
        custom_pipeline="lpw_stable_diffusion_xl",
        add_watermarker=False
    )
    return pipe  # Do NOT move to CUDA

pipe = load_model()

# Streamlit UI
st.title("Text-to-Image Generator (CPU)")

# User input
prompt = st.text_area("Enter your prompt:", "1girl, arima kana, oshi no ko, looking at viewer, smile")
width = st.slider("Width", 512, 1024, 832, step=64)
height = st.slider("Height", 512, 1280, 1216, step=64)
guidance_scale = st.slider("Guidance Scale", 1.0, 10.0, 6.0, step=0.5)
inference_steps = st.slider("Inference Steps", 10, 50, 25, step=5)

if st.button("Generate Image"):
    try:
        with st.spinner("Generating image..."):
            image = pipe(
                prompt,
                width=width,
                height=height,
                guidance_scale=guidance_scale,
                num_inference_steps=inference_steps
            ).images[0]

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

        image_path = "generated_image.png"
        image.save(image_path)

        with open(image_path, "rb") as file:
            st.download_button(label="Download Image", data=file, file_name="generated_image.png", mime="image/png")

    except Exception as e:
        st.error(f"Error generating image: {e}")