Update app.py

app.py

@@ -1,43 +1,31 @@
+import streamlit as st
+from transformers import pipeline
 from diffusers import StableDiffusionPipeline
 import torch
-from PIL import Image
-import streamlit as st
-
-# Load the Stable Diffusion pipeline
-@st.cache(allow_output_mutation=True)
-def load_pipeline():
-    pipeline = StableDiffusionPipeline.from_pretrained(
-        "CompVis/stable-diffusion-v1-4",
-        torch_dtype=torch.float32  # Use float32 for CPU support
-    )
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    pipeline.to(device)
-    return pipeline
-
-def main():
-    st.title("Stable Diffusion Image Generator")
-    st.write("Generate images from text prompts using Stable Diffusion")
-
-    # Initialize the pipeline
-    pipeline = load_pipeline()
-
-    # Text input for the prompt
-    prompt = st.text_input("Enter your text prompt", "")
-
-    # Generate button
-    if st.button("Generate"):
-        if not prompt:
-            st.warning("Please enter a prompt first.")
-            return  # ✅ Make sure this return is correctly indented
-
-        st.write("Generating your image...")
-        with torch.no_grad():
-            result = pipeline(prompt)
-            image = result.images[0]  # Extract the generated image
-
-        st.write("Generated Image:")
-        st.image(image, use_column_width=True)
 
-if __name__ == "__main__":
-    main()
+# Load the pre-trained Stable Diffusion model
+@st.cache_resource
+def load_model():
+    model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v-1-4-original", torch_dtype=torch.float16)
+    model.to("cuda")
+    return model
+
+# Initialize model
+model = load_model()
+
+# Streamlit Interface
+st.title("Text to Image Generator using Stable Diffusion")
+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", "A beautiful landscape with mountains and lakes.")
+
+if st.button("Generate Image"):
+    if user_input:
+        with st.spinner("Generating image..."):
+            # Generate image based on user input
+            generated_image = model(user_input).images[0]
+            st.image(generated_image)
+    else:
+        st.error("Please enter a text prompt!")