A2A-Video / scripts /upload_checkpoints_to_hub.py
Muhammad Uzair Khattak
Deploy A2A-Video demo
4bc559f
Raw
History Blame Contribute Delete
4.99 kB
"""
Run this ON THE MACHINE WHERE THE CHECKPOINTS ACTUALLY LIVE (the cluster) --
these paths (/scratch/uzair/..., /datasets/uzair/...) are not reachable from
a laptop or from the HF Space itself.
Uploads to two separate Hub repos, matching MODEL_REPO/TOKENIZERS_REPO in
hf_space_demo/inference.py:
- the main 4M checkpoint -> --model_repo_id (e.g. EPFL-VILAB/Video-4M-models)
- the 7 tokenizer ckpts -> --tokenizers_repo_id (e.g. EPFL-VILAB/Video-4M-tokenizers),
one folder per modality (rgb/, surface-normals/, depth/, opticalflow/, v-jepa-2/, dinov2/, siglip-2/)
They're split because the model checkpoint changes often during training
while the tokenizers are stable, so each gets its own upload/version history.
Usage:
python upload_checkpoints_to_hub.py \\
--model_repo_id EPFL-VILAB/Video-4M-models \\
--tokenizers_repo_id EPFL-VILAB/Video-4M-tokenizers \\
--model_path /scratch/uzair/A18_large_130_checkpoint.pth \\
--vidtok_rgb /scratch/uzair/clariden/vidok_rgb_tokenizer/rgb_normal_last_ckpt_580000_iter.ckpt \\
--vidtok_normal /scratch/uzair/clariden/vidtok_normal_tokenizer/second_run_lr_by_2/420000.ckpt \\
--vidtok_depth /scratch/uzair/clariden/vidtok_depth_tokenizers/depth_first_run_lsat_ckpt_500000.ckpt \\
--vidtok_opticalflow /scratch/uzair/clariden/vidtok_optical_flow_tokenizer_weights/first_run_full_finetuning.ckpt \\
--vjepa /datasets/uzair/weights_from_clariden/vjepa_fixed_tokenizer_weights/.../epoch=3-step=150000.ckpt \\
--dinov2 /scratch/uzair/clariden/tokenizers_for_feature_maps_michael/tokenizers/checkpoints/dinov2_l1_176_411_16807.ckpt \\
--siglipv2 /scratch/uzair/clariden/tokenizers_for_feature_maps_michael/tokenizers/checkpoints/siglip_ens_176_411_16807.ckpt
Pass --model_only or --tokenizers_only to upload just one side (e.g. once
the tokenizers are up, you'll only need --model_only on later runs whenever
the model checkpoint changes).
"""
import argparse
from huggingface_hub import HfApi, create_repo
# Must match MODEL_WEIGHT_FILES / TOKENIZER_WEIGHT_FILES in hf_space_demo/inference.py
MODEL_DEST_FILENAMES = {
"model_path": "main_model/checkpoint.pth",
}
TOKENIZER_DEST_FILENAMES = {
"vidtok_rgb": "rgb/ckpt.ckpt",
"vidtok_normal": "surface-normals/ckpt.ckpt",
"vidtok_depth": "depth/ckpt.ckpt",
"vidtok_opticalflow": "opticalflow/ckpt.ckpt",
"vjepa": "v-jepa-2/ckpt.ckpt",
"dinov2": "dinov2/ckpt.ckpt",
"siglipv2": "siglip-2/ckpt.ckpt",
}
def upload(api, repo_id, private, local_paths_by_key, dest_filenames):
create_repo(repo_id, repo_type="model", private=private, exist_ok=True)
for key, dest_filename in dest_filenames.items():
local_path = local_paths_by_key[key]
print(f"Uploading {local_path} -> {repo_id}:{dest_filename}")
api.upload_file(
path_or_fileobj=local_path,
path_in_repo=dest_filename,
repo_id=repo_id,
repo_type="model",
)
def main():
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--model_repo_id", default="EPFL-VILAB/Video-4M-models")
parser.add_argument("--tokenizers_repo_id", default="EPFL-VILAB/Video-4M-tokenizers")
parser.add_argument("--private", action="store_true", default=True, help="Create repos as private (default)")
parser.add_argument("--public", dest="private", action="store_false")
parser.add_argument("--model_only", action="store_true")
parser.add_argument("--tokenizers_only", action="store_true")
parser.add_argument("--model_path", help="Local path to the main 4M checkpoint")
for key in TOKENIZER_DEST_FILENAMES:
parser.add_argument(f"--{key}", help="Local path to this tokenizer checkpoint on this machine")
args = parser.parse_args()
do_model = not args.tokenizers_only
do_tokenizers = not args.model_only
if do_model and not args.model_path:
parser.error("--model_path is required unless --tokenizers_only is set")
if do_tokenizers and any(getattr(args, key) is None for key in TOKENIZER_DEST_FILENAMES):
missing = [key for key in TOKENIZER_DEST_FILENAMES if getattr(args, key) is None]
parser.error(f"missing tokenizer paths (or pass --model_only): {missing}")
api = HfApi()
if do_model:
upload(api, args.model_repo_id, args.private, {"model_path": args.model_path}, MODEL_DEST_FILENAMES)
if do_tokenizers:
tokenizer_paths = {key: getattr(args, key) for key in TOKENIZER_DEST_FILENAMES}
upload(api, args.tokenizers_repo_id, args.private, tokenizer_paths, TOKENIZER_DEST_FILENAMES)
print("\nDone. Set these env vars where the Space/app runs:")
if do_model:
print(f" FOURM_MODEL_REPO={args.model_repo_id}")
if do_tokenizers:
print(f" FOURM_TOKENIZERS_REPO={args.tokenizers_repo_id}")
if __name__ == "__main__":
main()