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

# Streamlit UI configuration - This must be the first Streamlit command
st.set_page_config(page_title="AI Image Variation Generator", layout="wide")

# Set device (CPU/GPU) and appropriate dtype
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32

# Load the Stable Diffusion XL model with mixed precision if on GPU
sd_pipe = StableDiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch_dtype
)
sd_pipe.to(device)

# Enable memory-efficient attention if available; fallback to attention slicing
try:
    sd_pipe.enable_xformers_memory_efficient_attention()
except Exception as e:
    st.warning("xFormers memory efficient attention not available. Falling back to attention slicing.")
    sd_pipe.enable_attention_slicing()

st.title("🎨 AI Image Variation Generator")
st.write("Upload an image and get AI-generated variations based on your style instructions!")

uploaded_file = st.file_uploader("Upload an Image", type=["png", "jpg", "jpeg"])
prompt = st.text_input("Enter style modification instructions (optional):")

if uploaded_file:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="Uploaded Image", use_column_width=True)

    if st.button("Generate Variations"):
        with st.spinner("Generating variations..."):
            variations = []
            # Use torch.no_grad() to reduce memory overhead during inference
            with torch.no_grad():
                for i in range(4):  # Generate 4 variations
                    seed = random.randint(0, 10000)
                    generator = torch.manual_seed(seed)
                    # Generate image using the pipeline with specified prompt and steps
                    img = sd_pipe(prompt=prompt, num_inference_steps=30, generator=generator).images[0]
                    filename = f"variation_{i+1}.png"
                    img.save(filename)
                    variations.append((img, filename))
                    # Clear GPU cache to free up memory between iterations
                    if device == "cuda":
                        torch.cuda.empty_cache()

            # Display the generated variations in two columns
            col1, col2 = st.columns(2)
            for idx, (img, filename) in enumerate(variations):
                with (col1 if idx % 2 == 0 else col2):
                    st.image(img, caption=f"Variation {idx+1}", use_column_width=True)
                    # Provide a download button for each variation
                    with open(filename, "rb") as file:
                        st.download_button("Download", file.read(), file_name=filename)

    if st.button("Regenerate"):
        st.experimental_rerun()