app.py

from diffusers.utils import load_image, make_image_grid
import gradio as gr
import cv2
import numpy as np
from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, UniPCMultistepScheduler
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16, use_safetensors=True)
pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True
).to(device)

pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()

def create_mask(img):
  img = cv2.imread(img, cv2.IMREAD_GRAYSCALE)

  # Create a binary mask
  mask = cv2.threshold(img, 254, 255, cv2.THRESH_BINARY)[1]

  # Invert the mask
  inverted_mask = cv2.bitwise_not(mask)
  return mask

def load_and_resize_images(image_path, mask_path, target_size=(512, 512)):
    init_image = load_image(image_path)
    init_image = init_image.resize(target_size)

    mask_image = load_image(mask_path)
    mask_image = mask_image.resize(target_size)

    return init_image, mask_image

def make_inpaint_condition(image, image_mask):
    image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
    image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0

    assert image.shape[0:1] == image_mask.shape[0:1]
    image[image_mask > 0.5] = -1.0  # set as masked pixel
    image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
    image = torch.from_numpy(image)
    return image

def generate_inpainting(init_image, mask_image, prompt, negative_prompt=None):
    control_image = make_inpaint_condition(init_image, mask_image)

    output = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=50,
        eta=1.0,
        image=init_image,
        mask_image=mask_image,
        control_image=control_image,
    ).images[0]

    return output

def process_images_and_generate_inpainting(image_path, mask_path, prompt, negative_prompt=None):
    init_image, mask_image = load_and_resize_images(image_path, mask_path)
    output_image = generate_inpainting(init_image, mask_image, prompt, negative_prompt)
    return output_image

maskGen = gr.Interface(
    fn=create_mask,
    inputs=gr.Image(type='filepath', label="Original Image"),
    outputs=gr.Image(label="Masked Image Image"),
    allow_flagging="never")

inpaintImg = gr.Interface(
    fn=process_images_and_generate_inpainting,
    inputs=[
        gr.Image(type='filepath', label="Original Image"),
        gr.Image(type='filepath', label="Masked Image"),
        gr.Textbox(label="Prompt"),
        gr.Textbox(label="Negative Prompt")
    ],
    outputs=[
        gr.Image(label="Result Image")
    ],
    allow_flagging="never")

demo = gr.TabbedInterface(
    [maskGen, inpaintImg],
    ["Generate Mask", "Make new Image"]
)

demo.launch(share=True)