Update app.py

app.py

@@ -1,13 +1,15 @@
 from diffusers import StableDiffusionPipeline
 import torch
 import streamlit as st
-from transformers import pipeline
 
-model_id = "runwayml/stable-diffusion-v1-5"
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
-pipe.to("cpu")
+# Load the model
+model_id = "CompVis/stable-diffusion-v1-4"
+device = "cuda"
 
-st.set_page_config(page_title="🔥 MemeForge - Text-to-Meme Generator 🔥", page_icon="🔥", layout="wide")
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pipe = pipe.to(device)
+
+st.set_page_config(page_title="🔥 Memeify - Text-to-Meme Generator 🔥", page_icon="🔥", layout="wide")
 
 # Custom styling for the page
 st.markdown("""
@@ -55,16 +57,19 @@ st.markdown("""
     </style>
 """, unsafe_allow_html=True)
 
+# Title and description
 st.markdown('<div class="title">🔥 MemeForge - Text-to-Meme Generator 🔥</div>', unsafe_allow_html=True)
 st.markdown('<div class="description">Enter a description, and the AI will generate a meme image for you!</div>', unsafe_allow_html=True)
 
+# User input for meme description
 input_text = st.text_input("Enter meme description:", key="input", placeholder="e.g., A funny dog looking confused.")
 
+# Generate meme image when button is clicked
 if st.button("Generate Meme", key="generate", help="Click to generate a meme from your description"):
     if input_text:
         modified_prompt = "MEME of an Image of " + input_text
 
-        # Generate image
+        # Generate image from the prompt
         image = pipe(modified_prompt, num_inference_steps=10).images[0]
 
         # Display the generated image
@@ -74,3 +79,4 @@ if st.button("Generate Meme", key="generate", help="Click to generate a meme fro
         image.save("generated_meme.png")
     else:
         st.warning("Please enter a description to generate a meme.")
+