Upload inference.py with huggingface_hub

inference.py

@@ -20,6 +20,10 @@ TRAJ_TEMPLATE_PATH = Path("./assets/template_trajectory.json")
 PATH_START_ID = 9
 PATH_POINT_INTERVAL = 10
 N_ACTION_TOKENS = 6
+WM_TOKENIZER_COMBINATION = {
+    "world_model": "lfq_tokenizer_B_256",
+    "world_model_v2": "lfq_tokenizer_B_256_ema",
+}
 
 # change here if you want to use your own images
 CONDITIONING_FRAMES_DIR = Path("./assets/conditioning_frames")
@@ -109,6 +113,7 @@ if __name__ == "__main__":
     parser.add_argument("--cmd", type=str, default="curving_to_left/curving_to_left_moderate")
     parser.add_argument("--num_frames", type=int, default=25)
     parser.add_argument("--num_overlapping_frames", type=int, default=3)
+    parser.add_argument("--model_name", type=str, default="world_model_v2")
     args = parser.parse_args()
 
     assert args.num_frames <= MAX_NUM_FRAMES, f"`num_frames` should be less than or equal to {MAX_NUM_FRAMES}"
@@ -122,8 +127,9 @@ if __name__ == "__main__":
     output_dir.mkdir(parents=True, exist_ok=True)
 
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    tokenizer = AutoModel.from_pretrained("turing-motors/Terra", subfolder="lfq_tokenizer_B_256", trust_remote_code=True).to(device).eval()
-    model = AutoModel.from_pretrained("turing-motors/Terra", subfolder="world_model", trust_remote_code=True).to(device).eval()
+    tokenizer_name = WM_TOKENIZER_COMBINATION[args.model_name]
+    tokenizer = AutoModel.from_pretrained("turing-motors/Terra", subfolder=tokenizer_name, trust_remote_code=True).to(device).eval()
+    model = AutoModel.from_pretrained("turing-motors/Terra", subfolder=args.model_name, trust_remote_code=True).to(device).eval()
 
     conditioning_frames = load_images(CONDITIONING_FRAMES_PATH_LIST, IMAGE_SIZE).to(device)
     with torch.inference_mode(), torch.autocast(device_type="cuda"):