acmyu/KeyframesAI
0
1import os2from PIL import Image3import numpy as np4from diffusers import UniPCMultistepScheduler5from src.models.stage2_inpaint_unet_2d_condition import Stage2_InapintUNet2DConditionModel6from src.pipelines.stage2_inpaint_pipeline import Stage2_InpaintDiffusionPipeline7import torch.nn.functional as F8from torchvision import transforms9from diffusers.models.controlnet import ControlNetConditioningEmbedding10from transformers import (11 CLIPVisionModelWithProjection,12 CLIPImageProcessor,13)14import argparse15from transformers import Dinov2Model16from typing import Any, Dict, List, Optional, Tuple, Union17from skimage.metrics import structural_similarity as compare_ssim18 19import torch20import torch.nn as nn21import torch.multiprocessing as mp22import json23import time24 25def split_list_into_chunks(lst, n):26 chunk_size = len(lst) // n27 chunks = [lst[i:i + chunk_size] for i in range(0, len(lst), chunk_size)]28 if len(chunks) > n:29 last_chunk = chunks.pop()30 chunks[-1].extend(last_chunk)31 return chunks32 33 34 35def image_grid(imgs, rows, cols):36 assert len(imgs) == rows * cols37 38 w, h = imgs[0].size39 grid = Image.new("RGB", size=(cols * w, rows * h))40 grid_w, grid_h = grid.size41 42 for i, img in enumerate(imgs):43 grid.paste(img, box=(i % cols * w, i // cols * h))44 return grid45 46 47 48class ImageProjModel_p(torch.nn.Module):49 """SD model with image prompt"""50 51 def __init__(self, in_dim, hidden_dim, out_dim, dropout = 0.):52 super().__init__()53 54 self.net = nn.Sequential(55 nn.Linear(in_dim, hidden_dim),56 nn.GELU(),57 nn.Dropout(dropout),58 nn.LayerNorm(hidden_dim),59 nn.Linear(hidden_dim, out_dim),60 nn.Dropout(dropout)61 )62 63 def forward(self, x):64 return self.net(x)65 66 67 68def inference(args):69 70 device = torch.device("cuda")71 generator = torch.Generator(device=device).manual_seed(args.seed_number)72 73 74 # save path75 save_dir = "{}/show_guidancescale{}_seed{}_numsteps{}/".format(args.save_path, args.guidance_scale, args.seed_number, args.num_inference_steps)76 save_dir_metric = "{}/guidancescale{}_seed{}_numsteps{}/".format(args.save_path, args.guidance_scale, args.seed_number, args.num_inference_steps)77 78 if not os.path.exists(save_dir):79 os.makedirs(save_dir, exist_ok=True)80 81 if not os.path.exists(save_dir_metric):82 os.makedirs(save_dir_metric, exist_ok=True)83 84 clip_image_processor = CLIPImageProcessor()85 86 img_transform = transforms.Compose([87 transforms.ToTensor(),88 transforms.Normalize([0.5], [0.5]),89 ])90 91 92 # model define93 image_proj_model_p_dict = {}94 pose_proj_dict = {}95 unet_dict = {}96 97 image_encoder_g = CLIPVisionModelWithProjection.from_pretrained(args.image_encoder_g_path).to(device).eval()98 image_encoder_p = Dinov2Model.from_pretrained(args.image_encoder_p_path).to(device).eval()99 100 image_proj_model_p = ImageProjModel_p(in_dim=1536, hidden_dim=768, out_dim=1024).to(device).eval()101 pose_proj = ControlNetConditioningEmbedding(320, 3, (16, 32, 96, 256)).to(device).eval()102 103 model_ckpt = args.weights_name104 model_sd = torch.load(model_ckpt, map_location="cpu")["module"]105 106 for k in model_sd.keys():107 if k.startswith("pose_proj"):108 109 pose_proj_dict[k.replace("pose_proj.", "")] = model_sd[k]110 111 elif k.startswith("image_proj_model_p"):112 image_proj_model_p_dict[k.replace("image_proj_model_p.", "")] = model_sd[k]113 114 elif k.startswith("unet"):115 unet_dict[k.replace("unet.", "")] = model_sd[k]116 117 else:118 print(k)119 120 pose_proj.load_state_dict(pose_proj_dict)121 image_proj_model_p.load_state_dict(image_proj_model_p_dict)122 123 pipe = Stage2_InpaintDiffusionPipeline.from_pretrained(args.pretrained_model_name_or_path,torch_dtype=torch.float16).to(device)124 125 pipe.unet= Stage2_InapintUNet2DConditionModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="unet",126 in_channels=9, class_embed_type="projection",127 projection_class_embeddings_input_dim=1024,torch_dtype=torch.float16,128 low_cpu_mem_usage=False, ignore_mismatched_sizes=True).to(device)129 130 pipe.unet.load_state_dict(unet_dict)131 132 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)133 pipe.enable_xformers_memory_efficient_attention()134 #print('====================== json_data: {}, model load finish ==================='.format((args.json_path).split('/')[-1]))135 136 137 data = {138 'source_image': 'sm.png',139 'target_image': 'pose2.png',140 }141 142 s_img_path = (args.img_path + data["source_image"].replace('.jpg', '.png'))143 s_pose_path = args.pose_path + data['source_image'].replace('.jpg', '_pose.jpg')144 145 t_img_path = (args.img_path + data["target_image"].replace('.jpg', '.png'))146 t_pose_path = (args.pose_path + data["target_image"].replace(".jpg", "_pose.jpg"))147 148 149 s_img = Image.open(s_img_path).convert("RGB").resize((args.img_width, args.img_height), Image.BICUBIC)150 t_img = Image.open(t_img_path).convert("RGB").resize((args.img_width, args.img_height), Image.BICUBIC)151 152 black_image = Image.new("RGB", s_img.size, (0, 0, 0))153 s_img_t_mask = Image.new("RGB", (s_img.width * 2, s_img.height))154 s_img_t_mask.paste(s_img, (0, 0))155 s_img_t_mask.paste(black_image, (s_img.width, 0))156 157 s_pose = Image.open(s_pose_path).convert("RGB").resize((args.img_width, args.img_height), Image.BICUBIC)158 t_pose = Image.open(t_pose_path).convert("RGB").resize((args.img_width, args.img_height), Image.BICUBIC)159 st_pose = Image.new("RGB", (s_pose.width * 2, s_pose.height))160 st_pose.paste(s_pose, (0, 0))161 st_pose.paste(t_pose, (s_pose.width, 0))162 163 164 clip_processor_s_img = clip_image_processor(images=s_img, return_tensors="pt").pixel_values165 s_img_f = image_encoder_p(clip_processor_s_img.to(device)).last_hidden_state166 s_img_proj_f = image_proj_model_p(s_img_f) # s_img167 168 169 vae_image = torch.unsqueeze(img_transform(s_img_t_mask), 0)170 171 172 cond_st_pose = torch.unsqueeze(img_transform(st_pose), 0)173 st_pose_f = pose_proj(cond_st_pose.to(device=device)) # t_pose174 175 mode = 'train' # args.json_path.split('/')[-1].split('_')[0]176 if mode == "train":177 clip_processor_s_img = clip_image_processor(images=t_img, return_tensors="pt").pixel_values178 pred_t_img_embed = (image_encoder_g(clip_processor_s_img.to(device)).image_embeds).unsqueeze(1)179 #180 elif mode == "test":181 pred_t_img_embed = torch.tensor(np.load('embed.npy')).to(device)182 pred_t_img_embed = pred_t_img_embed.unsqueeze(1)183 else:184 raise ValueError("Check the input JSON file path")185 186 187 output = pipe(188 height=args.img_height,189 width=args.img_width*2,190 guidance_rescale=0.0,191 vae_image=vae_image,192 s_img_proj_f=s_img_proj_f,193 st_pose_f=st_pose_f,194 pred_t_img_embed = pred_t_img_embed,195 num_images_per_prompt=4,196 guidance_scale=args.guidance_scale,197 generator=generator,198 num_inference_steps=args.num_inference_steps,199 )200 201 202 vis_st_pose = Image.new("RGB", (args.img_width*2, args.img_height))203 vis_st_pose.paste(Image.open(s_pose_path).convert("RGB").resize((args.img_width, args.img_height), Image.BICUBIC), (0, 0))204 vis_st_pose.paste(Image.open(t_pose_path).convert("RGB").resize((args.img_width, args.img_height), Image.BICUBIC), (args.img_width, 0))205 206 vis_st_image = Image.new("RGB", (args.img_width*2, args.img_height))207 vis_st_image.paste(Image.open(s_img_path).convert("RGB").resize((args.img_width, args.img_height), Image.BICUBIC), (0, 0))208 vis_st_image.paste(Image.open(t_img_path).convert("RGB").resize((args.img_width, args.img_height), Image.BICUBIC), (args.img_width, 0))209 210 211 if args.calculate_metrics:212 ssim_values = []213 for gen_img in output.images:214 gen_img = gen_img.crop((args.img_width,0, args.img_width*2,args.img_height))215 ssim_values.append(compare_ssim(np.array(t_img)*255.0, np.array(gen_img)*255.0,216 gaussian_weights=True, sigma=1.2,217 use_sample_covariance=False, multichannel=True, channel_axis=2,218 data_range=(np.array(gen_img)*255.0).max() - (np.array(gen_img)*255.0).min()219 ))220 max_value = max(ssim_values)221 all_ssim.append(max_value)222 max_index = ssim_values.index(max_value)223 grid_metric = output.images[max_index].crop((args.img_width,0, args.img_width*2,args.img_height))224 grid_metric.save(save_dir_metric + s_img_path.split("/")[-1].replace(".png", "") + "_to_" + t_img_path.split("/")[-1])225 else:226 output.images.insert(0, vis_st_pose)227 output.images.insert(0, vis_st_image)228 grid = image_grid(output.images, 2, 3)229 grid.save('coarse.png')230 231 232 if args.calculate_metrics:233 print(sum(all_ssim)/ len(all_ssim))234 235 236if __name__ == "__main__":237 238 parser = argparse.ArgumentParser(description="Simple example of an inpaint model of stage2 script.")239 parser.add_argument("--pretrained_model_name_or_path", type=str,240 default="stabilityai/stable-diffusion-2-1-base",241 help="Path to pretrained model or model identifier from huggingface.co/models.", )242 parser.add_argument("--image_encoder_g_path",type=str,default="laion/CLIP-ViT-H-14-laion2B-s32B-b79K", # openai/clip-vit-base-patch32243 help="Path to pretrained model or model identifier from huggingface.co/models.",)244 parser.add_argument("--image_encoder_p_path",type=str,default="facebook/dinov2-giant",245 help="Path to pretrained model or model identifier from huggingface.co/models.",)246 parser.add_argument("--img_path", type=str,default="imgs/", help="image path", )247 parser.add_argument("--pose_path", type=str,default="imgs/",help="pose path", )248 parser.add_argument("--json_path", type=str,default="./datasets/deepfashing/test_data.json",help="json path", )249 parser.add_argument("--target_embed_path", type=str,default="./logs/view_stage1/512_512/",help="t_img_embed path", )250 parser.add_argument("--save_path", type=str, default="./save_data/stage2", help="save path", ) # ./logs/view_stage2/512_512251 parser.add_argument("--guidance_scale",type=int,default=2.0,help="guidance_scale",)252 parser.add_argument("--seed_number",type=int,default=42,help="seed number",)253 parser.add_argument("--num_inference_steps",type=int,default=20,help="num_inference_steps",)254 parser.add_argument("--img_width",type=int,default=512,help="image width",)255 parser.add_argument("--img_height",type=int,default=512,help="image height",)256 parser.add_argument("--calculate_metrics", action='store_true', help="caculate ssim", )257 parser.add_argument("--weights_name", type=str, default="s2_512.pt",help="weights number", )258 args = parser.parse_args()259 print(args)260 261 inference(args)262 263 """264 num_devices = torch.cuda.device_count()265 print("using {} num_processes inference".format(num_devices))266 267 268 test_data = json.load(open(args.json_path))269 270 select_test_datas = test_data271 print(len(select_test_datas))272 273 mp.set_start_method("spawn")274 data_list = split_list_into_chunks(select_test_datas, num_devices)275 276 277 278 processes = []279 for rank in range(num_devices):280 p = mp.Process(target=inference, args=(args, rank, data_list[rank] ))281 processes.append(p)282 p.start()283 284 for rank, p in enumerate(processes):285 p.join()286 """287 288 289