repaint_person.py

import os
import json
import base64
import io
import requests
from PIL import Image
from animeinsseg import AnimeInsSeg
import numpy as np
import argparse
import traceback
import sys
from animeinsseg.inpainting import patch_match
import cv2
import os.path as osp
from omegaconf import OmegaConf
from utils.io_utils import find_all_imgs, submit_request, img2b64, save_encoded_image


INPAINTING_FILL_METHODS = ['fill', 'original', 'latent_noise', 'latent_nothing']

def repaint_img(args, img_path, detector: AnimeInsSeg):

    prompt_default = 'masterpiece, best quality, '
    if args.prompt:
        prompt_default = prompt_default + ',' + args.prompt

    nprompt = "lowres, ((bad anatomy)), ((bad hands)), missing finger, extra digits, fewer digits, ((mutated hands and fingers)), (poorly drawn face), ((mutation)), ((deformed face)), (ugly), ((bad proportions)), ((extra limbs)), extra face, (double head), (extra head), ((extra feet)), cropped, worst quality, low quality, normal quality, jpeg, humpbacked, long body, long neck, ((jpeg artifacts))"
    if args.negative_prompt:
        nprompt = nprompt + ", " + args.negative_prompt

    assert args.fill in INPAINTING_FILL_METHODS, \
        f'Fill method must be one of {INPAINTING_FILL_METHODS}.'

    resolution = args.resolution
    img = Image.open(img_path).convert('RGB')
    imgname = osp.basename(img_path)
    if img.height > img.width:
        W = resolution
        H = (img.height / img.width * resolution) // 32 * 32
        H = int(H)
    else:
        H = resolution
        W = (img.width / img.height * resolution) // 32 * 32 
        W = int(W)

    if args.to_grey:
        img = img.convert('L').convert('RGB')

    options_shared = {
        'width': W,
        'height': H,
        'steps': args.steps,
        'cfg_scale': args.cfg_scale,
        'sample_name': args.sample_name,
        'denoising_strength': args.denoising_strength,
        "alwayson_scripts": {
            "controlnet": {
            "args": [
                {
                    "input_image": "",
                    "module": "lineart_anime",
                    "model": "control_v11p_sd15s2_lineart_anime [3825e83e]",
                    "weight": 1,
                    "resize_mode": "Inner Fit (Scale to Fit)",
                    "lowvram": False,
                    "processor_res": resolution,
                    "threshold_a": 64,
                    "threshold_b": 64,
                    "guidance": 1,
                    "guidance_start": 0,
                    "guidance_end": 1,
                    "guessmode": False,
                    "pixel_perfect": True
                }
            ]
            }
        },
    }

    img = img.resize((W, H), resample=Image.Resampling.LANCZOS)

    if not args.onebyone:
        # detected tags got mixed during the rendering processing
        # imgbgr = cv2.cvtColor(np.array(img), cv2.COLOR_BGR2RGB)
        # tags, character_tags = detector.tagger.label_cv2_bgr(imgbgr)
        prompt = prompt_default
        # prompt = prompt_default
        
        img_b64 = img2b64(img)
        options_shared['alwayson_scripts']['controlnet']['args'][0]['input_image'] = img_b64
        data = {
            "init_images": [img_b64],
            "prompt": prompt,
            "negative_prompt": nprompt,
            **options_shared
        }
        data = json.dumps(data)

        print('runing default cldm ...')
        print(f'prompt: {prompt} \n negative prompt: {nprompt}')
        response = submit_request(args.url, data)
        output_img_b64 = response.json()['images'][0]

        save_encoded_image(output_img_b64, osp.join(save_dir, 'repaint-default-' + imgname+'.png'))
        return

    instances = detector.infer(img_path, output_type='numpy', infer_grey=args.to_grey)
    instances.remove_duplicated()
    im_cv = cv2.imread(img_path)
    ins_drawed = instances.draw_instances(im_cv, draw_bbox=False, mask_alpha=0.35, draw_tags=False)
    cv2.imwrite(osp.join(save_dir, imgname+'_instances.png'), ins_drawed)

    mask_fg = None
    masks = []
    if not instances.is_empty:
        for mask in instances.masks:
            mask = cv2.resize(mask.astype(np.uint8) * 255, (W, H), interpolation=cv2.INTER_AREA)
            mask = Image.fromarray(mask)
            masks.append(mask)
        mask_fg = cv2.resize(instances.compose_masks().astype(np.uint8) * 255, (W, H), interpolation=cv2.INTER_AREA)

    if mask_fg is not None:
        bg = patch_match.inpaint(np.array(img), mask_fg, patch_size=3)
        bg_tags, character_tags = detector.tagger.label_cv2_bgr(cv2.cvtColor(bg, cv2.COLOR_BGR2RGB))
        bg_prompt = ','.join(bg_tags) + ',' + prompt_default
        bg_nprompt = nprompt + ',' + args.bg_nprompt
        bg_prompt = bg_prompt.replace('portrait', '')
        bg = Image.fromarray(bg)
        bg_b64 = img2b64(bg)

        options_shared['alwayson_scripts']['controlnet']['args'][0]['input_image'] = bg_b64
        data = {
            "init_images": [bg_b64],
            "prompt": bg_prompt,
            "negative_prompt": bg_nprompt,
            **options_shared
        }
        data = json.dumps(data)

        print('fg...')
        print(f'bg prompt: {bg_prompt} \nbg negative prompt: {bg_nprompt}')
        response = submit_request(args.url, data)
        output_img_b64 = response.json()['images'][0]

        if args.save_intermediate:
            save_encoded_image(output_img_b64, osp.join(save_dir, 'repaint-bg-' + imgname+'.png'))
        bg_repainted = Image.open(io.BytesIO(base64.b64decode(output_img_b64)))
        img_bg_repainted = Image.composite(img, bg_repainted, Image.fromarray(mask_fg))
        # img_bg_repainted.save(osp.join(save_dir, 'repaint-bg-composed-' + imgname+'.png'))
        img = img_bg_repainted
    
    if instances.is_empty:
        return

    img_b64 = None
    num_fg = len(masks)
    for ii, mask in enumerate(masks):
        if img_b64 is None:
            img_b64 = img2b64(img)
        mask_b64 = img2b64(mask)
        tags = instances.tags[ii]
        prompt = tags.replace(' ', ',') + ',' + prompt_default
        # prompt = prompt_default

        options_shared['alwayson_scripts']['controlnet']['args'][0]['input_image'] = img_b64
        request = {
            "init_images": [img_b64],
            "mask": mask_b64,
            "prompt": prompt,
            "negative_prompt": nprompt,
            'mask_blur': args.mask_blur,
            'inpainting_fill': INPAINTING_FILL_METHODS.index(args.fill),
            'inpaint_full_res': args.inpaint_full_res,
            'inpaint_full_res_padding': args.inpaint_full_res_padding,
            **options_shared
        }

        data = json.dumps(request)
        print(f'runing fg repainting...{ii+1}/{num_fg}')
        print(f'fg prompt: {prompt} \n fg negative prompt: {nprompt}')
        response = submit_request(args.url, data)
        img_b64 = response.json()['images'][0]
        if args.save_intermediate or ii == num_fg - 1:
            save_encoded_image(img_b64, osp.join(save_dir, f'repaint-fg{ii}-' + imgname+'.png'))
        mask.save(osp.join(save_dir, f'repaint-fg{ii}-mask-' + imgname+'.png'))

