"""Export a slim, inference-only weight bundle for the Hugging Face Space. Builds the models exactly the way evaluate_with_realworld.py does, then serialises only what m2t inference needs: * lm.safetensors flan-t5-base + motion tokens, tied embedding copies dropped * vqvae.safetensors 2D encoder + quantizer (decoder dropped, `encode` only) * adapter.safetensors A_real * stats.npz mean/std for the 68-dim and 81-dim features * flan-t5-base/ config + tokenizer only (no 990MB model.safetensors) Everything stays fp32 — see the size report at the end for why fp16 is not needed. Run from the 2DMotionGPT directory: python space/export_bundle.py \ --vqvae_ckpt ./checkpoints/2d_vqvae_ver3/.../best_vqvae_epoch2960_valacc0.4471.tar \ --adapter_ckpt ./checkpoints/adapter/.../best_adapter_epoch690_valloss1.2251.tar """ import argparse import json import os import shutil import numpy as np import torch from omegaconf import OmegaConf from os.path import join as pjoin from safetensors.torch import save_file from src.config import parse_args from space.adapters import build_adapter from space.mgpt.mgpt_lm import MLM from space.mgpt.mgpt_vq import VQVae # config + tokenizer are all the Space needs; the weights come from lm.safetensors TOKENIZER_FILES = ["config.json", "generation_config.json", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"] def mb(x): return x / 1e6 def state_dict_bytes(sd): return sum(v.numel() * v.element_size() for v in sd.values()) def main(): p = argparse.ArgumentParser(add_help=False) p.add_argument("--vqvae_ckpt", required=True) p.add_argument("--adapter_ckpt", required=True) p.add_argument("--motiongpt_ckpt", default="./checkpoints/MotionGPT-base/motiongpt_s3_h3d.tar") p.add_argument("--adapter_type", default="residual") p.add_argument("--adapter_hidden", type=int, default=512) p.add_argument("--out", default="./space/bundle") a, _ = p.parse_known_args() cfg = parse_args(phase="train") os.makedirs(a.out, exist_ok=True) report = {} ckpt_base = torch.load(a.motiongpt_ckpt, map_location="cpu", weights_only=True) report["source_checkpoint_total"] = mb(state_dict_bytes(ckpt_base["state_dict"])) # ---- VQ-VAE: base weights, then the trained 2D encoder on top --------- vqvae = VQVae( nfeats=cfg.vq.default.params.nfeats, quantizer=cfg.vq.default.params.quantizer, code_num=cfg.vq.default.params.code_num, code_dim=cfg.vq.default.params.code_dim, output_emb_width=cfg.vq.default.params.output_emb_width, down_t=cfg.vq.default.params.down_t, stride_t=cfg.vq.default.params.stride_t, width=cfg.vq.default.params.width, depth=cfg.vq.default.params.depth, dilation_growth_rate=cfg.vq.default.params.dilation_growth_rate, norm=cfg.vq.default.params.norm, activation=cfg.vq.default.params.activation, ) vqvae.load_state_dict({k.replace("vae.", ""): v for k, v in ckpt_base["state_dict"].items() if "vae" in k}) ckpt_2d = torch.load(a.vqvae_ckpt, map_location="cpu", weights_only=False) raw_2d = ckpt_2d.get("model_state_dict", ckpt_2d) missing, unexpected = vqvae.load_state_dict( {k: v for k, v in raw_2d.items() if "encoder" in k}, strict=False) assert not unexpected, f"unexpected 2D encoder keys: {unexpected[:5]}" assert all("encoder" not in k for k in missing), "2D checkpoint did not cover the encoder" vq_full = vqvae.state_dict() vq_slim = {k: v.contiguous() for k, v in vq_full.items() if "decoder" not in k} report["vqvae_full"] = mb(state_dict_bytes(vq_full)) report["vqvae_slim"] = mb(state_dict_bytes(vq_slim)) save_file(vq_slim, pjoin(a.out, "vqvae.safetensors")) # ---- language model -------------------------------------------------- lm = MLM( model_path=cfg.lm.default.params.model_path, model_type=cfg.lm.default.params.model_type, stage=cfg.lm.default.params.stage, motion_codebook_size=cfg.lm.default.params.motion_codebook_size, ) lm.load_state_dict({k.replace("lm.", ""): v for k, v in ckpt_base["state_dict"].items() if "lm" in k}) lm_full = lm.state_dict() tied = ["language_model.encoder.embed_tokens.weight", "language_model.decoder.embed_tokens.weight"] shared = lm_full["language_model.shared.weight"] for k in tied: assert torch.equal(lm_full[k], shared), f"{k} is not tied to shared.weight" lm_slim = {k: v.contiguous() for k, v in lm_full.items() if k not in tied} report["lm_full"] = mb(state_dict_bytes(lm_full)) report["lm_slim"] = mb(state_dict_bytes(lm_slim)) save_file(lm_slim, pjoin(a.out, "lm.safetensors")) # ---- adapter --------------------------------------------------------- adapter = build_adapter(a.adapter_type, dim=81, hidden=a.adapter_hidden) ck_ad = torch.load(a.adapter_ckpt, map_location="cpu", weights_only=False) adapter.load_state_dict(ck_ad.get("model_state_dict", ck_ad)) ad_sd = {k: v.contiguous() for k, v in adapter.state_dict().items()} report["adapter"] = mb(state_dict_bytes(ad_sd)) save_file(ad_sd, pjoin(a.out, "adapter.safetensors")) # ---- normalisation stats -------------------------------------------- meta = pjoin(cfg.DATASET.HUMANML3D.MEAN_STD_PATH, "t2m", "VQVAEV3_CB1024_CMT_H1024_NRES3", "meta") np.savez(pjoin(a.out, "stats.npz"), mean_2d=np.load(pjoin(meta, "mean_2d_coco_normalized.npy")), std_2d=np.load(pjoin(meta, "std_2d_coco_normalized.npy")), mean_est=np.load(pjoin(meta, "mean_2d_coco_estimated_concatenate.npy")), std_est=np.load(pjoin(meta, "std_2d_coco_estimated_concatenate.npy"))) # ---- tokenizer / config (no pretrained weights) ---------------------- tok_out = pjoin(a.out, "flan-t5-base") os.makedirs(tok_out, exist_ok=True) src = cfg.lm.default.params.model_path for fn in TOKENIZER_FILES: sp = pjoin(src, fn) if os.path.exists(sp): shutil.copy(sp, pjoin(tok_out, fn)) else: print(f"[warn] missing {sp}") # ---- manifest -------------------------------------------------------- OmegaConf.save(OmegaConf.create({ "adapter_type": a.adapter_type, "adapter_hidden": a.adapter_hidden, "adapter_dim": 81, "nfeats": cfg.vq.default.params.nfeats, "unit_length": cfg.DATASET.HUMANML3D.UNIT_LEN, "max_motion_length": cfg.DATASET.HUMANML3D.MAX_MOTION_LEN, "motion_codebook_size": cfg.lm.default.params.motion_codebook_size, "vq": dict(cfg.vq.default.params), }), pjoin(a.out, "model_config.yaml")) # ---- report ---------------------------------------------------------- on_disk = {f: os.path.getsize(pjoin(a.out, f)) for f in os.listdir(a.out) if os.path.isfile(pjoin(a.out, f))} print("\n=== bundle ===") for f, b in sorted(on_disk.items(), key=lambda kv: -kv[1]): print(f" {f:24s} {mb(b):8.1f} MB") tok_bytes = sum(os.path.getsize(pjoin(tok_out, f)) for f in os.listdir(tok_out)) print(f" flan-t5-base/{'':11s} {mb(tok_bytes):8.1f} MB") print(f" {'TOTAL':24s} {mb(sum(on_disk.values()) + tok_bytes):8.1f} MB") print("\n=== what was dropped ===") print(f" source checkpoint state_dict {report['source_checkpoint_total']:8.1f} MB") print(f" lm {report['lm_full']:7.1f} -> {report['lm_slim']:7.1f} MB " f"(tied embedding copies)") print(f" vae {report['vqvae_full']:7.1f} -> {report['vqvae_slim']:7.1f} MB (decoder)") print(f" metrics.* evaluator dropped entirely") json.dump(report, open(pjoin(a.out, "export_report.json"), "w"), indent=2) if __name__ == "__main__": main()