Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,27 +1,26 @@
 import streamlit as st
-import requests
-from PIL import Image
-from io import BytesIO
-import os
+from diffusers import StableDiffusionPipeline
+import torch
 
 st.set_page_config(page_title="AI Design Generator", layout="centered")
-st.title("🎨 AI Poster/Packaging Design Generator")
-st.write("Enter a prompt to generate your custom design:")
 
-prompt = st.text_input("Prompt", "modern minimalist packaging with red and white theme")
+@st.cache_resource
+def load_model():
+    model = StableDiffusionPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5", 
+        torch_dtype=torch.float32
+    )
+    model.to("cpu")  # Space is on CPU unless you enable GPU manually
+    return model
 
-if st.button("Generate Design"):
-    with st.spinner("Generating..."):
-        API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2"
-        headers = {"Authorization": f"Bearer {os.environ['HF_TOKEN']}"}
+pipe = load_model()
 
-        payload = {"inputs": prompt}
-        response = requests.post(API_URL, headers=headers, json=payload)
+st.title("🧠 AI Design Generator")
+prompt = st.text_input("Enter your design prompt", "a modern packaging design with bright colors")
 
-        if response.status_code == 200:
-            image = Image.open(BytesIO(response.content))
-            st.image(image, caption="Generated Design")
-        else:
-            st.error("Failed to generate image. Try again later.")
+if st.button("Generate Image"):
+    with st.spinner("Generating..."):
+        image = pipe(prompt).images[0]
+        st.image(image, caption="Generated Design", use_column_width=True)