Update app.py

app.py

@@ -1,36 +1,30 @@
 import streamlit as st
-from diffusers import StableDiffusionPipeline
-import torch
 from PIL import Image
-from io import BytesIO
+from transformers import AutoProcessor, KandinskyDecoderPipeline
+import torch
 
-# Load the model only once
 @st.cache_resource
 def load_model():
-    model_id = "runwayml/stable-diffusion-v1-5"
-    pipe = StableDiffusionPipeline.from_pretrained(
-        model_id,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        safety_checker=None
+    processor = AutoProcessor.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")
+    model = KandinskyDecoderPipeline.from_pretrained(
+        "kandinsky-community/kandinsky-2-2-decoder",
+        torch_dtype=torch.float32,
     )
-    pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-    return pipe
-
-# Streamlit UI Code
-st.set_page_config(page_title="Thumbnail Generator", layout="centered")
-st.title("🖼️ AI Thumbnail Generator")
+    model.to("cpu")
+    return processor, model
 
-pipe = load_model()
+st.set_page_config(page_title="Thumbnail AI - Kandinsky", layout="centered")
+st.title("🎨 Kandinsky Thumbnail Generator")
 
-prompt = st.text_input("📥 Enter your prompt:", value="Minecraft NOOB vs PRO | bold text | fire background")
+prompt = st.text_input("📝 Enter your thumbnail idea:", "Minecraft Hardcore NOOB vs PRO battle")
 
-if st.button("🚀 Generate"):
-    with st.spinner("Generating thumbnail..."):
-        image = pipe(prompt, guidance_scale=7.5).images[0]
+if st.button("🎬 Generate Image"):
+    with st.spinner("Loading model and generating..."):
+        processor, model = load_model()
+        image = model(prompt=prompt).images[0]
         st.image(image, caption="Generated Thumbnail", use_column_width=True)
+        st.success("✅ Done!")
 
         # 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")
+        img_bytes = image.convert("RGB").tobytes("raw", "RGB")
+        st.download_button("Download Image", data=img_bytes, file_name="thumbnail.jpg", mime="image/jpeg")