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

# Load the diffusion pipeline model
@st.cache_resource
def load_pipeline():
    try:
        pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
        pipe.load_lora_weights("kothariyashhh/GenAi-Texttoimage")
        pipe = pipe.to("cuda")  # Move model to GPU if available
        return pipe
    except Exception as e:
        st.error(f"Error loading model: {e}")
        return None

pipe = load_pipeline()

# Streamlit app
st.title("Text-to-Image Generation App")

# User input for prompt
user_prompt = st.text_input("Enter your image prompt", value="a photo of Yash Kothari with bike")

# Button to generate the image
if st.button("Generate Image"):
    if user_prompt and pipe:
        with st.spinner("Generating image..."):
            try:
                # Generate the image
                image = pipe(user_prompt).images[0]
                
                # Display the generated image
                st.image(image, caption="Generated Image", use_column_width=True)
            except Exception as e:
                st.error(f"Error generating image: {e}")
    else:
        if not pipe:
            st.error("Model not loaded. Please check the logs.")
        else:
            st.error("Please enter a valid prompt.")