""" Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved. This source code is licensed under the license found in the LICENSE file in the root directory of this source tree. """ import os from typing import Callable, Dict, Union import numpy as np import torch from data_loaders.get_data import get_dataset_loader, load_local_data from diffusion.respace import SpacedDiffusion from model.cfg_sampler import ClassifierFreeSampleModel from model.diffusion import FiLMTransformer from torch.utils.data import DataLoader from utils.diff_parser_utils import generate_args from utils.misc import fixseed, prGreen from utils.model_util import create_model_and_diffusion, get_person_num, load_model def _construct_template_variables(unconstrained: bool) -> (str,): row_file_template = "sample{:02d}.mp4" all_file_template = "samples_{:02d}_to_{:02d}.mp4" if unconstrained: sample_file_template = "row{:02d}_col{:02d}.mp4" sample_print_template = "[{} row #{:02d} column #{:02d} | -> {}]" row_file_template = row_file_template.replace("sample", "row") row_print_template = "[{} row #{:02d} | all columns | -> {}]" all_file_template = all_file_template.replace("samples", "rows") all_print_template = "[rows {:02d} to {:02d} | -> {}]" else: sample_file_template = "sample{:02d}_rep{:02d}.mp4" sample_print_template = '["{}" ({:02d}) | Rep #{:02d} | -> {}]' row_print_template = '[ "{}" ({:02d}) | all repetitions | -> {}]' all_print_template = "[samples {:02d} to {:02d} | all repetitions | -> {}]" return ( sample_print_template, row_print_template, all_print_template, sample_file_template, row_file_template, all_file_template, ) def _replace_keyframes( model_kwargs: Dict[str, Dict[str, torch.Tensor]], model: Union[FiLMTransformer, ClassifierFreeSampleModel], ) -> torch.Tensor: B, T = ( model_kwargs["y"]["keyframes"].shape[0], model_kwargs["y"]["keyframes"].shape[1], ) with torch.no_grad(): tokens = model.transformer.generate( model_kwargs["y"]["audio"], T, layers=model.tokenizer.residual_depth, n_sequences=B, ) tokens = tokens.reshape((B, -1, model.tokenizer.residual_depth)) pred = model.tokenizer.decode(tokens).detach().cpu() assert ( model_kwargs["y"]["keyframes"].shape == pred.shape ), f"{model_kwargs['y']['keyframes'].shape} vs {pred.shape}" return pred def _run_single_diffusion( args, model_kwargs: Dict[str, Dict[str, torch.Tensor]], diffusion: SpacedDiffusion, model: Union[FiLMTransformer, ClassifierFreeSampleModel], inv_transform: Callable, gt: torch.Tensor, ) -> (torch.Tensor,): if args.data_format == "pose" and args.resume_trans is not None: model_kwargs["y"]["keyframes"] = _replace_keyframes(model_kwargs, model) sample_fn = diffusion.ddim_sample_loop with torch.no_grad(): sample = sample_fn( model, (args.batch_size, model.nfeats, 1, args.curr_seq_length), clip_denoised=False, model_kwargs=model_kwargs, init_image=None, progress=True, dump_steps=None, noise=None, const_noise=False, ) sample = inv_transform(sample.cpu().permute(0, 2, 3, 1), args.data_format).permute( 0, 3, 1, 2 ) curr_audio = inv_transform(model_kwargs["y"]["audio"].cpu().numpy(), "audio") keyframes = inv_transform(model_kwargs["y"]["keyframes"], args.data_format) gt_seq = inv_transform(gt.cpu().permute(0, 2, 3, 1), args.data_format).permute( 0, 3, 1, 2 ) return sample, curr_audio, keyframes, gt_seq def _generate_sequences( args, model_kwargs: Dict[str, Dict[str, torch.Tensor]], diffusion: SpacedDiffusion, model: Union[FiLMTransformer, ClassifierFreeSampleModel], test_data: torch.Tensor, gt: torch.Tensor, ) -> Dict[str, np.ndarray]: all_motions = [] all_lengths = [] all_audio = [] all_gt = [] all_keyframes = [] for rep_i in range(args.num_repetitions): print(f"### Sampling [repetitions #{rep_i}]") # add CFG scale to batch if args.guidance_param != 1: model_kwargs["y"]["scale"] = ( torch.ones(args.batch_size, device=args.device) * args.guidance_param ) model_kwargs["y"] = { key: val.to(args.device) if torch.is_tensor(val) else val for key, val in model_kwargs["y"].items() } sample, curr_audio, keyframes, gt_seq = _run_single_diffusion( args, model_kwargs, diffusion, model, test_data.dataset.inv_transform, gt ) all_motions.append(sample.cpu().numpy()) all_audio.append(curr_audio) all_keyframes.append(keyframes.cpu().numpy()) all_gt.append(gt_seq.cpu().numpy()) all_lengths.append(model_kwargs["y"]["lengths"].cpu().numpy()) print(f"created {len(all_motions) * args.batch_size} samples") return { "motions": np.concatenate(all_motions, axis=0), "audio": np.concatenate(all_audio, axis=0), "gt": np.concatenate(all_gt, axis=0), "lengths": np.concatenate(all_lengths, axis=0), "keyframes": np.concatenate(all_keyframes, axis=0), } def _render_pred( args, data_block: Dict[str, torch.Tensor], sample_file_template: str, audio_per_frame: int, ) -> None: from visualize.render_codes import BodyRenderer face_codes = None if args.face_codes is not None: face_codes = np.load(args.face_codes, allow_pickle=True).item() face_motions = face_codes["motions"] face_gts = face_codes["gt"] face_audio = face_codes["audio"] config_base = f"./checkpoints/ca_body/data/{get_person_num(args.data_root)}" body_renderer = BodyRenderer( config_base=config_base, render_rgb=True, ) for sample_i in range(args.num_samples): for rep_i in range(args.num_repetitions): idx = rep_i * args.batch_size + sample_i save_file = sample_file_template.format(sample_i, rep_i) animation_save_path = os.path.join(args.output_dir, save_file) # format data length = data_block["lengths"][idx] body_motion = ( data_block["motions"][idx].transpose(2, 0, 1)[:length].squeeze(-1) ) face_motion = face_motions[idx].transpose(2, 0, 1)[:length].squeeze(-1) assert np.array_equal( data_block["audio"][idx], face_audio[idx] ), "face audio is not the same" audio = data_block["audio"][idx, : length * audio_per_frame, :].T # set up render data block to pass into renderer render_data_block = { "audio": audio, "body_motion": body_motion, "face_motion": face_motion, } if args.render_gt: gt_body = data_block["gt"][idx].transpose(2, 0, 1)[:length].squeeze(-1) gt_face = face_gts[idx].transpose(2, 0, 1)[:length].squeeze(-1) render_data_block["gt_body"] = gt_body render_data_block["gt_face"] = gt_face body_renderer.render_full_video( render_data_block, animation_save_path, audio_sr=audio_per_frame * 30, render_gt=args.render_gt, ) def _reset_sample_args(args) -> None: # set the sequence length to match the one specified by user name = os.path.basename(os.path.dirname(args.model_path)) niter = os.path.basename(args.model_path).replace("model", "").replace(".pt", "") args.curr_seq_length = ( args.curr_seq_length if args.curr_seq_length is not None else args.max_seq_length ) # add the resume predictor model path resume_trans_name = "" if args.data_format == "pose" and args.resume_trans is not None: resume_trans_parts = args.resume_trans.split("/") resume_trans_name = f"{resume_trans_parts[1]}_{resume_trans_parts[-1]}" # reformat the output directory args.output_dir = os.path.join( os.path.dirname(args.model_path), "samples_{}_{}_seed{}_{}".format(name, niter, args.seed, resume_trans_name), ) assert ( args.num_samples <= args.batch_size ), f"Please either increase batch_size({args.batch_size}) or reduce num_samples({args.num_samples})" # set the batch size to match the number of samples to generate args.batch_size = args.num_samples def _setup_dataset(args) -> DataLoader: data_root = args.data_root data_dict = load_local_data( data_root, audio_per_frame=1600, flip_person=args.flip_person, ) test_data = get_dataset_loader( args=args, data_dict=data_dict, split="test", chunk=True, ) return test_data def _setup_model( args, ) -> (Union[FiLMTransformer, ClassifierFreeSampleModel], SpacedDiffusion): model, diffusion = create_model_and_diffusion(args, split_type="test") print(f"Loading checkpoints from [{args.model_path}]...") state_dict = torch.load(args.model_path, map_location="cpu") load_model(model, state_dict) if not args.unconstrained: assert args.guidance_param != 1 if args.guidance_param != 1: prGreen("[CFS] wrapping model in classifier free sample") model = ClassifierFreeSampleModel(model) model.to(args.device) model.eval() return model, diffusion def main(): args = generate_args() fixseed(args.seed) _reset_sample_args(args) print("Loading dataset...") test_data = _setup_dataset(args) iterator = iter(test_data) print("Creating model and diffusion...") model, diffusion = _setup_model(args) if args.pose_codes is None: # generate sequences gt, model_kwargs = next(iterator) data_block = _generate_sequences( args, model_kwargs, diffusion, model, test_data, gt ) os.makedirs(args.output_dir, exist_ok=True) npy_path = os.path.join(args.output_dir, "results.npy") print(f"saving results file to [{npy_path}]") np.save(npy_path, data_block) else: # load the pre generated results data_block = np.load(args.pose_codes, allow_pickle=True).item() # plot function only if face_codes exist and we are on pose prediction if args.plot: assert args.face_codes is not None, "need body and faces" assert ( args.data_format == "pose" ), "currently only supporting plot on pose stuff" print(f"saving visualizations to [{args.output_dir}]...") _, _, _, sample_file_template, _, _ = _construct_template_variables( args.unconstrained ) _render_pred( args, data_block, sample_file_template, test_data.dataset.audio_per_frame, ) if __name__ == "__main__": main()