import streamlit as st
from generator import load_model, generate_image, upscale_image
import io

# Set Streamlit page configuration
st.set_page_config(page_title="AI Image Generator", layout="centered")
st.title("🎨 AI Image Generator")

# Load model
@st.cache_resource
def get_pipe():
    return load_model()

pipe = get_pipe()

# Prompt input
prompt = st.text_input("Enter your prompt", "A scenic landscape with mountains and a river")

# Optional parameters
with st.expander("⚙️ Advanced Options"):
    steps = st.slider("Inference Steps", 1, 50, 8)
    guidance = st.slider("Guidance Scale", 1.0, 20.0, 2.5)
    upscale = st.checkbox("🔍 Upscale Image (2x resolution)", value=True)

# Generate button
if st.button("Generate"):
    with st.spinner("Generating image..."):
        # Generate image from prompt
        image = generate_image(pipe, prompt, steps, guidance)
        st.image(image, caption="🖼️ Original Image", use_column_width=True)

        # Upscale if selected
        if upscale:
            upscaled = upscale_image(image)
            st.image(upscaled, caption="🔍 Upscaled Image", use_column_width=True)
            img_to_download = upscaled
        else:
            img_to_download = image

        # Prepare image for download
        buf = io.BytesIO()
        img_to_download.save(buf, format="PNG")
        st.download_button(
            label="📥 Download Image",
            data=buf.getvalue(),
            file_name="generated_image.png",
            mime="image/png"
        )