[feat] Port daVinci-MagiHuman base model with AV output

Ports GAIR-NLP/daVinci-MagiHuman's 15B-param joint-AV DiT into FastVideo:
text + video + audio joint denoise in one flat token stream, 32-step
FlowUniPC w/ CFG=2, Wan 2.2 TI2V-5B video VAE, Stable Audio Open 1.0
audio VAE (first-class port at fastvideo/models/vaes/oobleck.py, no
runtime diffusers import), T5-Gemma 9B UL2 text encoder, auto-muxed
mp4 with h264 video + aac stereo audio.

Parity tests all pass on converted weights: DiT (diff_median=2.6e-3),
text encoder (exact), video VAE (8e-4), audio VAE (exact), pipeline
latent (5e-2).

See .agents/skills/add-model/REVIEW.md for porting procedure + open
review items.

scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py

@@ -0,0 +1,479 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Convert daVinci-MagiHuman (GAIR-NLP) weights to a Diffusers-format repo.
+
+MagiHuman publishes weights in a raw layout on HuggingFace
+(https://huggingface.co/GAIR/daVinci-MagiHuman). The layout is:
+
+    base/                       <- DiT safetensors (sharded)
+    distill/                    <- distilled DiT (out of scope for the base port)
+    540p_sr/, 1080p_sr/         <- super-resolution DiTs (out of scope)
+    turbo_vae/                  <- optional fast VAE decoder (out of scope for first cut)
+
+The base DiT uses Wan-AI/Wan2.2-TI2V-5B's VAE and google/t5gemma-9b-9b-ul2's
+encoder at inference time; neither is bundled upstream.
+
+This converter takes the raw MagiHuman base DiT and emits a Diffusers-style
+directory so `VideoGenerator.from_pretrained(...)` can load it standalone:
+
+    <output>/
+        model_index.json
+        transformer/
+            config.json
+            diffusion_pytorch_model-00001-of-00N.safetensors (+ index)
+        vae/
+            config.json
+            diffusion_pytorch_model.safetensors      (pulled from Wan-AI/Wan2.2-TI2V-5B)
+        scheduler/
+            scheduler_config.json                    (FlowUniPC default)
+        text_encoder/                                (optional; see --bundle-text-encoder)
+            config.json, model-*.safetensors
+        tokenizer/                                   (optional; see --bundle-text-encoder)
+
+The DiT key names pass through unchanged — the FastVideo `MagiHumanDiT` module
+mirrors the reference module tree (`adapter.*`, `block.layers.*`, `final_*`),
+so no regex remapping is needed. The conversion is effectively a reshard +
+Diffusers wrapper.
+
+Example:
+    python scripts/checkpoint_conversion/convert_magi_human_to_diffusers.py \\
+        --source GAIR/daVinci-MagiHuman \\
+        --output converted_weights/magi_human_base \\
+        --bundle-vae \\
+        --bundle-text-encoder
+
+Skip `--bundle-text-encoder` to keep the artifact size down; consumers can
+point at `google/t5gemma-9b-9b-ul2` directly via `text_encoder_path`.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import shutil
+from collections import OrderedDict
+from pathlib import Path
+
+import torch
+from huggingface_hub import hf_hub_download, snapshot_download
+from safetensors.torch import load_file, save_file
+
+
+# MagiHuman base arch — keys must be valid `MagiHumanArchConfig` fields.
+# FastVideo's `TransformerLoader.load` calls `ArchConfig.update_model_arch`
+# with this dict (minus `_class_name`, `_diffusers_version`) and rejects
+# any key that isn't a declared field. Pipeline-level knobs (steps, CFG,
+# guidance scales, flow_shift) live on `MagiHumanT2VConfig` and do NOT
+# belong here — they'd silently shadow the ArchConfig loader otherwise.
+MAGI_HUMAN_BASE_ARCH: dict = {
+    "_class_name": "MagiHumanDiT",
+    "_diffusers_version": "0.33.0",
+    # Transformer shape (upstream `ModelConfig`, `inference/common/config.py`).
+    "num_layers": 40,
+    "hidden_size": 5120,
+    "head_dim": 128,
+    "num_query_groups": 8,
+    # Modality channels.
+    "video_in_channels": 192,            # 48 (VAE z_dim) * patch_size product 1*2*2
+    "audio_in_channels": 64,
+    "text_in_channels": 3584,            # T5Gemma-9B encoder hidden size
+    # Block-level switches.
+    "mm_layers": [0, 1, 2, 3, 36, 37, 38, 39],
+    "local_attn_layers": [],
+    "gelu7_layers": [0, 1, 2, 3],
+    "post_norm_layers": [],
+    "enable_attn_gating": True,
+    "activation_type": "swiglu7",
+    # DiT patching / positional.
+    "patch_size": [1, 2, 2],
+    "spatial_rope_interpolation": "extra",
+    # TReAD (flattened; upstream nests as `tread_config`).
+    "tread_selection_rate": 0.5,
+    "tread_start_layer_idx": 2,
+    "tread_end_layer_idx": 25,
+}
+
+
+SCHEDULER_CONFIG: dict = {
+    "_class_name": "FlowUniPCMultistepScheduler",
+    "_diffusers_version": "0.33.0",
+    "num_train_timesteps": 1000,
+    "solver_order": 2,
+    "prediction_type": "flow_prediction",
+    "shift": 5.0,
+    "predict_x0": True,
+    "solver_type": "bh2",
+    "lower_order_final": True,
+    "disable_corrector": [],
+    "flow_shift": 5.0,
+}
+
+
+MAX_SHARD_BYTES = 5 * 1024 * 1024 * 1024  # 5 GB shards, matches HF defaults
+
+
+def _download_dit_shards(source: Path | str, subfolder: str = "base") -> list[Path]:
+    """Return local paths to all safetensors shards for the DiT."""
+    source = str(source)
+    if os.path.isdir(source):
+        shard_dir = Path(source) / subfolder
+        shards = sorted(shard_dir.glob("*.safetensors"))
+        if not shards:
+            raise FileNotFoundError(f"No safetensors under {shard_dir}")
+        return shards
+
+    # Remote HF repo — pull just the base subfolder.
+    local_dir = snapshot_download(
+        repo_id=source,
+        allow_patterns=[f"{subfolder}/*.safetensors", f"{subfolder}/*.json"],
+    )
+    shard_dir = Path(local_dir) / subfolder
+    return sorted(shard_dir.glob("*.safetensors"))
+
+
+def _load_all_shards(
+    shards: list[Path],
+    cast_bf16: bool = False,
+) -> "OrderedDict[str, torch.Tensor]":
+    """Load all safetensors shards into a single state dict.
+
+    When `cast_bf16` is True, fp32 tensors whose names match the transformer
+    core (attention / mlp / final_linear_* / adapter.{video,text,audio}_embedder)
+    are cast to bfloat16. fp32 is preserved for norms, rope bands, and any
+    other tensor where precision matters. This is the right default for
+    the distill checkpoint, which upstream ships as fp32 master weights
+    (61 GB) — casting yields a 30 GB Diffusers artifact that matches the
+    base checkpoint format.
+    """
+    # Tensors that must stay in float32 regardless of cast_bf16. These are
+    # the dtypes that appear as fp32 in the BASE checkpoint, which is the
+    # ground-truth shape of a "runtime-loadable" MagiHuman repo.
+    _FP32_KEEP_SUFFIXES = (
+        ".pre_norm.weight",
+        ".q_norm.weight",
+        ".k_norm.weight",
+        ".attn_post_norm.weight",
+        ".mlp_post_norm.weight",
+        "final_norm_video.weight",
+        "final_norm_audio.weight",
+        "adapter.rope.bands",
+    )
+    _FP32_KEEP_FULL = {"adapter.rope.bands"}
+
+    def _keep_fp32(k: str) -> bool:
+        if k in _FP32_KEEP_FULL:
+            return True
+        return any(k.endswith(s) for s in _FP32_KEEP_SUFFIXES)
+
+    state: OrderedDict[str, torch.Tensor] = OrderedDict()
+    for shard in shards:
+        piece = load_file(str(shard))
+        for k, v in piece.items():
+            if k in state:
+                raise RuntimeError(f"Duplicate key across shards: {k}")
+            if cast_bf16 and v.dtype == torch.float32 and not _keep_fp32(k):
+                v = v.to(torch.bfloat16)
+            state[k] = v
+        print(f"  loaded {shard.name} ({len(piece)} tensors)")
+    return state
+
+
+def _validate_state(state: dict[str, torch.Tensor]) -> None:
+    """Sanity-check required top-level modules are present."""
+    required_prefixes = (
+        "adapter.video_embedder.",
+        "adapter.text_embedder.",
+        "adapter.audio_embedder.",
+        "adapter.rope.bands",
+        "final_norm_video.",
+        "final_norm_audio.",
+        "final_linear_video.",
+        "final_linear_audio.",
+    )
+    for pref in required_prefixes:
+        if not any(k.startswith(pref) for k in state):
+            raise RuntimeError(f"Missing expected key prefix: {pref}")
+    # Layer count
+    layer_ids = {int(k.split(".")[2]) for k in state if k.startswith("block.layers.")}
+    if layer_ids != set(range(40)):
+        raise RuntimeError(f"Expected layers 0..39, got {sorted(layer_ids)}")
+
+
+def _shard_state_dict(
+    state: dict[str, torch.Tensor],
+    max_bytes: int = MAX_SHARD_BYTES,
+) -> tuple[list[dict[str, torch.Tensor]], dict[str, str]]:
+    """Greedy shard-packing: produce N shards of <= max_bytes, plus index."""
+    shards: list[dict[str, torch.Tensor]] = []
+    index: dict[str, str] = {}
+    cur: dict[str, torch.Tensor] = {}
+    cur_bytes = 0
+    shard_idx = 0
+    total = len(state)
+    for k, v in state.items():
+        t_bytes = v.numel() * v.element_size()
+        if cur and cur_bytes + t_bytes > max_bytes:
+            shards.append(cur)
+            cur = {}
+            cur_bytes = 0
+            shard_idx += 1
+        cur[k] = v
+        cur_bytes += t_bytes
+    if cur:
+        shards.append(cur)
+    n = len(shards)
+    for i, shard in enumerate(shards, start=1):
+        shard_name = f"diffusion_pytorch_model-{i:05d}-of-{n:05d}.safetensors"
+        for k in shard:
+            index[k] = shard_name
+    assert sum(len(s) for s in shards) == total
+    return shards, index
+
+
+def _write_transformer(out_dir: Path, state: dict[str, torch.Tensor], arch: dict) -> None:
+    transformer_dir = out_dir / "transformer"
+    transformer_dir.mkdir(parents=True, exist_ok=True)
+
+    shards, weight_map = _shard_state_dict(state)
+    n = len(shards)
+    total_bytes = sum(v.numel() * v.element_size() for v in state.values())
+    for i, shard in enumerate(shards, start=1):
+        shard_name = f"diffusion_pytorch_model-{i:05d}-of-{n:05d}.safetensors"
+        save_file(shard, str(transformer_dir / shard_name))
+        print(f"  wrote {shard_name} ({len(shard)} tensors)")
+
+    index = {"metadata": {"total_size": total_bytes}, "weight_map": weight_map}
+    with (transformer_dir / "diffusion_pytorch_model.safetensors.index.json").open("w") as f:
+        json.dump(index, f, indent=2)
+        f.write("\n")
+
+    with (transformer_dir / "config.json").open("w") as f:
+        json.dump(arch, f, indent=2)
+        f.write("\n")
+    print(f"  wrote transformer/config.json ({len(arch)} keys)")
+
+
+def _write_scheduler(out_dir: Path) -> None:
+    scheduler_dir = out_dir / "scheduler"
+    scheduler_dir.mkdir(parents=True, exist_ok=True)
+    with (scheduler_dir / "scheduler_config.json").open("w") as f:
+        json.dump(SCHEDULER_CONFIG, f, indent=2)
+        f.write("\n")
+    print(f"  wrote scheduler/scheduler_config.json")
+
+
+def _write_model_index(
+    out_dir: Path,
+    bundle_vae: bool,
+    bundle_text: bool,
+    bundle_audio_vae: bool = False,
+) -> None:
+    index = {
+        "_class_name": "MagiHumanPipeline",
+        "_diffusers_version": "0.33.0",
+        "transformer": ["diffusers", "MagiHumanDiT"],
+        "scheduler": ["diffusers", "FlowUniPCMultistepScheduler"],
+    }
+    if bundle_vae:
+        index["vae"] = ["diffusers", "AutoencoderKLWan"]
+    if bundle_audio_vae:
+        index["audio_vae"] = ["diffusers", "AutoencoderOobleck"]
+    if bundle_text:
+        index["text_encoder"] = ["transformers", "T5GemmaEncoderModel"]
+        index["tokenizer"] = ["transformers", "GemmaTokenizer"]
+    with (out_dir / "model_index.json").open("w") as f:
+        json.dump(index, f, indent=2)
+        f.write("\n")
+    print(f"  wrote model_index.json")
+
+
+def _bundle_wan_vae(out_dir: Path, source_repo: str = "Wan-AI/Wan2.2-TI2V-5B-Diffusers") -> None:
+    """Download the Wan 2.2 TI2V 5B VAE component into <out_dir>/vae/.
+
+    The `-Diffusers` variant has the canonical `vae/config.json` +
+    `vae/diffusion_pytorch_model.safetensors` layout. The plain
+    `Wan-AI/Wan2.2-TI2V-5B` repo ships the VAE as a single `.pth` at the
+    root, which is not `from_pretrained`-friendly.
+    """
+    print(f"  fetching VAE from {source_repo} ...")
+    local = snapshot_download(
+        repo_id=source_repo,
+        allow_patterns=["vae/*"],
+    )
+    src_vae = Path(local) / "vae"
+    if not src_vae.exists():
+        raise FileNotFoundError(f"No vae/ subdir in {source_repo}")
+    dst_vae = out_dir / "vae"
+    if dst_vae.exists():
+        shutil.rmtree(dst_vae)
+    shutil.copytree(src_vae, dst_vae)
+    print(f"  copied {src_vae} -> {dst_vae}")
+
+
+def _bundle_sa_audio_vae(out_dir: Path, source_repo: str = "stabilityai/stable-audio-open-1.0") -> None:
+    """Download the Stable Audio Open 1.0 VAE component into <out_dir>/audio_vae/.
+
+    Stability ships the VAE at `vae/config.json` +
+    `vae/diffusion_pytorch_model.safetensors` inside the main repo, so
+    the bundle is just a copy of that subdir. The repo is gated — the
+    caller's HF token must have accepted terms on
+    https://huggingface.co/stabilityai/stable-audio-open-1.0.
+    """
+    print(f"  fetching audio VAE from {source_repo} (gated) ...")
+    token = (
+        os.environ.get("HF_TOKEN")
+        or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+        or os.environ.get("HF_API_KEY")
+    )
+    local = snapshot_download(
+        repo_id=source_repo, token=token, allow_patterns=["vae/*"],
+    )
+    src = Path(local) / "vae"
+    if not src.exists():
+        raise FileNotFoundError(f"No vae/ subdir in {source_repo}")
+    dst = out_dir / "audio_vae"
+    if dst.exists():
+        shutil.rmtree(dst)
+    shutil.copytree(src, dst)
+    print(f"  copied {src} -> {dst}")
+
+
+def _bundle_text_encoder(out_dir: Path, source_repo: str = "google/t5gemma-9b-9b-ul2") -> None:
+    """Download the T5Gemma encoder + tokenizer.
+
+    T5Gemma is a Google gated repo; this step requires a write-scoped token with
+    accepted terms of use for the repo.
+    """
+    print(f"  fetching text encoder from {source_repo} (gated) ...")
+    token = (
+        os.environ.get("HF_TOKEN")
+        or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+        or os.environ.get("HF_API_KEY")
+    )
+    local = snapshot_download(
+        repo_id=source_repo,
+        token=token,
+        allow_patterns=[
+            "*.json",
+            "*.model",
+            "*.safetensors",
+            "*.safetensors.index.json",
+        ],
+    )
+    # Encoder-only bundling: keep tokenizer at the root and encoder weights
+    # under text_encoder/. HF's T5GemmaEncoderModel.from_pretrained(<dir>) on
+    # the whole repo works, but we split to match Diffusers layout.
+    src = Path(local)
+    dst_encoder = out_dir / "text_encoder"
+    dst_tokenizer = out_dir / "tokenizer"
+    if dst_encoder.exists():
+        shutil.rmtree(dst_encoder)
+    if dst_tokenizer.exists():
+        shutil.rmtree(dst_tokenizer)
+    dst_encoder.mkdir(parents=True, exist_ok=True)
+    dst_tokenizer.mkdir(parents=True, exist_ok=True)
+
+    for fname in src.iterdir():
+        if fname.name in {"tokenizer.model", "tokenizer.json", "tokenizer_config.json",
+                          "special_tokens_map.json", "spiece.model"}:
+            shutil.copy(fname, dst_tokenizer / fname.name)
+        else:
+            shutil.copy(fname, dst_encoder / fname.name)
+    print(f"  staged text_encoder and tokenizer from {src}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    parser.add_argument(
+        "--source",
+        default="GAIR/daVinci-MagiHuman",
+        help="HF repo id or local directory containing base/*.safetensors shards.",
+    )
+    parser.add_argument(
+        "--subfolder",
+        default="base",
+        choices=["base", "distill", "540p_sr", "1080p_sr"],
+        help="Which MagiHuman variant to convert (scope of this skill: base).",
+    )
+    parser.add_argument(
+        "--output",
+        required=True,
+        help="Destination directory for the Diffusers-format repo.",
+    )
+    parser.add_argument(
+        "--bundle-vae",
+        action="store_true",
+        help="Download Wan-AI/Wan2.2-TI2V-5B VAE into <output>/vae/.",
+    )
+    parser.add_argument(
+        "--cast-bf16",
+        action="store_true",
+        help=(
+            "Cast fp32 DiT weights to bfloat16 on save. Recommended for the "
+            "distill subfolder (61 GB fp32 upstream -> 30 GB bf16 artifact). "
+            "Keeps norms, RoPE bands, and other precision-sensitive tensors "
+            "in fp32."
+        ),
+    )
+    parser.add_argument(
+        "--bundle-text-encoder",
+        action="store_true",
+        help="Download google/t5gemma-9b-9b-ul2 into <output>/text_encoder/ and tokenizer/. "
+             "Requires a write-scoped HF token with accepted terms of use.",
+    )
+    parser.add_argument(
+        "--bundle-audio-vae",
+        action="store_true",
+        help="Download stabilityai/stable-audio-open-1.0 VAE into <output>/audio_vae/. "
+             "Requires HF terms accepted for the Stability AI gated repo.",
+    )
+    args = parser.parse_args()
+
+    out_dir = Path(args.output)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"-> DiT shards from {args.source}/{args.subfolder}")
+    shards = _download_dit_shards(args.source, subfolder=args.subfolder)
+    print(f"  found {len(shards)} shard(s)")
+
+    print(f"-> loading DiT state dict (cast_bf16={args.cast_bf16})")
+    state = _load_all_shards(shards, cast_bf16=args.cast_bf16)
+    print(f"  total keys: {len(state)}")
+    _validate_state(state)
+    print(f"  state dict validation passed")
+
+    print(f"-> writing {out_dir}/transformer/")
+    _write_transformer(out_dir, state, MAGI_HUMAN_BASE_ARCH)
+
+    print(f"-> writing {out_dir}/scheduler/")
+    _write_scheduler(out_dir)
+
+    if args.bundle_vae:
+        print(f"-> bundling video VAE (Wan 2.2 TI2V-5B)")
+        _bundle_wan_vae(out_dir)
+
+    if args.bundle_audio_vae:
+        print(f"-> bundling audio VAE (Stable Audio Open 1.0)")
+        _bundle_sa_audio_vae(out_dir)
+
+    if args.bundle_text_encoder:
+        print(f"-> bundling text encoder")
+        _bundle_text_encoder(out_dir)
+
+    print(f"-> writing model_index.json")
+    _write_model_index(
+        out_dir,
+        bundle_vae=args.bundle_vae,
+        bundle_text=args.bundle_text_encoder,
+        bundle_audio_vae=args.bundle_audio_vae,
+    )
+
+    print(f"\nDone. Output at: {out_dir}")
+    if not args.bundle_vae:
+        print("  (remember to fetch Wan-AI/Wan2.2-TI2V-5B VAE separately or re-run with --bundle-vae)")
+    if not args.bundle_text_encoder:
+        print("  (remember to fetch google/t5gemma-9b-9b-ul2 separately or re-run with --bundle-text-encoder)")
+
+
+if __name__ == "__main__":
+    main()

