246 lines
10 KiB
Python
246 lines
10 KiB
Python
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)
|
|
|
|
|
|
|
|
|