Update app.py

app.py

@@ -2,13 +2,17 @@ import streamlit as st
 from diffusers import DiffusionPipeline
 import torch
 from PIL import Image
+from huggingface_hub import login
+
+# Optional: Log in to Hugging Face (if using private models)
+# login(token="your_huggingface_token")
 
 # Streamlit app title
 st.title("Stable Diffusion Image Generator")
 
-# Load the diffusion pipeline
+# Load the diffusion pipeline (make sure you have access to the model)
 pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float32)
-pipe.to("cpu")  # Use CPU for inference
+pipe.to("cpu")  # Ensure it's using CPU
 
 # Get user input (prompt)
 prompt = st.text_input("What do you want to see?", "A beautiful landscape")
@@ -16,10 +20,14 @@ prompt = st.text_input("What do you want to see?", "A beautiful landscape")
 # Button to generate image
 if st.button('Generate Image'):
     with st.spinner('Generating...'):
-        # Generate image
-        image = pipe(prompt).images[0]
-        
-        # Display image in Streamlit
-        st.image(image, caption="Generated Image", use_column_width=True)
+        try:
+            # Generate image
+            image = pipe(prompt).images[0]
+
+            # Display image in Streamlit
+            st.image(image, caption="Generated Image", use_column_width=True)
+
+            st.success("Image generated successfully!")
+        except Exception as e:
+            st.error(f"An error occurred: {str(e)}")
 
-    st.success("Image generated successfully!")