Update app.py

app.py

@@ -1,35 +1,26 @@
-import torch
 import streamlit as st
-from diffusers import StableDiffusionPipeline
-
-# Check if CUDA (GPU) is available, otherwise use CPU
-device = "cuda" if torch.cuda.is_available() else "cpu"
+import torch
+from diffusers import DiffusionPipeline
+from PIL import Image
 
-# Load the pre-trained Stable Diffusion model and move it to the correct device
-pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
-pipe.to(device)
+# Load the pre-trained model and LoRA weights
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev")
+pipe.load_lora_weights("prithivMLmods/Flux-Meme-Xd-LoRA")
 
-# Title and description
-st.markdown('<div class="title">🔥 Memeify - 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)
+# Streamlit app
+st.title("Memeify")
 
-# User input for meme description
-input_text = st.text_input("Enter meme description:", key="input", placeholder="e.g., dog dancing")
+# Input field for the meme prompt
+prompt = st.text_area("Enter your meme prompt:", "meme, A medium-sized painting of a white T-rex in the middle of a dark, stormy night. The t-rex is facing towards the left side of the frame, its head turned towards the right. Its mouth is open, revealing its sharp teeth. A rooster is standing in the foreground of the painting, with a red cap on its head. The roosters head is turned to the right, and the word \"Remember who you are\" is written in white text above it. The background is a deep blue, with dark gray clouds and a crescent moon in the upper left corner of the image. There are mountains in the background, and a few other animals can be seen in the lower right corner.")
 
-# 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:
-        try:
-            # Modify the input prompt for meme generation
-            modified_prompt = "MEME of an Image of " + input_text
-            
-            # Generate the image from the modified prompt
-            image = pipe(modified_prompt).images[0]
-            
-            # Display the generated image in the Streamlit app
-            st.image(image, caption=f"Meme: {input_text}", use_column_width=True)
+# Button to generate the meme
+if st.button("Generate Meme"):
+    with st.spinner("Generating meme..."):
+        # Generate the image using the provided prompt
+        image = pipe(prompt).images[0]
+        
+        # Convert the image to PIL format
+        pil_image = Image.fromarray(image.astype('uint8'))
         
-        except Exception as e:
-            st.error(f"An error occurred while generating the meme: {e}")
-    else:
-        st.warning("Please enter a description to generate a meme.")
+        # Display the generated image
+        st.image(pil_image, caption="Generated Meme", use_column_width=True)