Add safetensors -> FastVideo loader converter

convert_to_fastvideo_pt.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+"""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()