from share import * import config import os import cv2 import einops import gradio as gr import numpy as np import torch import random import re from datetime import datetime from glob import glob import argparse from pytorch_lightning import seed_everything from torchvision.transforms import ToPILImage from annotator.util import pad_image, resize_image, HWC3 from annotator.openpose import OpenposeDetector from cldm.model import create_model, load_state_dict from cldm.ddim_hacked import DDIMSampler from pathlib import Path from PIL import Image from omegaconf import OmegaConf from ldm.util import instantiate_from_config, log_txt_as_img from visconet.segm import ATRSegmentCropper as SegmentCropper from huggingface_hub import snapshot_download # supply directory of visual prompt images HF_REPO = 'soonyau/visconet' GALLERY_PATH = Path('./fashion/') WOMEN_GALLERY_PATH = GALLERY_PATH/'WOMEN' MEN_GALLERY_PATH = GALLERY_PATH/'MEN' DEMO = True LOG_SAMPLES = False APP_FILES_PATH = Path('./app_files') VISCON_IMAGE_PATH = APP_FILES_PATH/'default_images' LOG_PATH = APP_FILES_PATH/'logs' SAMPLE_IMAGE_PATH = APP_FILES_PATH/'samples' DEFAULT_CONTROL_SCALE = 1.0 SCALE_CONFIG = { 'Default': [DEFAULT_CONTROL_SCALE]*13, 'DeepFakes':[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.5, 0.5, 0.5, 0.0, 0.0, 0.0, 0.0,], 'Faithful':[1,1,1, 1,1,1, 1,1,0.5, 0.5,0.5,0,0], 'Painting':[0.0,0.0,0.0, 0.5,0.5,0.5, 0.5,0.5,0.5, 0.5,0,0,0], 'Pose': [0.0,0.0,0.0, 0.0,0.0,0.0, 0.0,0.0,0.5, 0.0,0.0,0,0], 'Texture Transfer': [1.0,1.0,1.0, 1.0,1.0,1.0, 0.5,0.0,0.5, 0.0,0.0,0,0] } DEFAULT_SCALE_CONFIG = 'Default' ignore_style_list = ['headwear', 'accesories', 'shoes'] global device global segmentor global apply_openpose global style_encoder global model global ddim_sampler def convert_fname(long_name): gender = 'MEN' if long_name[7:10] == 'MEN' else 'WOMEN' input_list = long_name.replace('fashion','').split('___') # Define a regular expression pattern to match the relevant parts of each input string if gender == 'MEN': pattern = r'MEN(\w+)id(\d+)_(\d)(\w+)' else: pattern = r'WOMEN(\w+)id(\d+)_(\d)(\w+)' # Use a list comprehension to extract the matching substrings from each input string, and format them into the desired output format output_list = [f'{gender}/{category}/id_{id_num[:8]}/{id_num[8:]}_{view_num}_{view_desc}' for (category, id_num, view_num, view_desc) in re.findall(pattern, ' '.join(input_list))] # Print the resulting list of formatted strings return [f +'.jpg' for f in output_list] def fetch_deepfashion(deepfashion_names): src_name, dst_name = convert_fname(deepfashion_names) input_image = np.array(Image.open(image_root/src_name)) pose_image = np.array(Image.open(str(pose_root/dst_name))) mask_image = Image.open(str(mask_root/dst_name).replace('.jpg','_mask.png')) temp = src_name.replace('.jpg','').split('/') lastfolder = temp.pop(-1).replace('_','/', 1) style_folder = style_root/('/'.join(temp+[lastfolder])) viscon_images = [] for style_name in style_names: f_path = style_folder/f'{style_name}.jpg' if os.path.exists(str(f_path)): viscon_images.append(np.array(Image.open(f_path))) else: viscon_images.append(None) return [input_image, pose_image, mask_image, *viscon_images] def select_gallery_image(evt: gr.SelectData): return evt.target.value[evt.index]['name'] def select_default_strength(strength_config): return SCALE_CONFIG[strength_config] def change_all_scales(scale): return [float(scale)]*13 def encode_style_images(style_images): style_embeddings = [] for style_name, style_image in zip(style_names, style_images): if style_image == None: style_image = Image.fromarray(np.zeros((224, 224, 3), dtype=np.uint8)) #style_image = style_image.resize((224,224)) style_image = style_encoder.preprocess(style_image).to(device) style_emb = style_encoder.postprocess(style_encoder(style_image)[0]) style_embeddings.append(style_emb) styles = torch.tensor(np.array(style_embeddings)).squeeze(-2).unsqueeze(0).float().to(device) return styles def save_viscon_images(*viscon_images): ret_images = [] for image, name in zip(viscon_images, style_names): fname = str(VISCON_IMAGE_PATH/name)+'.jpg' if image: image = image.resize((224,224)) if os.path.exists(fname): os.remove(fname) image.save(fname) ret_images.append(image) return ret_images def extract_pose_mask(input_image, detect_resolution, ignore_head=True, ignore_hair=False): # skeleton input_image = pad_image(input_image, min_aspect_ratio=0.625) detected_map, _ = apply_openpose(resize_image(input_image, detect_resolution), hand=True) detected_map = HWC3(detected_map) # human mask cropped = segmentor(input_image, ignore_head=ignore_head, ignore_hair=ignore_hair) mask = cropped['human_mask'] mask = Image.fromarray(np.array(mask*255, dtype=np.uint8), mode='L') return [detected_map, mask] def extract_fashion(input_image): # style images cropped = segmentor(input_image) cropped_images = [] for style_name in style_names: if style_name in cropped and style_name not in ignore_style_list: cropped_images.append(cropped[style_name]) else: cropped_images.append(None) return [*cropped_images] def get_image_files(image_path, ret_image=True, exts=['.jpg','.jpeg','.png']): images = [] for ext in exts: images += [x for x in glob(str(Path(image_path)/f'*{ext}'))] if ret_image: images = [Image.open(x) for x in images] return images def log_sample(seed, results, prompt, skeleton_image, mask_image, control_scales, *viscon_images): time_str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") log_dir = LOG_PATH/time_str os.makedirs(str(log_dir), exist_ok=True) # save result concat = np.hstack((skeleton_image, *results)) Image.fromarray(skeleton_image).save(str(log_dir/'skeleton.jpg')) Image.fromarray(mask_image).save(str(log_dir/'mask.png')) for i, result in enumerate(results): Image.fromarray(result).save(str(log_dir/f'result_{i}.jpg')) # save text with open(str(log_dir/'info.txt'),'w') as f: f.write(f'prompt: {prompt} \n') f.write(f'seed: {seed}\n') control_str = [str(x) for x in control_scales] f.write(','.join(control_str) + '\n') # save vison images for style_name, style_image in zip(style_names, viscon_images): if style_image is not None: style_image.save(str(log_dir/f'{style_name}.jpg')) def process(prompt, a_prompt, n_prompt, num_samples, ddim_steps, scale, seed, eta, mask_image, pose_image, c12, c11, c10, c9, c8, c7, c6, c5, c4, c3, c2, c1, c0, *viscon_images): with torch.no_grad(): control_scales = [c12, c11, c10, c9, c8, c7, c6, c5, c4, c3, c2, c1, c0] mask = torch.tensor(mask_image.mean(-1)/255.,dtype=torch.float) #(512,512), [0,1] mask = mask.unsqueeze(0).to(device) # (1, 512, 512) style_emb = encode_style_images(viscon_images) # fix me detected_map = HWC3(pose_image) #detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_NEAREST) H, W, C = detected_map.shape control = torch.from_numpy(detected_map.copy()).float().to(device) / 255.0 control = torch.stack([control for _ in range(num_samples)], dim=0) control = einops.rearrange(control, 'b h w c -> b c h w').clone() if seed == -1: seed = random.randint(0, 65535) seed_everything(seed) if config.save_memory: model.low_vram_shift(is_diffusing=False) new_style_shape = [num_samples] + [1] * (len(style_emb.shape)-1) cond = {"c_concat": [control], "c_crossattn": [style_emb.repeat(new_style_shape)], "c_text": [model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)], 'c_concat_mask': [mask.repeat(num_samples, 1, 1, 1)]} un_cond = {"c_concat": [control], "c_crossattn": [torch.zeros_like(style_emb).repeat(new_style_shape)], "c_text":[model.get_learned_conditioning([n_prompt] * num_samples)], 'c_concat_mask': [torch.zeros_like(mask).repeat(num_samples, 1, 1, 1)]} shape = (4, H // 8, W // 8) if config.save_memory: model.low_vram_shift(is_diffusing=True) model.control_scales = control_scales samples, _ = ddim_sampler.sample(ddim_steps, num_samples, shape, cond, verbose=False, eta=eta, unconditional_guidance_scale=scale, unconditional_conditioning=un_cond) if config.save_memory: model.low_vram_shift(is_diffusing=False) x_samples = model.decode_first_stage(samples) x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8) results = [x_samples[i] for i in range(num_samples)] if LOG_SAMPLES: log_sample(seed, results, prompt, detected_map, mask_image, control_scales, *viscon_images) return results def get_image(name, file_ext='.jpg'): fname = str(VISCON_IMAGE_PATH/name)+file_ext if not os.path.exists(fname): return None return Image.open(fname) def get_image_numpy(name, file_ext='.png'): fname = str(VISCON_IMAGE_PATH/name)+file_ext if not os.path.exists(fname): return None return np.array(Image.open(fname)) def create_app(): block = gr.Blocks().queue() with block: with gr.Row(): gr.Markdown("## ViscoNet: Visual ControlNet with Human Pose and Fashion
[Video tutorial](https://youtu.be/85NyIuLeV00)") with gr.Row(): with gr.Column(): with gr.Accordion("Get pose and mask", open=False): with gr.Row(): input_image = gr.Image(source='upload', type="numpy", label='input image', value=np.array(get_image_numpy('ref'))) pose_image = gr.Image(source='upload', type="numpy", label='pose', value=np.array(get_image_numpy('pose'))) mask_image = gr.Image(source='upload', type="numpy", label='mask', value=np.array(get_image_numpy('mask'))) with gr.Accordion("Samples", open=False): with gr.Tab('Female'): samples = get_image_files(str(SAMPLE_IMAGE_PATH/'pose/WOMEN/')) female_pose_gallery = gr.Gallery(label='pose', show_label=False, value=samples).style(grid=3, height='auto') with gr.Tab('Male'): samples = get_image_files(str(SAMPLE_IMAGE_PATH/'pose/MEN/')) male_pose_gallery = gr.Gallery(label='pose', show_label=False, value=samples).style(grid=3, height='auto') with gr.Row(): #pad_checkbox = gr.Checkbox(label='Pad pose to square', value=True) ignorehead_checkbox = gr.Checkbox(label='Ignore face in masking (for DeepFake)', value=True) ignorehair_checkbox = gr.Checkbox(label='Ignore hair in masking', value=False, visible=True) with gr.Row(): #ignore_head_checkbox = gr.Checkbox(label='Ignore head', value=False) get_pose_button = gr.Button(label="Get pose", value='Get pose') get_fashion_button = gr.Button(label="Get visual", value='Get visual prompt') with gr.Accordion("Visual Conditions", open=False): gr.Markdown('Drag-and-drop, or click from samples below.') with gr.Column(): viscon_images = [] viscon_images_names2index = {} viscon_len = len(style_names) v_idx = 0 with gr.Row(): for _ in range(8): viscon_name = style_names[v_idx] vis = False if viscon_name in ignore_style_list else True viscon_images.append(gr.Image(source='upload', type="pil", min_height=112, min_width=112, label=viscon_name, value=get_image(viscon_name), visible=vis)) viscon_images_names2index[viscon_name] = v_idx v_idx += 1 viscon_button = gr.Button(value='Save as Default',visible=False if DEMO else True) viscon_galleries = [] with gr.Column(): with gr.Accordion("Female", open=False): for garment, number in zip(['hair', 'top', 'bottom', 'outer'], [150, 500, 500, 250]): with gr.Tab(garment): samples = [] if WOMEN_GALLERY_PATH and os.path.exists(WOMEN_GALLERY_PATH): samples = glob(os.path.join(WOMEN_GALLERY_PATH, f'**/{garment}.jpg'), recursive=True) #samples = glob(f'/home/soon/datasets/deepfashion_inshop/styles_default/WOMEN/**/{garment}.jpg', recursive=True) samples = random.choices(samples, k=number) viscon_gallery = gr.Gallery(label='hair', allow_preview=False, show_label=False, value=samples).style(grid=4, height='auto') viscon_galleries.append({'component':viscon_gallery, 'inputs':[garment]}) with gr.Accordion("Male", open=False): for garment, number in zip(['hair', 'top', 'bottom', 'outer'], [150, 500, 500, 250]): with gr.Tab(garment): samples = [] if MEN_GALLERY_PATH and os.path.exists(MEN_GALLERY_PATH): samples = glob(os.path.join(MEN_GALLERY_PATH, f'**/{garment}.jpg'), recursive=True) samples = random.choices(samples, k=number) viscon_gallery = gr.Gallery(label='hair', allow_preview=False, show_label=False, value=samples).style(grid=4, height='auto') viscon_galleries.append({'component':viscon_gallery, 'inputs':[garment]}) with gr.Accordion("Control Strength Scaling", open=False): gr.Markdown("smaller value for stronger textual influence. c12 is highest spatial resolution controlling textures") with gr.Row(): strength_select = gr.Dropdown(list(SCALE_CONFIG.keys()), label='strength settings', value=DEFAULT_SCALE_CONFIG) scale_all = gr.Slider(label=f'set all scales', minimum=0, maximum=1, value=DEFAULT_CONTROL_SCALE, step=0.05) scale_values = SCALE_CONFIG[DEFAULT_SCALE_CONFIG] control_scales = [] c_idx = 12 with gr.Accordion("Advanced settings", open=False): with gr.Row(): for _ in range(3): control_scales.append(gr.Slider(label=f'c{c_idx}', minimum=0, maximum=1, value=scale_values[12-c_idx], step=0.05)) c_idx -= 1 with gr.Row(): for _ in range(3): control_scales.append(gr.Slider(label=f'c{c_idx}', minimum=0, maximum=1, value=scale_values[12-c_idx], step=0.05)) c_idx -= 1 with gr.Row(): for _ in range(3): control_scales.append(gr.Slider(label=f'c{c_idx}', minimum=0, maximum=1, value=scale_values[12-c_idx], step=0.05)) c_idx -= 1 with gr.Row(): for _ in range(4): control_scales.append(gr.Slider(label=f'c{c_idx}', minimum=0, maximum=1, value=scale_values[12-c_idx], step=0.05)) c_idx -= 1 with gr.Accordion("Advanced options", open=False): with gr.Row(): detect_resolution = gr.Slider(label="OpenPose Resolution", minimum=128, maximum=512, value=512, step=1) ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=20, step=1) scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=12.0, step=0.1) eta = gr.Number(label="eta (DDIM)", value=0.0, visible=False) a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed') n_prompt = gr.Textbox(label="Negative Prompt", value='longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, sunglasses, hat') with gr.Column(): result_gallery = gr.Gallery(label='Output', show_label=False, show_download_button=True, elem_id="gallery").style(grid=1, height='auto') with gr.Row(): max_samples = 8 if not DEMO else 4 num_samples = gr.Slider(label="Images", minimum=1, maximum=max_samples, value=1, step=1) seed = gr.Slider(label="Seed (-1 for random)", minimum=-1, maximum=2147483647, step=1, value=1561194236)#randomize=True) #value=1561194234) if not DEMO: DF_DEMO = 'fashionWOMENTees_Tanksid0000762403_1front___fashionWOMENTees_Tanksid0000762403_1front' DF_EVAL = 'fashionWOMENBlouses_Shirtsid0000035501_1front___fashionWOMENBlouses_Shirtsid0000035501_1front' DF_RESULT ="fashionWOMENTees_Tanksid0000796209_1front___fashionWOMENTees_Tanksid0000796209_2side" deepfashion_names = gr.Textbox(label='Deepfashion name', value=DF_EVAL) gr.Markdown("Default config reconstruct image faithful to pose, mask and visual condition. Reduce control strength to tip balance towards text prompt for more creativity.") prompt = gr.Textbox(label="Text Prompt", value="") run_button = gr.Button(label="Run") female_pose_gallery.select(fn=select_gallery_image, inputs=None, outputs=input_image) male_pose_gallery.select(fn=select_gallery_image, inputs=None, outputs=input_image) for vision_gallery in viscon_galleries: viscon_idx = viscon_images_names2index[vision_gallery['inputs'][0]] vision_gallery['component'].select(fn=select_gallery_image, inputs=None, outputs=viscon_images[viscon_idx]) ips = [prompt, a_prompt, n_prompt, num_samples, ddim_steps, scale, seed, eta, mask_image, pose_image, *control_scales, *viscon_images] run_button.click(fn=process, inputs=ips, outputs=[result_gallery]) prompt.submit(fn=process, inputs=ips, outputs=[result_gallery]) get_pose_button.click(fn=extract_pose_mask, inputs=[input_image, detect_resolution, ignorehead_checkbox, ignorehair_checkbox], outputs=[pose_image, mask_image]) get_fashion_button.click(fn=extract_fashion, inputs=input_image, outputs=[*viscon_images]) viscon_button.click(fn=save_viscon_images, inputs=[*viscon_images], outputs=[*viscon_images]) strength_select.select(fn=select_default_strength, inputs=[strength_select], outputs=[*control_scales]) scale_all.release(fn=change_all_scales, inputs=[scale_all], outputs=[*control_scales]) if not DEMO: deepfashion_names.submit(fn=fetch_deepfashion, inputs=[deepfashion_names], outputs=[input_image, pose_image, mask_image, *viscon_images]) return block if __name__ == "__main__": parser = argparse.ArgumentParser(description='Calculate image-text similarity score.') parser.add_argument('--gpu', type=int, default=0, help='GPU id') parser.add_argument('--config', type=str, default='./configs/visconet_v1.yaml') parser.add_argument('--ckpt', type=str, default='./models/visconet_v1.pth') parser.add_argument('--public_link', action='store_true', default='', help='Create public link') args = parser.parse_args() global device global segmentor global apply_openpose global style_encoder global model global ddim_sampler device = f'cuda:{args.gpu}' if torch.cuda.is_available() else 'cpu' config_file = args.config model_ckpt = args.ckpt proj_config = OmegaConf.load(config_file) style_names = proj_config.dataset.train.params.style_names data_root = Path(proj_config.dataset.train.params.image_root) image_root = data_root/proj_config.dataset.train.params.image_dir style_root = data_root/proj_config.dataset.train.params.style_dir pose_root = data_root/proj_config.dataset.train.params.pose_dir mask_root = data_root/proj_config.dataset.train.params.mask_dir segmentor = SegmentCropper() apply_openpose = OpenposeDetector() snapshot_download(repo_id=HF_REPO, local_dir='./models', allow_patterns=os.path.basename(model_ckpt)) style_encoder = instantiate_from_config(proj_config.model.style_embedding_config).to(device) model = create_model(config_file).cpu() model.load_state_dict(load_state_dict(model_ckpt, location=device)) model = model.to(device) model.cond_stage_model.device = device ddim_sampler = DDIMSampler(model) if not GALLERY_PATH.exists(): zip_name = 'fashion.zip' snapshot_download(repo_id=HF_REPO, allow_patterns=zip_name, local_dir='.') from zipfile import ZipFile with ZipFile(zip_name, 'r') as zip_ref: zip_ref.extractall('.') os.remove(zip_name) # Calling the main function with parsed arguments block = create_app() block.launch(show_api=False, share=True)