import streamlit as st
import torch
from diffusers import DiffusionPipeline
from PIL import Image

# Load the pre-trained model and LoRA weights
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")

# Streamlit app
st.title("Memeify")

# 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.")

# 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'))
        
        # Display the generated image
        st.image(pil_image, caption="Generated Meme", use_column_width=True)