# 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()