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

# 1. Page Config
st.set_page_config(page_title="AI Image Generator")
st.title("🎨 AI Text-to-Image Generator")

# 2. Load the Model
@st.cache_resource
def load_model():
    model_id = "runwayml/stable-diffusion-v1-5"
    
    # Check if GPU (CUDA) is available
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Load the pipeline
    if device == "cuda":
        # If GPU, use fast float16 precision
        pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
    else:
        # If CPU, use standard float32 precision
        pipe = StableDiffusionPipeline.from_pretrained(model_id)
    
    pipe.to(device)
    return pipe

st.info("Loading AI Model... The first time may take a few minutes.")
gen_pipe = load_model()

# 3. User Interface
prompt = st.text_input("Describe the image you want to see:", placeholder="A futuristic city with flying cars at sunset")

if st.button("Generate Image"):
    if prompt:
        with st.spinner("AI is painting... please wait (this takes longer on CPU)"):
            # Generate the image
            # num_inference_steps=20 makes it faster for testing; 50 is better quality
            image = gen_pipe(prompt, num_inference_steps=20).images[0]
            
            # Display the image
            st.image(image, caption=f"Generated: {prompt}", use_column_width=True)
            
            # Allow user to download
            image.save("generated_img.png")
            with open("generated_img.png", "rb") as file:
                st.download_button("Download Image", data=file, file_name="ai_image.png", mime="image/png")
    else:
        st.warning("Please enter a prompt first!")