runtime/single_file_transformer.py

"""Load a supported single-file ComfyUI Mage-Flow XPO3 transformer."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn
from safetensors import safe_open

from packed_artifact import (
    assign_tensor_by_name,
    instantiate_mage_transformer_on_meta,
    materialize_mage_rope_tensor_attributes,
    set_child_module,
    unregistered_meta_tensor_attribute_names,
)
from standard_transformer import (
    _apply_runtime_defaults,
    _quantized_keys,
    _target_specs_for_config,
    fail,
)
from torch_ops_native import (
    PackedNvfp4LinearNativeOp,
    initialize_native_sm120_op,
)


def parse_metadata_json(
    metadata: dict[str, str],
    key: str,
) -> dict[str, Any]:
    try:
        value = json.loads(metadata[key])
    except (KeyError, json.JSONDecodeError) as exc:
        fail(f"single-file checkpoint has invalid {key} metadata: {exc}")
    if not isinstance(value, dict):
        fail(f"single-file checkpoint metadata {key} is not an object")
    return value


def load_single_file_native_transformer(
    checkpoint_path: str | Path,
    *,
    support_root: str | Path,
    device: torch.device,
) -> tuple[nn.Module, dict[str, Any]]:
    checkpoint_path = Path(checkpoint_path).resolve()
    support_root = Path(support_root).resolve()
    if not checkpoint_path.is_file():
        fail(f"single-file checkpoint is missing: {checkpoint_path}")
    if device.type != "cuda":
        fail("the native resident transformer requires a CUDA destination")

    with safe_open(
        checkpoint_path,
        framework="pt",
        device="cpu",
    ) as handle:
        metadata = handle.metadata() or {}
        if metadata.get("mage_flow.component") != "transformer":
            fail("single-file checkpoint is not a Mage-Flow transformer")
        supported_variants = {
            "turbo-balanced-v2-fused-qkv",
            "edit-turbo-xpo3-v1-fused-qkv",
        }
        if metadata.get("mage_flow.variant") not in supported_variants:
            fail(
                "unsupported Mage-Flow single-file variant: "
                f"{metadata.get('mage_flow.variant')!r}"
            )
        config = parse_metadata_json(
            metadata,
            "mage_flow.transformer_config",
        )
        quant_config = config.get("quantization_config")
        if not isinstance(quant_config, dict):
            fail("single-file transformer config has no quantization_config")
        if quant_config.get("quant_method") not in {
            "mage_flow_nvfp4",
            "xpo3_nvfp4",
        }:
            fail(
                "single-file transformer does not declare a supported "
                "Mage-Flow/XPO3 NVFP4 method"
            )
        if quant_config.get("quant_algo") != "NVFP4":
            fail("single-file transformer does not declare NVFP4")
        nvfp4_metadata = parse_metadata_json(
            metadata,
            "mage_flow.nvfp4_metadata",
        )
        attention_metadata = parse_metadata_json(
            metadata,
            "mage_flow.attention_nvfp4",
        )

        runtime_defaults = _apply_runtime_defaults(quant_config)
        if not initialize_native_sm120_op(
            allow_python_schema_fallback=False
        ):
            fail("compiled native SM120 torch op did not load")

        depth = int(config.get("depth", 0))
        target_specs = _target_specs_for_config(depth, quant_config)
        expected_targets = [spec.module_key for spec in target_specs]
        recorded_targets = nvfp4_metadata.get("targets")
        if not isinstance(recorded_targets, list):
            fail("single-file NVFP4 metadata has no target list")
        if [
            entry.get("module_key")
            for entry in recorded_targets
            if isinstance(entry, dict)
        ] != expected_targets:
            fail("single-file MLP targets do not match transformer config")

        non_target_keys = nvfp4_metadata.get("non_target_keys")
        if not isinstance(non_target_keys, list) or not all(
            isinstance(key, str) for key in non_target_keys
        ):
            fail("single-file NVFP4 metadata has no non-target key list")

        attention_groups = attention_metadata.get("groups")
        if (
            attention_metadata.get("mode") != "fused_qkv"
            or not isinstance(attention_groups, list)
            or len(attention_groups) != 24
        ):
            fail("single-file attention metadata is not 24 fused QKV groups")

        attention_source_keys = {
            key
            for group in attention_groups
            if isinstance(group, dict)
            for field in ("source_weight_keys", "source_bias_keys")
            for key in group.get(field, [])
            if isinstance(key, str)
        }
        fused_attention_keys = {
            f"{group['module_key']}.{suffix}"
            for group in attention_groups
            if isinstance(group, dict)
            for suffix in (
                "packed_weight",
                "weight_scales",
                "weight_scale",
                "bias",
            )
        }
        quantized_mlp_keys = {
            key
            for spec in target_specs
            for key in _quantized_keys(spec.module_key).values()
        }
        expected_keys = (
            (set(non_target_keys) - attention_source_keys)
            | quantized_mlp_keys
            | fused_attention_keys
        )
        actual_keys = set(handle.keys())
        if actual_keys != expected_keys:
            missing = sorted(expected_keys - actual_keys)
            unexpected = sorted(actual_keys - expected_keys)
            fail(
                "single-file tensor coverage mismatch; "
                f"missing={missing[:1]}, unexpected={unexpected[:1]}"
            )

        model = instantiate_mage_transformer_on_meta(support_root)

        def tensor(key: str) -> torch.Tensor:
            if key not in actual_keys:
                fail(f"tensor is absent from single-file checkpoint: {key}")
            return handle.get_tensor(key)

        for spec in target_specs:
            original = model.get_submodule(spec.module_key)
            if not isinstance(original, nn.Linear):
                fail(
                    f"expected target {spec.module_key} to be nn.Linear, "
                    f"found {type(original).__name__}"
                )
            keys = _quantized_keys(spec.module_key)
            replacement = PackedNvfp4LinearNativeOp(
                in_features=int(original.in_features),
                out_features=int(original.out_features),
                packed_weight=tensor(keys["packed_weight"]).to(device),
                weight_scales=tensor(keys["weight_scales"]).to(device),
                weight_scale=tensor(keys["weight_scale"]).to(device),
                bias=tensor(keys["bias"]).to(device),
            )
            set_child_module(model, spec.module_key, replacement)

        installed_attention: list[str] = []
        for group in attention_groups:
            if not isinstance(group, dict):
                fail("single-file attention group is malformed")
            module_key = str(group.get("module_key", ""))
            parts = module_key.split(".")
            if (
                len(parts) != 4
                or parts[0] != "transformer_blocks"
                or parts[2] != "attn"
                or parts[3] not in {"to_qkv", "add_qkv_proj"}
            ):
                fail(f"invalid fused attention module key: {module_key}")
            block_index = int(parts[1])
            attention = model.transformer_blocks[block_index].attn
            replacement = PackedNvfp4LinearNativeOp(
                in_features=int(group["in_features"]),
                out_features=int(group["out_features"]),
                packed_weight=tensor(
                    f"{module_key}.packed_weight"
                ).to(device),
                weight_scales=tensor(
                    f"{module_key}.weight_scales"
                ).to(device),
                weight_scale=tensor(
                    f"{module_key}.weight_scale"
                ).to(device),
                bias=tensor(f"{module_key}.bias").to(device),
            )
            setattr(attention, parts[3], replacement)
            source_names = (
                ("to_q", "to_k", "to_v")
                if parts[3] == "to_qkv"
                else ("add_q_proj", "add_k_proj", "add_v_proj")
            )
            for source_name in source_names:
                setattr(attention, source_name, None)
            installed_attention.append(module_key)

        retained_keys = sorted(
            set(non_target_keys) - attention_source_keys
        )
        for key in retained_keys:
            assign_tensor_by_name(model, key, tensor(key).to(device))

    materialized = materialize_mage_rope_tensor_attributes(model)
    meta_parameters = [
        name for name, value in model.named_parameters() if value.is_meta
    ]
    meta_buffers = [
        name for name, value in model.named_buffers() if value.is_meta
    ]
    unregistered_meta = unregistered_meta_tensor_attribute_names(model)
    if meta_parameters or meta_buffers or unregistered_meta:
        fail(
            "single-file loader left unresolved meta tensors: "
            f"{(meta_parameters + meta_buffers + unregistered_meta)[:4]}"
        )

    return model.eval().requires_grad_(False), {
        "layout": "comfyui_single_file_fused_qkv",
        "checkpoint": str(checkpoint_path),
        "checkpoint_tensor_count": len(actual_keys),
        "loaded_non_target_tensor_count": len(retained_keys),
        "loaded_quantized_mlp_projection_count": len(target_specs),
        "loaded_fused_attention_projection_count": len(
            installed_attention
        ),
        "bf16_attention_weight_reads": 0,
        "bf16_mlp_target_weight_reads": 0,
        "materialized_unregistered_tensor_attribute_names": materialized,
        "runtime_defaults_applied": runtime_defaults,
    }


__all__ = ["load_single_file_native_transformer"]