Update app.py

app.py

@@ -1,30 +1,33 @@
 import streamlit as st
-from diffusers import DiffusionPipeline
+from transformers import pipeline
+from diffusers import StableDiffusionPipeline
 import torch
 
-# Load the pre-trained model
+# Load the pre-trained Stable Diffusion model
 @st.cache_resource
 def load_model():
-    # Load the Stable Diffusion model from Hugging Face
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16)
-    pipe.to("cuda")  # Make sure to move the model to GPU if available
-    return pipe
-
-# Initialize the model
-pipe = load_model()
-
-# Streamlit Interface
-st.title("Text to Image Generator using Stable Diffusion 2.1")
-st.write("Enter a description below, and the model will generate an image based on it.")
-
-# Text input field for user to enter prompt
-user_input = st.text_area("Enter the text prompt", "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k")
-
-if st.button("Generate Image"):
-    if user_input:
-        with st.spinner("Generating image..."):
-            # Generate image based on user input
-            generated_image = pipe(user_input).images[0]
-            st.image(generated_image)
-    else:
-        st.error("Please enter a text prompt!")
+    # You can load any Stable Diffusion model here
+    model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v-1-4-original")
+    model.to("cpu")  # Ensure it uses CPU for inference
+    return model
+
+# Initialize Streamlit components
+st.title("Text-to-Image Generator")
+
+st.write("Enter a description and generate an image!")
+
+# User input for text prompt
+prompt = st.text_input("Enter your text prompt here:")
+
+if prompt:
+    st.write("Generating image... Please wait.")
+
+    # Load model once (from the Hugging Face model hub)
+    model = load_model()
+
+    # Generate the image from the text prompt
+    with torch.no_grad():
+        image = model(prompt).images[0]
+
+    # Show the generated image in the app
+    st.image(image, caption="Generated Image", use_column_width=True)