"""Requantize MiniMax-H3 pruned_int8_convrot -> NVFP4.

The StarNodes converter cannot do this: it passes non-floating-point tensors through
untouched, so an already-int8 model comes out byte-identical. This walks the state dict
and does a real dequantize -> requantize on the 200 attn/mlp layers, leaving AdaLN, the
norms, embedders and token_refiner exactly as they are.

Per-layer input:   .weight (I8), .weight_scale (F32), .comfy_quant (U8 json bytes)
Per-layer output:  .weight + NVFP4 layout suffixes, .comfy_quant = {"format": "nvfp4"}
"""
import json, os, sys, time, collections
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from comfy_kitchen import tensor as ckt

SRC = "/ComfyUI/models/diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors"
DST = "/ComfyUI/models/diffusion_models/minimax_h3_ref2va_pruned_nvfp4.safetensors"
DEV = "cuda"
FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2)

f = safe_open(SRC, "pt")
keys = list(f.keys())

# layers that carry a comfy_quant blob are the quantized ones
quant_bases = sorted({k[: -len(".comfy_quant")] for k in keys if k.endswith(".comfy_quant")})
print(f"quantized layers found: {len(quant_bases)}", flush=True)

out, stats = {}, collections.Counter()
consumed = set()
t0 = time.time()

for i, base in enumerate(quant_bases):
    wk, sk, ck_ = f"{base}.weight", f"{base}.weight_scale", f"{base}.comfy_quant"
    cfg = json.loads(bytes(f.get_tensor(ck_).tolist()).decode())
    if cfg.get("format") != "int8_tensorwise":
        print(f"  SKIP {base}: unexpected format {cfg}", flush=True)
        stats["skipped_unknown_format"] += 1
        continue

    qdata = f.get_tensor(wk).to(DEV)
    scale = f.get_tensor(sk).to(DEV)

    params = ckt.TensorWiseINT8Layout.Params(
        scale=scale,
        orig_dtype=torch.bfloat16,
        orig_shape=tuple(qdata.shape),
        is_weight=True,
        convrot=cfg.get("convrot", False),
        convrot_groupsize=cfg.get("convrot_groupsize", 256),
    )
    deq = ckt.TensorWiseINT8Layout.dequantize(qdata, params)          # -> bf16

    nq, nparams = ckt.TensorCoreNVFP4Layout.quantize(deq.float().contiguous())
    tensors = ckt.TensorCoreNVFP4Layout.state_dict_tensors(nq, nparams)

    for suffix, t in tensors.items():
        name = f"{base}.weight{suffix}"
        if t.dtype == torch.float8_e8m0fnu:
            out[name] = t.view(torch.uint8).cpu()
        elif t.dtype in FP8_DTYPES:
            out[name] = t.view(torch.uint8).cpu().view(t.dtype)
        else:
            out[name] = t.cpu()

    blob = json.dumps({"format": "nvfp4"}).encode()
    out[ck_] = torch.tensor(list(blob), dtype=torch.uint8)

    consumed.update({wk, sk, ck_})
    stats["nvfp4"] += 1
    del qdata, scale, deq, nq
    if (i + 1) % 25 == 0:
        torch.cuda.empty_cache()
        print(f"  {i+1}/{len(quant_bases)}  {time.time()-t0:.0f}s", flush=True)

# copy everything untouched
for k in keys:
    if k in consumed or k in out:
        continue
    out[k] = f.get_tensor(k)
    stats["copied"] += 1

print(f"stats: {dict(stats)}", flush=True)
print(f"writing {DST} ...", flush=True)
save_file(out, DST)
gb = os.path.getsize(DST) / 1e9
print(f"DONE {gb:.2f} GB in {time.time()-t0:.0f}s", flush=True)

c = collections.Counter(str(v.dtype) for v in out.values())
print("dtype census:", dict(c), flush=True)