| |
| """Convert a FastH3 NVFP4 decoder ``decoder.safetensors`` into the ``.pt`` file FastVideo loads. |
| |
| FastVideo's ``FASTVIDEO_H3_VAE_NVFP4_DECODER`` loader reads a |
| ``fastvideo_h3_decoder_nvfp4_deploy_v1`` torch checkpoint. This script rebuilds |
| that checkpoint from the safetensors release without changing any tensor. |
| |
| python convert_to_fastvideo_pt.py quality-36L/decoder.safetensors quality-36L.pt |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
|
|
| import torch |
| from safetensors import safe_open |
| from safetensors.torch import load_file |
|
|
| DEPLOY_FORMAT = "fastvideo_h3_decoder_nvfp4_deploy_v1" |
| GROUPS = ("decoder", "post_quant_conv") |
|
|
|
|
| def safetensors_to_checkpoint(path: str) -> dict: |
| with safe_open(path, framework="pt") as handle: |
| header = handle.metadata() or {} |
| if header.get("format") != DEPLOY_FORMAT: |
| raise ValueError(f"{path} is not a {DEPLOY_FORMAT} file (format: {header.get('format')!r})") |
| tensors = load_file(path) |
| checkpoint: dict = {"format": DEPLOY_FORMAT, "metadata": json.loads(header["nvfp4_metadata"])} |
| for group in GROUPS: |
| prefix = f"{group}." |
| checkpoint[group] = {key[len(prefix):]: value for key, value in tensors.items() if key.startswith(prefix)} |
| unknown = [key for key in tensors if not key.startswith(tuple(f"{g}." for g in GROUPS))] |
| if unknown: |
| raise ValueError(f"unexpected tensors in {path}: {unknown[:5]}") |
| return checkpoint |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument("safetensors", help="quality-36L/decoder.safetensors or fast-8L/decoder.safetensors") |
| parser.add_argument("output", help="destination .pt for FASTVIDEO_H3_VAE_NVFP4_DECODER") |
| args = parser.parse_args() |
| checkpoint = safetensors_to_checkpoint(args.safetensors) |
| torch.save(checkpoint, args.output) |
| print(f"wrote {args.output}: {len(checkpoint['decoder'])} decoder tensors, " |
| f"{checkpoint['metadata'].get('student_layers')} blocks") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
