Update app.py

app.py

@@ -2,8 +2,9 @@ import streamlit as st
 from diffusers import StableDiffusionPipeline
 import torch
 from PIL import Image
+from io import BytesIO
 
-# Load model (cached so it doesn't reload every time)
+# Load the model only once
 @st.cache_resource
 def load_model():
     model_id = "runwayml/stable-diffusion-v1-5"
@@ -15,30 +16,21 @@ def load_model():
     pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
     return pipe
 
-# Page config
-st.set_page_config(page_title="AI Thumbnail Generator", layout="centered")
-st.title("🖼️ AI Thumbnail Generator (Stable Diffusion)")
+# Streamlit UI Code
+st.set_page_config(page_title="Thumbnail Generator", layout="centered")
+st.title("🖼️ AI Thumbnail Generator")
 
-# Load model
 pipe = load_model()
 
-# Prompt input
-prompt = st.text_input("📝 Enter Thumbnail Prompt:", value="Minecraft NOOB vs PRO | fire background | bold text")
+prompt = st.text_input("📥 Enter your prompt:", value="Minecraft NOOB vs PRO | bold text | fire background")
 
-# Generate button
-if st.button("🎨 Generate Thumbnail"):
-    with st.spinner("Generating image... Please wait"):
-        result = pipe(prompt, guidance_scale=7.5)
-        image = result.images[0]
-        st.image(image, caption="🎉 Your AI Thumbnail", use_column_width=True)
+if st.button("🚀 Generate"):
+    with st.spinner("Generating thumbnail..."):
+        image = pipe(prompt, guidance_scale=7.5).images[0]
+        st.image(image, caption="Generated Thumbnail", use_column_width=True)
 
-        # Download option
-        st.download_button("📥 Download Image", image_to_bytes(image), file_name="thumbnail.jpg", mime="image/jpeg")
-
-# Utility function
-def image_to_bytes(img: Image.Image):
-    from io import BytesIO
-    buf = BytesIO()
-    img.save(buf, format="JPEG")
-    byte_im = buf.getvalue()
-    return byte_im
+        # Download button
+        buf = BytesIO()
+        image.save(buf, format="JPEG")
+        byte_im = buf.getvalue()
+        st.download_button("Download Thumbnail", byte_im, file_name="thumbnail.jpg", mime="image/jpeg")