Update app.py

app.py

@@ -1,30 +1,25 @@
 import streamlit as st
-from diffusers import StableDiffusionPipeline
+from transformers import CLIPTextModel, CLIPTokenizer
 import torch
+from torch import autocast
+from diffusers import StableDiffusionPipeline
 
-# Function to load the model
-@st.cache_resource
+# Load the model
+@st.cache(allow_output_mutation=True)
 def load_model():
     model_id = "CompVis/stable-diffusion-v1-4"
-    pipe = StableDiffusionPipeline.from_pretrained(model_id)
-    pipe.to("cpu")  # Ensure the model runs on CPU
+    device = "cpu"
+    pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True)
+    pipe.to(device)
     return pipe
 
-# Load the model
 pipe = load_model()
 
-# Streamlit UI
-st.title("Text-to-Image Generator")
-st.write("Enter a description, and the AI model will generate an image!")
-
-# User input
-text_input = st.text_input("Enter your description:")
+# Streamlit app
+st.title("Text to Image Generator")
+prompt = st.text_input("Enter your prompt:", "A scenic landscape of mountains during sunset")
 
-# Generate and display image
-if text_input:
-    with st.spinner("Generating image... Please wait ⏳"):
-        try:
-            image = pipe(text_input).images[0]  # Generate image
-            st.image(image, caption="Generated Image", use_column_width=True)
-        except Exception as e:
-            st.error(f"Error generating image: {e}")
+if st.button("Generate Image"):
+    with st.spinner("Generating..."):
+        image = pipe(prompt)["sample"][0]
+        st.image(image, caption=prompt)