diff --git a/animation/animate-x/inference_cli.py b/animation/animate-x/inference_cli.py index c8315ac..8311ada 100755 --- a/animation/animate-x/inference_cli.py +++ b/animation/animate-x/inference_cli.py @@ -29,6 +29,12 @@ def parse_args(): default=None, required=True ) + parser.add_argument( + "--model_checkpoints_directory", + type=str, + default=None, + required=False + ) parser.add_argument( "--pickle_data_file", type=str, @@ -59,8 +65,16 @@ def parse_args(): default=14, required=False ) + parser.add_argument( + "--round", + type=int, + default=1, + required=False + ) - # TODO(bt): Means to load model files + # TODO(bt): Means to control output filename + + # TODO(bt): Don't emit side-by-side comparison video #parser.add_argument( # "--pretrained_model_name_or_path", @@ -91,7 +105,7 @@ def main(): raise Exception(f"Pickle data file does not exist: #{args.pickle_data_file}") cfg_update.cfg_dict['test_list_path'] = [[ - 4, # ??? + 4, # (???) Not sure what this parameter is! args.image_file, args.pose_directory, args.frame_directory, @@ -99,12 +113,22 @@ def main(): args.seed, ]] + # NB(bt): I think this controls extra inference rounds. + cfg_update.cfg_dict['round'] = args.round + if args.max_frames: cfg_update.cfg_dict['max_frames'] = args.max_frames if args.width and args.height: cfg_update.cfg_dict['resolution'] = [args.width, args.height] + model_dir = 'checkpoints/' if not args.model_checkpoints_directory else args.model_checkpoints_directory + + cfg_update.cfg_dict['test_model'] = os.path.join(model_dir, 'animate-x_ckpt.pth') + cfg_update.cfg_dict['embedder']['pretrained'] = os.path.join(model_dir, 'open_clip_pytorch_model.bin') + cfg_update.cfg_dict['auto_encoder']['pretrained'] = os.path.join(model_dir, 'v2-1_512-ema-pruned.ckpt') + + print("Configurations:\n\n", cfg_update.cfg_dict, "\n\n") INFER_ENGINE.build(dict(type=cfg_update.TASK_TYPE), cfg_update=cfg_update.cfg_dict) diff --git a/animation/animate-x/run_inference.sh b/animation/animate-x/run_inference.sh index 1c623bb..20cd4d4 100755 --- a/animation/animate-x/run_inference.sh +++ b/animation/animate-x/run_inference.sh @@ -4,12 +4,13 @@ python inference_cli.py \ --cfg configs/Animate_X_infer.yaml \ - --image_file data_test_copy/images/future_girl_2.png \ - --pose_directory data_test_copy/saved_pose/dance_1 \ - --frame_directory data_test_copy/saved_frames/dance_1 \ - --pickle_data_file data_test_copy/saved_pkl/dance_1.pkl \ - --width 768 \ - --height 512 + --image_file data_test_copy/images/jeff_goldblum.jpg \ + --pose_directory data_test_copy/saved_pose/dance_2 \ + --frame_directory data_test_copy/saved_frames/dance_2 \ + --pickle_data_file data_test_copy/saved_pkl/dance_2.pkl \ + --model_checkpoints_directory checkpoints2 \ + --height 768 \ + --width 512 # --height 512 # --height 768