scripts/checkpoint_conversion/push_magi_human_to_hf.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Push a converted daVinci-MagiHuman Diffusers-format directory to the Hub.
+
+This is a thin wrapper around `huggingface_hub.create_repo` + `upload_folder`,
+dedicated to the MagiHuman upload flow. It does NOT modify `create_hf_repo.py`
+(which is LTX-2-oriented and rewrites component weights inside an existing
+Diffusers repo).
+
+Example (one-shot per variant):
+    python scripts/checkpoint_conversion/push_magi_human_to_hf.py \\
+        --local-dir converted_weights/magi_human_base \\
+        --repo-id FastVideo/MagiHuman-Base-Diffusers \\
+        --public
+
+    python scripts/checkpoint_conversion/push_magi_human_to_hf.py \\
+        --local-dir converted_weights/magi_human_distill \\
+        --repo-id FastVideo/MagiHuman-Distilled-Diffusers \\
+        --public
+
+After upload, the local directory can be deleted — the HF repo is the
+source of truth. `VideoGenerator.from_pretrained("FastVideo/...")` pulls
+shards on demand.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+from huggingface_hub import HfApi, create_repo, upload_folder
+
+
+def _validate_local_dir(local_dir: Path) -> None:
+    required = ["model_index.json", "transformer"]
+    missing = [r for r in required if not (local_dir / r).exists()]
+    if missing:
+        sys.exit(
+            f"Error: {local_dir} is missing {missing}. Run "
+            f"convert_magi_human_to_diffusers.py first."
+        )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("--local-dir", required=True, help="Path to the converted Diffusers directory.")
+    parser.add_argument("--repo-id", required=True, help="Target HF repo id, e.g. FastVideo/MagiHuman-Base-Diffusers.")
+    parser.add_argument(
+        "--public",
+        action="store_true",
+        help="Create the repo as public (default: private). Mutually exclusive with --private.",
+    )
+    parser.add_argument(
+        "--private",
+        action="store_true",
+        help="Create the repo as private. Default when neither --public nor --private is set.",
+    )
+    parser.add_argument(
+        "--commit-message",
+        default="Initial upload of daVinci-MagiHuman Diffusers-format conversion.",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Describe what would happen without creating a repo or uploading.",
+    )
+    args = parser.parse_args()
+
+    if args.public and args.private:
+        sys.exit("Error: --public and --private are mutually exclusive.")
+    private = args.private or not args.public
+
+    local_dir = Path(args.local_dir).resolve()
+    _validate_local_dir(local_dir)
+
+    token = (
+        os.environ.get("HF_TOKEN")
+        or os.environ.get("HUGGINGFACE_HUB_TOKEN")
+        or os.environ.get("HF_API_KEY")
+    )
+    if not token:
+        sys.exit(
+            "Error: no HF token in env (set HF_TOKEN / HUGGINGFACE_HUB_TOKEN / HF_API_KEY)."
+        )
+
+    api = HfApi()
+    me = api.whoami(token=token)
+    print(f"token user: {me.get('name')}")
+    print(f"source:     {local_dir}")
+    print(f"target:     {args.repo_id}")
+    print(f"visibility: {'private' if private else 'public'}")
+    if args.dry_run:
+        print("(dry run — not creating or uploading)")
+        return
+
+    print(f"-> create_repo (exist_ok=True)")
+    create_repo(
+        repo_id=args.repo_id,
+        token=token,
+        private=private,
+        exist_ok=True,
+        repo_type="model",
+    )
+
+    print(f"-> upload_folder (this can take a while for 30 GB)")
+    upload_folder(
+        repo_id=args.repo_id,
+        folder_path=str(local_dir),
+        token=token,
+        commit_message=args.commit_message,
+        repo_type="model",
+    )
+    print(f"Done. https://huggingface.co/{args.repo_id}")
+
+
+if __name__ == "__main__":
+    main()