"""Generate StarNodes converter profiles for MiniMax-H3 from the bf16 safetensors header.

Two variants:
  minimax_h3_nvfp4_mixed.json  - AdaLN kept at FP8, attn/mlp at NVFP4  (recommended)
  minimax_h3_nvfp4_full.json   - everything incl. AdaLN at NVFP4       (aggressive)

Kept at BF16 in both: all biases, all norms, rope freqs, patch/time/condition
embedders and the final output heads (~0.04B params total, 0.1% of the model).
"""
import json, struct, sys, os, datetime

SRC = "/workspace/ComfyUI/models/diffusion_models/minimax_h3_ref2va_bf16.safetensors"
OUT_DIR = "/workspace/ComfyUI/custom_nodes/comfyui-starnodes-modelconverter/profiles"

with open(SRC, "rb") as f:
    n = struct.unpack("<Q", f.read(8))[0]
    hdr = json.loads(f.read(n))
keys = sorted(k for k in hdr if k != "__metadata__")


def classify(key, adaln_fmt):
    # non-weight tensors and everything tiny stays bf16
    if not key.endswith(".weight"):
        return "BF16"
    if "norm" in key or key.endswith("inv_freq"):
        return "BF16"
    if any(t in key for t in ("patch_proj", "time_embedder", "condition_proj", "final_layer")):
        return "BF16"
    if "adaln" in key:
        return adaln_fmt
    if ".attn." in key or ".mlp." in key:
        return "NVFP4"
    return "BF16"


def build(adaln_fmt, name):
    layers, counts = {}, {}
    for k in keys:
        fmt = classify(k, adaln_fmt)
        layers[k] = fmt
        counts[fmt] = counts.get(fmt, 0) + 1
        if fmt != "BF16":
            base = k[: -len(".weight")]
            layers[f"{base}.weight_scale"] = "FP32_SCALE"
            layers[f"{base}.comfy_quant"] = "METADATA"
    prof = {
        "__metadata__": {
            "original_model_name": name,
            "original_model_path": SRC,
            "timestamp": datetime.datetime.now().isoformat(),
            "total_layers": len(layers),
            "created_by": "hand-authored for MiniMax-H3 (33.12B: adaln 39.4%, mlp 36.3%, attn 24.2%)",
        },
        "layers": layers,
    }
    path = os.path.join(OUT_DIR, f"{name}.json")
    with open(path, "w") as f:
        json.dump(prof, f, indent=1)
    print(f"{name}: {counts}  -> {path}")


build("FP8_E4M3FN + SCALE", "minimax_h3_nvfp4_mixed")
build("NVFP4", "minimax_h3_nvfp4_full")