#!/usr/bin/env python3 r"""Export a self-contained, merged single-file student checkpoint for inference. Why this exists --------------- The Self-Forcing student is *saved as a trainable-only delta* (LoRA + selectively-unfrozen modules — ~5 GB at 14B; see ``checkpointer.py``), so running it normally requires assembling THREE weight sources at load time: base Wan 2.1 T2V + OmniAvatar-LS V2V adapter + this SF delta This script performs that assembly once, merges the LoRA into the base, and writes ONE fully-baked ``.pth`` (plain keys, no PEFT wrappers) that the inference scripts load with ``--ckpt_path`` ALONE — no ``--base_model_paths`` or ``--omniavatar_ckpt_path`` needed. This is the recommended artifact to publish for the inference-only audience. Run on a machine that has the three sources + enough RAM/VRAM (e.g. one H200). It re-uses the exact assembly path from ``scripts/inference/_loader.py`` so the exported weights are numerically identical to a normal 3-source load. Example ------- python scripts/export_merged_checkpoint.py \ --model_size 14B \ --ckpt_path /path/to/0000600.pth \ # SF delta (.pth or .net_model dir) --base_model_paths /path/wan2.1_t2v_14b-00001.safetensors,/path/...-00002.safetensors \ --omniavatar_ckpt_path /path/to/omniavatar_ls_14b.pt \ --output_path /path/to/lipforcing_14b.pth Then inference needs only: python scripts/inference/inference_streaming.py --model_size 14B \ --ckpt_path /path/to/lipforcing_14b.pth \ --vae_path ... --wav2vec_path ... --mask_path ... --taehv_ckpt ... """ import argparse import os import sys import torch # --- Path setup (mirror the inference scripts) ---------------------------- SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) LIPFORCING_ROOT = os.path.abspath(os.path.join(SCRIPT_DIR, "..")) sys.path.insert(0, LIPFORCING_ROOT) sys.path.insert(0, os.path.join(SCRIPT_DIR, "inference")) # for `import _loader` os.environ.setdefault("OMNIAVATAR_ROOT", LIPFORCING_ROOT) from _loader import load_diffusion_model # noqa: E402 _DTYPES = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} def _to_plain_merged_state_dict(core: torch.nn.Module) -> dict: """Map a post-merge PEFT ``_core`` state_dict to plain (unwrapped) keys. After ``LoraLayer.merge()`` the merged weights live in ``.base_layer.weight`` while the now-redundant ``lora_A``/``lora_B`` tensors remain. A model constructed WITHOUT the OmniAvatar adapter has no PEFT wrappers, so its keys are plain (``.weight``). We therefore drop the lora_* tensors and strip the ``base_layer.`` infix. (For the 1.3B path the OmniAvatar LoRA is already folded at construction, so ``_core`` has plain keys to begin with and this is a near no-op.) """ raw = core.state_dict() out, n_drop, n_unwrap, n_plain = {}, 0, 0, 0 for k, v in raw.items(): if ".lora_A." in k or ".lora_B." in k or k.endswith(".lora_A") or k.endswith(".lora_B"): n_drop += 1 continue if ".base_layer." in k: k = k.replace(".base_layer.", ".") n_unwrap += 1 else: n_plain += 1 out[k] = v print(f" merged state_dict: {len(out)} tensors " f"({n_unwrap} unwrapped from PEFT base_layer, {n_plain} plain, {n_drop} lora_* dropped)") return out def _verify_reload(plain_sd: dict, args) -> bool: """Build a fresh model WITHOUT base/omni and confirm it accepts the merged state_dict with no missing parameters and no unexpected keys. This is the exact structure the inference scripts construct when only ``--ckpt_path`` is given, so a clean result here means single-flag inference will load fully. """ from lipforcing.networks.OmniAvatar.network_causal import CausalOmniAvatarWan probe = CausalOmniAvatarWan( model_size=args.model_size, in_dim=65, mode="v2v", use_audio=True, audio_hidden_size=32, chunk_size=args.chunk_size, total_num_frames=21, base_model_paths=None, omniavatar_ckpt_path=None, merge_lora=False, lora_rank=128, lora_alpha=64, net_pred_type="flow", schedule_type="rf", mask_all_frames=True, dtype=args.dtype, local_attn_size=args.local_attn_size, sink_size=args.sink_size, use_dynamic_rope=args.use_dynamic_rope, ) prefixed = {"_core." + k: v for k, v in plain_sd.items()} missing, unexpected = probe.load_state_dict(prefixed, strict=False) param_keys = set(dict(probe.named_parameters()).keys()) missing_params = [m for m in missing if m in param_keys] print(f" verify reload: {len(missing_params)} missing params, {len(unexpected)} unexpected keys") if missing_params: print(f" WARN missing params (first 10): {missing_params[:10]}") if unexpected: print(f" WARN unexpected keys (first 10): {unexpected[:10]}") return not missing_params and not unexpected def main(): ap = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--model_size", default="14B", choices=["14B", "1.3B"]) ap.add_argument("--ckpt_path", required=True, help="Self-Forcing student delta: .pth, .net_model distcp dir, or .pth + sibling dir.") ap.add_argument("--base_model_paths", required=True, help="Comma-separated base Wan 2.1 T2V safetensor paths (1.3B or 14B).") ap.add_argument("--omniavatar_ckpt_path", required=True, help="OmniAvatar-LS V2V adapter checkpoint (.pt/.pth) — LoRA + audio + patch-embed.") ap.add_argument("--output_path", required=True, help="Output single-file merged .pth.") ap.add_argument("--dtype", default="bf16", choices=list(_DTYPES), help="Compute dtype for assembly (default bf16, matches training/inference).") ap.add_argument("--save_dtype", default=None, choices=list(_DTYPES), help="Dtype to store floating tensors in (default: same as --dtype).") ap.add_argument("--chunk_size", type=int, default=3) ap.add_argument("--local_attn_size", type=int, default=-1) ap.add_argument("--sink_size", type=int, default=0) ap.add_argument("--use_dynamic_rope", action="store_true", default=False) ap.add_argument("--device", default="cuda") ap.add_argument("--no_verify", action="store_true", help="Skip the reload sanity check.") args = ap.parse_args() args.merge_lora_post_load = True # force the LoRA->base merge for 14B dtype = _DTYPES[args.dtype] device = torch.device(args.device) print(f"Assembling {args.model_size} student (base Wan + OmniAvatar-LS + SF delta) ...") model = load_diffusion_model(args, device, dtype) # constructs + loads 3 sources + merges LoRA plain = _to_plain_merged_state_dict(model._core) save_dtype = _DTYPES[args.save_dtype] if args.save_dtype else dtype plain = {k: (v.to(save_dtype) if v.is_floating_point() else v).cpu() for k, v in plain.items()} if not args.no_verify: ok = _verify_reload(plain, args) print(f" verify: {'OK — single-flag inference will load fully.' if ok else 'MISMATCH — inspect the warnings above before publishing.'}") out_dir = os.path.dirname(os.path.abspath(args.output_path)) if out_dir: os.makedirs(out_dir, exist_ok=True) print(f"Writing merged checkpoint -> {args.output_path} ...") torch.save(plain, args.output_path) sz = os.path.getsize(args.output_path) / 1e9 print(f"Done. {len(plain)} tensors, {sz:.1f} GB. Load it with --ckpt_path alone (no base/omni flags).") if __name__ == "__main__": main()