#!/usr/bin/env python3
"""Convert the MiniMax H3 bf16 DiT to NVFP4 in ComfyUI's native quant layout.

Runs ON the instance (needs /workspace/ComfyUI + comfy_kitchen + a GPU).

Mirrors the quantization policy of the released int8_convrot DiT exactly:
only the 50 main blocks' attn.qkv_proj / attn.out_proj / mlp.fc1 / mlp.fc2
are quantized (200 layers); norms, adaln, patch/condition projections,
token_refiner and final layers stay bf16.

Per quantized layer (identical tensor layout to the released nvfp4_awq TE):
  .weight          U8   [out, in/2]    packed FP4 E2M1 pairs
  .weight_scale    F8_E4M3 [out, in/16] per-16-block scales
  .weight_scale_2  F32  []             global scale (amax / (448*6))
  .comfy_quant     U8   [n]            JSON layer config

Usage:
  python3 convert_nvfp4.py --src models/diffusion_models/minimax_h3_ref2va_bf16.safetensors \
      --dst models/diffusion_models/minimax_h3_ref2va_nvfp4.safetensors [--fpmm]

  --fpmm  adds {"full_precision_matrix_mult": true} (dequant->bf16 GEMM, the
          quality-safe path the official TE uses). Without it, comfy_kitchen's
          native FP4 tensor-core GEMM path is used (faster, more quality risk).
"""
import argparse
import json
import sys
import time

sys.path.insert(0, "/workspace/ComfyUI")

import torch  # noqa: E402
from safetensors import safe_open  # noqa: E402
from safetensors.torch import save_file  # noqa: E402
from comfy.quant_ops import (  # noqa: E402
    TensorCoreConvRotW4A4Layout,
    TensorCoreNVFP4Layout,
)

TARGET_SUFFIXES = (".attn.qkv_proj.weight", ".attn.out_proj.weight",
                   ".mlp.fc1.weight", ".mlp.fc2.weight")


def should_quantize(key, shape):
    return (key.startswith("blocks.")
            and key.endswith(TARGET_SUFFIXES)
            and len(shape) == 2)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True)
    ap.add_argument("--dst", required=True)
    ap.add_argument("--algo", choices=["nvfp4", "convrot_w4a4"], default="nvfp4")
    ap.add_argument("--fpmm", action="store_true",
                    help="nvfp4 only: full_precision_matrix_mult=true "
                         "(dequant->bf16 GEMM)")
    args = ap.parse_args()

    if args.algo == "convrot_w4a4":
        # int4 weights + int4 activations, rotation-assisted (groupsize 256)
        cfg = {"format": "convrot_w4a4", "convrot_groupsize": 256,
               "linear_dtype": "int4"}
    else:
        cfg = {"format": "nvfp4"}
        if args.fpmm:
            cfg["full_precision_matrix_mult"] = True
    cfg_tensor = torch.tensor(list(json.dumps(cfg).encode("utf-8")),
                              dtype=torch.uint8)

    out, n_q, n_keep, t0 = {}, 0, 0, time.time()
    with safe_open(args.src, framework="pt", device="cpu") as f:
        keys = list(f.keys())
        for i, k in enumerate(keys):
            t = f.get_tensor(k)
            if should_quantize(k, t.shape):
                layer = k[:-len(".weight")]
                w = t.cuda()
                if args.algo == "convrot_w4a4":
                    qdata, params = TensorCoreConvRotW4A4Layout.quantize(
                        w, convrot_groupsize=256, linear_dtype="int4")
                    out[layer + ".weight"] = qdata.contiguous().cpu()
                    out[layer + ".weight_scale"] = params.scale.contiguous().cpu()
                else:
                    qdata, params = TensorCoreNVFP4Layout.quantize(w)
                    out[layer + ".weight"] = qdata.contiguous().cpu()
                    out[layer + ".weight_scale"] = params.block_scale.contiguous().cpu()
                    out[layer + ".weight_scale_2"] = params.scale.to(torch.float32).cpu()
                out[layer + ".comfy_quant"] = cfg_tensor.clone()
                del w, qdata, params
                n_q += 1
                if n_q % 20 == 0:
                    torch.cuda.empty_cache()
                    print(f"[{time.time()-t0:6.0f}s] quantized {n_q} layers "
                          f"({i+1}/{len(keys)} tensors)", flush=True)
            else:
                out[k] = t
                n_keep += 1

    print(f"quantized {n_q} layers, kept {n_keep} tensors; saving {args.dst}")
    save_file(out, args.dst)
    size = sum(v.numel() * v.element_size() for v in out.values())
    print(f"done in {time.time()-t0:.0f}s — ~{size/1e9:.1f} GB "
          f"(config: {json.dumps(cfg)})")


if __name__ == "__main__":
    main()