Update app.py

app.py

@@ -1,22 +1,48 @@
 import streamlit as st
-from generator import load_model, generate_image
-from PIL import Image
+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 (CPU Optimized)")
+st.title("🎨 AI Image Generator")
 
-pipe = load_model()
+# Load model
+@st.cache_resource
+def get_pipe():
+    return load_model()
 
-prompt = st.text_input("Enter prompt:", value="a magical landscape with floating islands")
-steps = st.slider("Quality (steps)", min_value=4, max_value=15, value=8)
-guidance = st.slider("Prompt Accuracy (guidance scale)", min_value=0.0, max_value=10.0, value=2.5)
+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="Generated Image", use_column_width=True)
+        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()
-        image.save(buf, format="PNG")
-        st.download_button("Download Image", buf.getvalue(), "generated_image.png", "image/png")
+        img_to_download.save(buf, format="PNG")
+        st.download_button(
+            label="📥 Download Image",
+            data=buf.getvalue(),
+            file_name="generated_image.png",
+            mime="image/png"
+        )