app.py

import streamlit as st
from diffusers import AutoPipelineForInpainting, StableDiffusionInpaintingPipeline
from diffusers.utils import load_image
import torch
from PIL import Image

st.title("Image Inpainting with Stable Diffusion")

# Model Loading
@st.cache_resource
def load_model():
    return AutoPipelineForInpainting.from_pretrained(
        "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16"
    ).to("cuda")

pipe = load_model()

# Input Fields
img_url = st.text_input("Enter image URL:", "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png")
mask_url = st.text_input("Enter mask image URL:", "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png")
prompt = st.text_input("Describe the change:", "a tiger sitting on a park bench")

# Parameters
guidance_scale = st.slider("Guidance Scale", 0.0, 10.0, 8.0, 0.5)
strength = st.slider("Strength", 0.0, 1.0, 0.99, 0.01)
num_steps = st.slider("Inference Steps", 10, 50, 20)

if st.button("Run Inpainting"):
    if not all([img_url, mask_url, prompt]):
        st.warning("Please fill all input fields.")
    else:
        try:
            image = load_image(img_url).resize((1024, 1024))
            mask_image = load_image(mask_url).resize((1024, 1024))
            generator = torch.Generator(device="cuda").manual_seed(0)  # Optional for reproducibility

            with torch.no_grad():
                output_image = pipe(
                    prompt=prompt,
                    image=image,
                    mask_image=mask_image,
                    guidance_scale=guidance_scale,
                    num_inference_steps=num_steps,
                    strength=strength, 
                    generator=generator
                ).images[0]

            st.image([Image.open(img_url), Image.open(mask_url), output_image], caption=["Original Image", "Mask", "Result"])
        except Exception as e:
            st.error(f"An error occurred: {e}")