from diffusers import DiffusionPipeline
import torch
import streamlit as st

# Model paths for Hugging Face
sdxl_base_model_path = "stabilityai/stable-diffusion-xl-base-1.0"
sdxl_refiner_model_path = "stabilityai/stable-diffusion-xl-refiner-1.0"

@st.cache_resource
def load_pipeline():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    try:
        # Load the base pipeline
        pipe = DiffusionPipeline.from_pretrained(
            sdxl_base_model_path,
            torch_dtype=torch.float16 if device == "cuda" else torch.float32
        )
        pipe.to(device)
        return pipe
    except ValueError as e:
        st.error(f"Error loading the pipeline: {str(e)}")
        return None


def image_generation(pipe, prompt):
    if not pipe:
        return None
    try:
        # Generate the image
        image = pipe(
            prompt=prompt,
            negative_prompt="blurred, ugly, watermark, low resolution",
            num_inference_steps=20,
            guidance_scale=9.0
        ).images[0]
        return image
    except Exception as e:
        st.error(f"Error generating image: {str(e)}")
        return None


# Streamlit app interface
st.title("Project 11: Image Generation using SD XL")
prompt = st.text_input("Enter your prompt", value="A futuristic superhero cat")

pipeline = load_pipeline()

if pipeline and st.button("Generate Image"):
    with st.spinner("Generating your Image..."):
        image = image_generation(pipeline, prompt)
        if image:
            st.image(image)