Update app.py

app.py

@@ -2,64 +2,18 @@ from diffusers import StableDiffusionPipeline
 import torch
 import streamlit as st
 
-# Function to load the model and cache it for better performance
-@st.cache_resource
-def load_model():
-    model_id = "CompVis/stable-diffusion-v1-4"
-    device = "cuda" if torch.cuda.is_available() else "cpu"  # Check for GPU, otherwise use CPU
-    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-    pipe = pipe.to(device)
-    return pipe
 
-# Load model
-pipe = load_model()
+# Load the model
+model_id = "CompVis/stable-diffusion-v1-4"
+device = "cuda"
 
-# Set up Streamlit page configuration
-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("""
-    <style>
-        .title {
-            font-size: 3em;
-            font-weight: bold;
-            color: #FF6347;
-            text-align: center;
-            padding: 20px;
-        }
-        .description {
-            font-size: 1.2em;
-            text-align: center;
-            color: #666;
-        }
-        .input-box {
-            font-size: 1.5em;
-            padding: 10px;
-            border-radius: 8px;
-            border: 2px solid #FF6347;
-            margin: 10px auto;
-            display: block;
-            width: 60%;
-            background-color: #fff3e3;
-        }
-        .generate-button {
-            background-color: #FF6347;
-            color: white;
-            font-size: 1.2em;
-            padding: 10px 30px;
-            border: none;
-            border-radius: 8px;
-            cursor: pointer;
-        }
-        .generate-button:hover {
-            background-color: #FF4500;
-        }
-        .footer {
-            text-align: center;
-            font-size: 1.2em;
-            color: #888;
-            padding: 30px;
-        }
     </style>
 """, unsafe_allow_html=True)
 
@@ -67,26 +21,24 @@ st.markdown("""
 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)
 
-# Input for meme description
+# User input for meme description
 input_text = st.text_input("Enter meme description:", key="input", placeholder="e.g., A funny dog looking confused.")
 
-# Button to generate meme
+# 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:
-        # Modify prompt for meme generation
         modified_prompt = "MEME of an Image of " + input_text
 
-        # Generate image with fewer inference steps for faster generation
-        image = pipe(modified_prompt, num_inference_steps=5).images[0]
-
-        # Resize the image for faster loading and display
-        image = image.resize((512, 512))
+        # Show a loading spinner while the image is being generated
+        with st.spinner('Generating your meme...'):
+            # Generate image from the prompt
+            image = pipe(modified_prompt, num_inference_steps=10).images[0]
 
-        # Display the generated image
-        st.image(image, caption="Your Generated Meme", use_column_width=True)
+            # Display the generated image
+            st.image(image, caption="Your Generated Meme", use_column_width=True)
 
-        # Save the image
-        image.save("generated_meme.png")
+            # Save the image
+            image.save("generated_meme.png")
     else:
         st.warning("Please enter a description to generate a meme.")