import argparse import datetime import glob import json import math import os import sys import time from collections import OrderedDict import cv2 import numpy as np import torch import torchvision ## note: decord should be imported after torch from omegaconf import OmegaConf from pytorch_lightning import seed_everything from tqdm import tqdm sys.path.insert(1, os.path.join(sys.path[0], '..', '..')) from lvdm.models.samplers.ddim import DDIMSampler from main.evaluation.motionctrl_prompts_camerapose_trajs import ( both_prompt_camerapose_traj, cmcm_prompt_camerapose, omom_prompt_traj) from utils.utils import instantiate_from_config DEFAULT_NEGATIVE_PROMPT = 'blur, haze, deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, '\ 'sketch, cartoon, drawing, anime, mutated hands and fingers, deformed, distorted, '\ 'disfigured, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, '\ 'floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, amputation' post_prompt = 'Ultra-detail, masterpiece, best quality, cinematic lighting, 8k uhd, dslr, soft lighting, film grain, Fujifilm XT3' def load_model_checkpoint(model, ckpt, adapter_ckpt=None): if adapter_ckpt: ## main model state_dict = torch.load(ckpt, map_location="cpu") if "state_dict" in list(state_dict.keys()): state_dict = state_dict["state_dict"] result = model.load_state_dict(state_dict, strict=False) else: # deepspeed new_pl_sd = OrderedDict() for key in state_dict['module'].keys(): new_pl_sd[key[16:]]=state_dict['module'][key] result = model.load_state_dict(new_pl_sd, strict=False) print(result) print('>>> model checkpoint loaded.') ## adapter state_dict = torch.load(adapter_ckpt, map_location="cpu") if "state_dict" in list(state_dict.keys()): state_dict = state_dict["state_dict"] model.adapter.load_state_dict(state_dict, strict=True) print('>>> adapter checkpoint loaded.') else: state_dict = torch.load(ckpt, map_location="cpu") if "state_dict" in list(state_dict.keys()): state_dict = state_dict["state_dict"] model.load_state_dict(state_dict, strict=False) else: # deepspeed new_pl_sd = OrderedDict() for key in state_dict['module'].keys(): new_pl_sd[key[16:]]=state_dict['module'][key] model.load_state_dict(new_pl_sd) print('>>> model checkpoint loaded.') return model def load_trajs(cond_dir, trajs): traj_files = [f'{cond_dir}/trajectories/{traj}.npy' for traj in trajs] data_list = [] traj_name = [] for idx in range(len(traj_files)): traj_name.append(traj_files[idx].split('/')[-1].split('.')[0]) data_list.append(torch.tensor(np.load(traj_files[idx])).permute(3, 0, 1, 2).float()) # [t,h,w,c] -> [c,t,h,w] return data_list, traj_name def load_camera_pose(cond_dir, camera_poses): pose_file = [f'{cond_dir}/camera_poses/{pose}.json' for pose in camera_poses] pose_sample_num = len(pose_file) data_list = [] pose_name = [] for idx in range(pose_sample_num): cur_pose_name = camera_poses[idx].replace('test_camera_', '') pose_name.append(cur_pose_name) with open(pose_file[idx], 'r') as f: pose = json.load(f) pose = np.array(pose) # [t, 12] pose = torch.tensor(pose).float() # [t, 12] data_list.append(pose) return data_list, pose_name def save_results(samples, filename, savedir, fps=10): ## save prompt ## save video videos = [samples] savedirs = [savedir] for idx, video in enumerate(videos): if video is None: continue # b,c,t,h,w video = video.detach().cpu() video = torch.clamp(video.float(), -1., 1.) n = video.shape[0] video = video.permute(2, 0, 1, 3, 4) # t,n,c,h,w frame_grids = [torchvision.utils.make_grid(framesheet, nrow=int(n)) for framesheet in video] #[3, 1*h, n*w] grid = torch.stack(frame_grids, dim=0) # stack in temporal dim [t, 3, n*h, w] grid = (grid + 1.0) / 2.0 grid = (grid * 255).to(torch.uint8).permute(0, 2, 3, 1) path = os.path.join(savedirs[idx], "%s.mp4"%filename) torchvision.io.write_video(path, grid, fps=fps, video_codec='h264', options={'crf': '10'}) def motionctrl_sample( model, prompts, noise_shape, camera_poses=None, trajs=None, n_samples=1, unconditional_guidance_scale=1.0, unconditional_guidance_scale_temporal=None, ddim_steps=50, ddim_eta=1., **kwargs): ddim_sampler = DDIMSampler(model) batch_size = noise_shape[0] ## get condition embeddings (support single prompt only) if isinstance(prompts, str): prompts = [prompts] for i in range(len(prompts)): prompts[i] = f'{prompts[i]}, {post_prompt}' cond = model.get_learned_conditioning(prompts) if camera_poses is not None: RT = camera_poses[..., None] else: RT = None if trajs is not None: traj_features = model.get_traj_features(trajs) else: traj_features = None if unconditional_guidance_scale != 1.0: # prompts = batch_size * [""] prompts = batch_size * [DEFAULT_NEGATIVE_PROMPT] uc = model.get_learned_conditioning(prompts) if traj_features is not None: un_motion = model.get_traj_features(torch.zeros_like(trajs)) else: un_motion = None uc = {"features_adapter": un_motion, "uc": uc} else: uc = None batch_variants = [] for _ in range(n_samples): if ddim_sampler is not None: samples, _ = ddim_sampler.sample(S=ddim_steps, conditioning=cond, batch_size=noise_shape[0], shape=noise_shape[1:], verbose=False, unconditional_guidance_scale=unconditional_guidance_scale, unconditional_conditioning=uc, eta=ddim_eta, temporal_length=noise_shape[2], conditional_guidance_scale_temporal=unconditional_guidance_scale_temporal, features_adapter=traj_features, pose_emb=RT, **kwargs ) ## reconstruct from latent to pixel space batch_images = model.decode_first_stage(samples) batch_variants.append(batch_images) ## variants, batch, c, t, h, w batch_variants = torch.stack(batch_variants) return batch_variants.permute(1, 0, 2, 3, 4, 5) def run_inference(args, gpu_num, gpu_no): ## model config config = OmegaConf.load(args.base) model_config = config.pop("model", OmegaConf.create()) model = instantiate_from_config(model_config) model = model.cuda(gpu_no) assert os.path.exists(args.ckpt_path), f"Error: checkpoint {args.ckpt_path} Not Found!" print(f"Loading checkpoint from {args.ckpt_path}") model = load_model_checkpoint(model, args.ckpt_path, args.adapter_ckpt) model.eval() ## run over data assert (args.height % 16 == 0) and (args.width % 16 == 0), "Error: image size [h,w] should be multiples of 16!" ## latent noise shape h, w = args.height // 8, args.width // 8 channels = model.channels frames = model.temporal_length noise_shape = [args.bs, channels, frames, h, w] savedir = os.path.join(args.savedir, "samples") os.makedirs(savedir, exist_ok=True) if args.condtype == 'camera_motion': prompt_list = cmcm_prompt_camerapose['prompts'] camera_pose_list, pose_name = load_camera_pose(args.cond_dir, cmcm_prompt_camerapose['camera_poses']) traj_list = None save_name_list = [] for i in range(len(pose_name)): save_name_list.append(f"{pose_name[i]}__{prompt_list[i].replace(' ', '_').replace(',', '')}") elif args.condtype == 'object_motion': prompt_list = omom_prompt_traj['prompts'] traj_list, traj_name = load_trajs(args.cond_dir, omom_prompt_traj['trajs']) camera_pose_list = None save_name_list = [] for i in range(len(traj_name)): save_name_list.append(f"{traj_name[i]}__{prompt_list[i].replace(' ', '_').replace(',', '')}") elif args.condtype == 'both': prompt_list = both_prompt_camerapose_traj['prompts'] camera_pose_list, pose_name = load_camera_pose(args.cond_dir, both_prompt_camerapose_traj['camera_poses']) traj_list, traj_name = load_trajs(args.cond_dir, both_prompt_camerapose_traj['trajs']) save_name_list = [] for i in range(len(pose_name)): save_name_list.append(f"{pose_name[i]}__{traj_name[i]}__{prompt_list[i].replace(' ', '_').replace(',', '')}") num_samples = len(prompt_list) samples_split = num_samples // gpu_num print('Prompts testing [rank:%d] %d/%d samples loaded.'%(gpu_no, samples_split, num_samples)) #indices = random.choices(list(range(0, num_samples)), k=samples_per_device) indices = list(range(samples_split*gpu_no, samples_split*(gpu_no+1))) prompt_list_rank = [prompt_list[i] for i in indices] camera_pose_list_rank = None if camera_pose_list is None else [camera_pose_list[i] for i in indices] traj_list_rank = None if traj_list is None else [traj_list[i] for i in indices] save_name_list_rank = [save_name_list[i] for i in indices] start = time.time() for idx, indice in tqdm(enumerate(range(0, len(prompt_list_rank), args.bs)), desc='Sample Batch'): prompts = prompt_list_rank[indice:indice+args.bs] camera_poses = None if camera_pose_list_rank is None else camera_pose_list_rank[indice:indice+args.bs] trajs = None if traj_list_rank is None else traj_list_rank[indice:indice+args.bs] save_name = save_name_list_rank[indice:indice+args.bs] print(f'Processing {save_name}') if camera_poses is not None: camera_poses = torch.stack(camera_poses, dim=0).to("cuda") if trajs is not None: trajs = torch.stack(trajs, dim=0).to("cuda") batch_samples = motionctrl_sample( model, prompts, noise_shape, camera_poses=camera_poses, trajs=trajs, n_samples=args.n_samples, unconditional_guidance_scale=args.unconditional_guidance_scale, unconditional_guidance_scale_temporal=args.unconditional_guidance_scale_temporal, ddim_steps=args.ddim_steps, ddim_eta=args.ddim_eta, cond_T = args.cond_T, ) ## save each example individually for nn, samples in enumerate(batch_samples): ## samples : [n_samples,c,t,h,w] prompt = prompts[nn] name = save_name[nn] if len(name) > 90: name = name[:90] filename = f'{name}_{idx*args.bs+nn:04d}_randk{gpu_no}' save_results(samples, filename, savedir, fps=10) if args.save_imgs: parts = save_name[nn].split('__') if len(parts) == 2: cond_name = parts[0] prname = prompts[nn].replace(' ', '_').replace(',', '') cur_outdir = os.path.join(savedir, cond_name, prname) elif len(parts) == 3: poname, trajname, _ = save_name[nn].split('__') prname = prompts[nn].replace(' ', '_').replace(',', '') cur_outdir = os.path.join(savedir, poname, trajname, prname) else: raise NotImplementedError os.makedirs(cur_outdir, exist_ok=True) save_images(samples, cur_outdir) if nn % 100 == 0: print(f'Finish {nn}/{len(batch_samples)}') print(f"Saved in {args.savedir}. Time used: {(time.time() - start):.2f} seconds") def save_images(samples, savedir): ## samples : [n_samples,c,t,h,w] n_samples, c, t, h, w = samples.shape samples = torch.clamp(samples, -1.0, 1.0) samples = (samples + 1.0) / 2.0 samples = (samples * 255).detach().cpu().numpy().astype(np.uint8) for i in range(n_samples): cur_outdir = os.path.join(savedir, f'{i}/images') os.makedirs(cur_outdir, exist_ok=True) for j in range(t): img = samples[i,:,j,:,:] img = np.transpose(img, (1,2,0)) img = img[:,:,::-1] # BGR to RGB path = os.path.join(cur_outdir, f'{j:04d}.png') cv2.imwrite(path, img) def get_parser(): parser = argparse.ArgumentParser() parser.add_argument("--savedir", type=str, default=None, help="results saving path") parser.add_argument("--ckpt_path", type=str, default=None, help="checkpoint path") parser.add_argument("--adapter_ckpt", type=str, default=None, help="adapter checkpoint path") parser.add_argument("--base", type=str, help="config (yaml) path") parser.add_argument("--condtype", default='frame', type=str, help="conditon type: {frame, depth, adapter}") parser.add_argument("--prompt_dir", type=str, default=None, help="a data dir containing videos and prompts") parser.add_argument("--n_samples", type=int, default=1, help="num of samples per prompt",) parser.add_argument("--ddim_steps", type=int, default=50, help="steps of ddim if positive, otherwise use DDPM",) parser.add_argument("--ddim_eta", type=float, default=1.0, help="eta for ddim sampling (0.0 yields deterministic sampling)",) parser.add_argument("--bs", type=int, default=1, help="batch size for inference") parser.add_argument("--height", type=int, default=512, help="image height, in pixel space") parser.add_argument("--width", type=int, default=512, help="image width, in pixel space") parser.add_argument("--unconditional_guidance_scale", type=float, default=1.0, help="prompt classifier-free guidance") parser.add_argument("--unconditional_guidance_scale_temporal", type=float, default=None, help="temporal consistency guidance") parser.add_argument("--seed", type=int, default=20230211, help="seed for seed_everything") parser.add_argument("--cond_T", default=800, type=int, help="Steps smaller than cond_T will not contain condition") parser.add_argument("--save_imgs", action='store_true', help="save condition") parser.add_argument("--cond_dir", type=str, default=None, help="condition dir") return parser if __name__ == '__main__': now = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S") print("@CoLVDM cond-Inference: %s"%now) parser = get_parser() args, unkown = parser.parse_known_args() # args = parser.parse_args() seed_everything(args.seed) rank, gpu_num = 0, 1 run_inference(args, gpu_num, rank)