configure model location

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)

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