import os import json import base64 import io import requests from PIL import Image import numpy as np import argparse import random from tqdm import tqdm import os.path as osp from omegaconf import OmegaConf from utils.io_utils import find_all_imgs, submit_request, img2b64, save_encoded_image from random import randint from requests.auth import HTTPBasicAuth import cv2 from copy import deepcopy from pathlib import Path import math INPAINTING_FILL_METHODS = ['fill', 'original', 'latent_noise', 'latent_nothing'] def run_sdinpaint(img: Image.Image, mask: Image.Image, data: dict, prompt: str = '', nprompt: str = '', url='', auth=None) -> str: if isinstance(img, Image.Image): img_b64 = img2b64(img) else: assert isinstance(img, str) img_b64 = img mask_b64 = img2b64(mask) data['init_images'] = [img_b64] data['alwayson_scripts']['controlnet']['args'][0]['input_image'] = img_b64 data['mask'] = mask_b64 data['negative_prompt'] = nprompt data['prompt'] = prompt response = submit_request(url, json.dumps(data), auth=auth) img_b64 = response.json()['images'][0] return img_b64 def long_side_to(H, W, long_side): asp = H / W if asp > 1: H = long_side H = int(round(H / 32)) * 32 W = int(round(H / asp / 32)) * 32 else: W = long_side W = int(round(W / 32)) * 32 H = int(round(W * asp / 32)) * 32 return H, W if __name__ == '__main__': parser = argparse.ArgumentParser(description='Inpaint instances of people using stable ' 'diffusion.') # workspace\forpaper\eval\kenburns\original\1.png parser.add_argument('--img_path', type=str, help='Path to input image. Can be directory.') # parser.add_argument('--img_path', type=str, default=r'workspace/style/original/', help='Path to input image.') parser.add_argument('--onebyone', type=bool, default=True, help='repainting person one by one') 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/txt2img', help='img2img url') parser.add_argument('--cfg', type=str, default='configs/3d_pixar.yaml', help='repaint config path') parser.add_argument('--bg_nprompt', type=str, default='', 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) parser.add_argument('--style-prompt', default='') parser.add_argument('--global-nprompt', default='') parser.add_argument('--apply-bg-tagger', default=False) parser.add_argument('--apply-fg-tagger', default=True) args = parser.parse_args() args = OmegaConf.create(vars(args)) args.merge_with(OmegaConf.load(args.cfg)) data = { **OmegaConf.to_container(args.sd_params), # "init_images": [img_b64] } auth = None if 'username' in args: username = args.pop('username') password = args.pop('password') auth = HTTPBasicAuth(username, password) img_path = args.img_path if osp.isfile(img_path): imglist = [img_path] else: imglist = find_all_imgs(img_path, abs_path=True) imglist = imglist[::-1] detector = None for ii, img_path in enumerate(imglist): print(f'repainting {img_path} ... {ii+1}/{len(imglist)}') imname = osp.basename(img_path).replace(Path(img_path).suffix, '') cimg = Image.open(img_path).convert('RGB') W, H = cimg.width, cimg.height H, W = long_side_to(H, W, args.long_side) data['width'], data['height'] = W, H img_resized = cimg.resize((W, H), resample=Image.Resampling.LANCZOS) if not osp.exists(args.save_dir): os.makedirs(args.save_dir) if args.onebyone: repaint_args = { '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, 'denoising_strength': args.denoising_strength } data_inpaint = deepcopy(data) data_inpaint.update(repaint_args) from utils.io_utils import json2dict, dict2json, find_all_imgs cache_masks_dir = args.cache_masks_dir mask_fg = None masks = [] if not osp.exists(cache_masks_dir): os.makedirs(cache_masks_dir) promptp = osp.join(cache_masks_dir, f'{imname}_prompts.json') bg_prompt = '' fg_prompts = [] if not osp.exists(promptp): if detector is None: from animeinsseg import AnimeInsSeg import numpy as np from animeinsseg.inpainting import patch_match detector = AnimeInsSeg(args.detector_ckpt, device='cuda') detector.init_tagger() instances = detector.infer(img_path, output_type='numpy', infer_tags=True) if not instances.is_empty: prompts_dict = {} for ii, mask in enumerate(instances.masks): mask = cv2.resize(mask.astype(np.uint8) * 255, (W, H), interpolation=cv2.INTER_AREA) mask = Image.fromarray(mask) savename = imname + '_' + str(ii).zfill(3) + '.png' mask.save(osp.join(cache_masks_dir, savename)) masks.append(mask) tags = instances.tags[ii].split(' ') ctags = instances.character_tags[ii] for ctag in ctags: if ctag in tags: tags.remove(ctag) prompt = ','.join(tags).replace('_', ' ') prompts_dict[savename] = prompt fg_prompts.append(prompt) mask_fg = cv2.resize(instances.compose_masks().astype(np.uint8) * 255, (W, H), interpolation=cv2.INTER_AREA) bg = patch_match.inpaint(np.array(img_resized), mask_fg, patch_size=3) savep = osp.join(cache_masks_dir, f'{imname}_bg_repaint.png') Image.fromarray(bg).save(savep) mask_fg = Image.fromarray(mask_fg) mask_fg.save(osp.join(cache_masks_dir, f'{imname}_mask_fg.png')) bg_tags, character_tags = detector.tagger.label_cv2_bgr(cv2.cvtColor(bg, cv2.COLOR_BGR2RGB)) for ii, t in enumerate(bg_tags): bg_tags[ii] = t.replace('_', ' ') bg_prompt = ','.join(bg_tags) prompts_dict[f'{imname}_bg_repaint.png'] = bg_prompt dict2json(prompts_dict, promptp) else: maskp_list = find_all_imgs(cache_masks_dir, abs_path=False) prompts_dict = json2dict(promptp) for maskn in prompts_dict.keys(): maskp = osp.join(cache_masks_dir, maskn) mask = Image.open(maskp) if maskn.endswith('bg_repaint.png'): bg_prompt = prompts_dict[maskn] bg = mask else: mask = mask.convert('L') fg_prompts.append(prompts_dict[maskn]) masks.append(mask) mask_fg = osp.join(cache_masks_dir, f'{imname}_mask_fg.png') mask_fg = Image.open(mask_fg).convert('L') if len(masks) == 0: print('no fg is found') continue for ii in tqdm(range(args.niter)): if args.random_seed: data['seed'] = randint(0, 65536) else: data['seed'] += ii seed = data['seed'] if args.onebyone: data_inpaint['seed'] = seed if ii == 0: img_b64 = img2b64(bg) img_repainted = img_resized else: img_b64 = output_img_b64 if ii == 0: nprompt = args.bg_nprompt prompt = args.style_prompt + ',' if args.apply_bg_tagger: prompt += bg_prompt + ',' prompt = prompt.strip(',') data['alwayson_scripts']['controlnet']['args'][0]['input_image'] = img_b64 data['init_images'] = [img_b64] data['negative_prompt'] = nprompt data['prompt'] = prompt response = submit_request(args.url, json.dumps(data), auth=auth) output_img_b64 = response.json()['images'][0] bg_repainted = Image.open(io.BytesIO(base64.b64decode(output_img_b64))) # bg_repainted.show() img_repainted = Image.composite(img_repainted, bg_repainted, mask_fg) # img_repainted.show() for jj, (fg_prompt, mask) in enumerate(zip(fg_prompts, masks)): nprompt = args.global_nprompt prompt = args.style_prompt + ',' if args.apply_fg_tagger: prompt += fg_prompt + ',' print(prompt) prompt = prompt.strip(',') output_img_b64 = run_sdinpaint(img_repainted, mask, data_inpaint, prompt=prompt, nprompt=nprompt, url=args.url, auth=auth) img_repainted = Image.open(io.BytesIO(base64.b64decode(output_img_b64))) # img_repainted.show() save_encoded_image(output_img_b64, osp.join(args.save_dir, f'{imname}_onebyone_niter{ii}_output_{seed}.png')) else: img_b64 = img2b64(cimg) data['alwayson_scripts']['controlnet']['args'][0]['input_image'] = img_b64 data['init_images'] = [img_b64] prompt = args.style_prompt + ',' prompt = prompt.strip(',') data['prompt'] = prompt data['negative_prompt'] = args.global_nprompt response = submit_request(args.url, json.dumps(data), auth=auth) output_img_b64 = response.json()['images'][0] imgsavep = osp.join(args.save_dir, f'{imname}_niter{ii}_output_{seed}.png') save_encoded_image(output_img_b64, imgsavep) cimg = Image.open(io.BytesIO(base64.b64decode(output_img_b64)))