if __name__ == '__main__':
    
    parser = argparse.ArgumentParser(description='Inpaint instances of people using stable '
                                                 'diffusion.')
    parser.add_argument('--img_path', type=str, default='', help='Path to input image.')
    parser.add_argument('--onebyone', type=bool, default=True, help='repainting person one by one')
    parser.add_argument('-p', '--prompt', type=str, default='',
                        help='Stable diffusion prompt to use.')
    parser.add_argument('-n', '--negative_prompt', type=str, default='',
                        help='Stable diffusion negative prompt.')
    parser.add_argument('-W', '--width', type=int, default=768, help='Width of output image.')
    parser.add_argument('-H', '--height', type=int, default=768, help='Height of output image.')
    parser.add_argument('-s', '--steps', type=int, default=24, help='Number of diffusion steps.')
    parser.add_argument('-c', '--cfg_scale', type=int, default=9, help='Classifier free guidance '
                        'scale, i.e. how strongly the image should conform to prompt.')
    parser.add_argument('-S', '--sample_name', type=str, default='Euler a', help='Name of sampler '
                        'to use.')
    parser.add_argument('-d', '--denoising_strength', type=float, default=0.75, help='How much to '
                        'disregard original image.')
    parser.add_argument('-f', '--fill', type=str, default=INPAINTING_FILL_METHODS[1],
                        help='The fill method to use for inpainting.')
    parser.add_argument('-b', '--mask_blur', type=int, default=4, help='Blur radius of Gaussian '
                        'filter to apply to mask.')
    parser.add_argument('-r', '--resolution', type=int, default=640, help='inpainting resolution')
    parser.add_argument('--save_dir', type=str, default='repaint_output')
    parser.add_argument('--url', type=str, default='http://127.0.0.1:7860/sdapi/v1/img2img', help='img2img url')
    parser.add_argument('--cfg', type=str, default='', help='repaint config path')
    parser.add_argument('--bg_nprompt', type=str, default='((person)), character, 1girl, 1boy', help='background negative prompt')
    parser.add_argument('--inpaint_full_res', type=int, default=1)
    parser.add_argument('--inpaint_full_res_padding', type=int, default=32)
    parser.add_argument('--detector_ckpt', type=str, default='models/AnimeInstanceSegmentation/rtmdetl_e60.ckpt')
    parser.add_argument('--save_intermediate', type=bool, default=False)
    parser.add_argument('--to-grey', type=bool, default=False)
    parser.add_argument('--infer-tagger', type=bool, default=True)
    args = parser.parse_args()

    if osp.exists(args.cfg):
        args = OmegaConf.create(vars(args))
        args.merge_with(OmegaConf.load(args.cfg))
        print(args)

    save_dir = args.save_dir
    if not osp.exists(save_dir):
        os.makedirs(save_dir)

    detector = AnimeInsSeg(args.detector_ckpt, device='cuda')
    detector.init_tagger()
    imgp_list = args.img_path
    if isinstance(imgp_list, str):
        if osp.isdir(imgp_list):
            imgp_list = find_all_imgs(imgp_list, abs_path=True)
        else:
            imgp_list = [imgp_list]
    for imgp in imgp_list:
        repaint_img(args, imgp, detector)