"""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", "edit-xpo3-v1-fused-qkv-direct-hnd-steps7-29", } 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"]