376e445f0e
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
70 lines
3.0 KiB
Python
70 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Export a pre-packed NVFP4 decoder (``fastvideo_h3_decoder_nvfp4_packed_v1`` safetensors).
|
|
|
|
Loads an exported bf16-master decoder (``export_deploy.py`` .pt or its safetensors) into a dense
|
|
H3 VAE exactly as the pipeline loader does, freezes it on a Blackwell GPU, and writes the frozen
|
|
layers' packed tensors. The loader then skips the bf16 masters entirely.
|
|
|
|
python export_packed.py decoder.safetensors /path/to/MiniMax-H3/vae decoder-nvfp4.safetensors
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import glob
|
|
import json
|
|
import os
|
|
|
|
import torch
|
|
|
|
from fastvideo.models.vaes.minimax_h3_nvfp4_checkpoint import (
|
|
packed_nvfp4_decoder_checkpoint, write_nvfp4_decoder_safetensors)
|
|
from fastvideo.models.vaes.minimax_h3_nvfp4_decoder import (
|
|
apply_nvfp4_decoder_checkpoint, load_nvfp4_decoder_checkpoint)
|
|
|
|
|
|
def _load_h3_vae(vae_dir: str, device: torch.device) -> torch.nn.Module:
|
|
"""The dense H3 VAE the decoder checkpoint was distilled against."""
|
|
from safetensors.torch import load_file
|
|
|
|
from fastvideo.configs.models.vaes.minimax_h3_video import MiniMaxH3VideoVAEConfig
|
|
from fastvideo.models.vaes.minimax_h3_int8_convrot import dense_vae_safetensors
|
|
from fastvideo.models.vaes.minimax_h3_video import AutoencoderKLMiniMaxH3
|
|
|
|
config = MiniMaxH3VideoVAEConfig()
|
|
with open(os.path.join(vae_dir, "config.json")) as handle:
|
|
config.update_model_arch(json.load(handle))
|
|
vae = AutoencoderKLMiniMaxH3(config).to(device)
|
|
state: dict[str, torch.Tensor] = {}
|
|
for path in dense_vae_safetensors(sorted(glob.glob(os.path.join(vae_dir, "*.safetensors")))):
|
|
state.update(load_file(path))
|
|
vae.load_state_dict(state, strict=True)
|
|
return vae.eval().requires_grad_(False)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
parser.add_argument("decoder", help="bf16-master decoder (.pt from export_deploy.py, or .safetensors)")
|
|
parser.add_argument("host_vae_dir", help="MiniMax-H3 vae/ folder with at least the decoder's block count")
|
|
parser.add_argument("output", help="destination .safetensors")
|
|
args = parser.parse_args()
|
|
if not torch.cuda.is_available():
|
|
raise SystemExit("packing needs a CUDA Blackwell GPU (FlashInfer FP4 quantization)")
|
|
|
|
device = torch.device("cuda", torch.cuda.current_device())
|
|
vae = _load_h3_vae(args.host_vae_dir, device)
|
|
metadata = apply_nvfp4_decoder_checkpoint(vae, load_nvfp4_decoder_checkpoint(args.decoder), freeze=True)
|
|
import flashinfer
|
|
|
|
major, minor = torch.cuda.get_device_capability(device)
|
|
packed_on = f"{torch.cuda.get_device_name(device)} sm_{major}{minor}"
|
|
packed = packed_nvfp4_decoder_checkpoint(vae, {
|
|
**metadata, "packed_on": packed_on,
|
|
"flashinfer_version": flashinfer.__version__
|
|
})
|
|
write_nvfp4_decoder_safetensors(args.output, packed)
|
|
print(json.dumps({"output": args.output, "bytes": os.path.getsize(args.output), "metadata": packed["metadata"]}))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|