app.py

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!")