Update app.py

app.py

@@ -2,8 +2,12 @@ import torch
 import streamlit as st
 from diffusers import DiffusionPipeline
 
-# Load the pre-trained Stable Diffusion model
+# Hugging Face Spaces may not always have GPU, so we'll check if it's available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load the pre-trained Stable Diffusion model and move it to the correct device
 pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
+pipe.to(device)
 
 # Title and description
 st.markdown('<div class="title">🔥 Memeify - Text-to-Meme Generator 🔥</div>', unsafe_allow_html=True)
@@ -23,3 +27,4 @@ if st.button("Generate Meme", key="generate", help="Click to generate a meme fro
         st.image(image, caption=f"Meme: {input_text}", use_column_width=True)
     else:
         st.warning("Please enter a description to generate a meme.")
+