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

# Set up title and description
st.title("Insert Yourself Into Historical Photos with Stable Diffusion!")
st.write("Upload a historical photo and a mask, then describe how you'd like to place yourself into the scene.")

# Load model (for inpainting)
@st.cache_resource
def load_pipeline():
    pipe = StableDiffusionInpaintPipeline.from_pretrained(
        "stabilityai/stable-diffusion-2-inpainting",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
    )
    pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
    return pipe

pipe = load_pipeline()

# File uploads for base image and mask
base_image = st.file_uploader("Upload the Historical Photo", type=["jpg", "jpeg", "png"])
mask_image = st.file_uploader("Upload the Mask (Black & White Image)", type=["jpg", "jpeg", "png"])

# Prompt input
prompt = st.text_input("Describe the scene. How do you want to place yourself?")

if base_image and mask_image and prompt:
    # Display input images
    st.image(base_image, caption="Historical Photo", use_column_width=True)
    st.image(mask_image, caption="Mask", use_column_width=True)

    # Load images using PIL
    base_image = Image.open(base_image).convert("RGB")
    mask_image = Image.open(mask_image).convert("RGB")

    # Generate inpainting
    st.write("Generating the image... This may take a moment.")
    result = pipe(prompt=prompt, image=base_image, mask_image=mask_image).images[0]

    # Display result
    st.image(result, caption="Inpainted Image", use_column_width=True)
    st.write("Right-click on the image to download!")