#!/usr/bin/env python3
from diffusers import AutoPipelineForInpainting, AutoPipelineForImage2Image, UNet2DConditionModel
from diffusers.utils import load_image
import torch
import hf_image_uploader as hfi

torch.backends.cuda.matmul.allow_tf32 = True
torch_device = "cuda" if torch.cuda.is_available() else "cpu"

# pipe = AutoPipelineForInpainting.from_pretrained("runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant="fp16", use_safetensors=True) 
unet = UNet2DConditionModel.from_pretrained("valhalla/sdxl-inpaint-ema", torch_dtype=torch.float16)
pipe = AutoPipelineForInpainting.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")

pipe = pipe.to(torch_device)

img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
img_url = "https://huggingface.co/datasets/patrickvonplaten/images_1024_1024/resolve/main/jdn_jacques_de_nuce-1024x1024.jpg"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"

image = load_image(img_url).resize((1024, 1024))
mask_image = load_image(mask_url).resize((1024, 1024))

prompt = "dslr photography of an empty bench, high quality"
generator = torch.Generator(device="cuda").manual_seed(0)

image = pipe(
    prompt=prompt,
    image=image,
    mask_image=mask_image,
    guidance_scale=8.0,
    num_inference_steps=20,
    strength=0.90,
    generator=generator,
).images[0]


image = image.resize((1024, 1024))
hfi.upload(image, "patrickvonplaten/images")