#!/usr/bin/env python
"""
Offline unit tests for pipelines.native_transformer.

Covers the pure helpers that own per-arch knob handling:

- ``drop_companion_keys`` for filtering bundled TE/VAE families out of all-in-one files
- ``strip_prefix`` for single/multi prefix detection and mixed-prefix rejection
- ``partition_siblings`` for inline-sibling key partitioning
- ``check_forbidden_markers`` for structural-mismatch rejection
- ``detect_comfy_quant`` for comfy_quant marker detection and format gating
- ``remap_comfy_quant`` for comfy_quant -> SDNQ key translation
- ``is_noop_converter`` for diffusers no-op lambda detection
- ``validate_state_dict_load`` for unexpected / missing key handling
- ``make_default_spec`` default-spec synthesis with diffusers converter pickup
- ``auto_pickup_converter`` for diffusers ``SINGLE_FILE_LOADABLE_CLASSES`` integration
- ``TransformerSpec`` / ``SiblingSpec`` defaults

Plus one end-to-end ``load`` test against a tiny mock module that exercises
the read -> strip -> convert -> from_config -> load_state_dict -> validate
pipeline without needing a real diffusers transformer or hf_hub_download.

No running server required.

Usage:
    python test/test-native-transformer.py
"""

import os
import sys
import tempfile

import torch
import safetensors.torch

script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, script_dir)
os.chdir(script_dir)

os.environ['SD_INSTALL_QUIET'] = '1'

# Bootstrap cmd_args before any module that pulls in shared.py.
import modules.cmd_args  # pylint: disable=wrong-import-position
import installer  # pylint: disable=wrong-import-position
_orig_argv = sys.argv
sys.argv = [sys.argv[0]]
try:
    modules.cmd_args.parse_args()
finally:
    sys.argv = _orig_argv
installer.add_args(modules.cmd_args.parser)
modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])

from modules.errors import log                          # pylint: disable=wrong-import-position
from pipelines import native_transformer as nt          # pylint: disable=wrong-import-position


# ============================================================
# Test infrastructure
# ============================================================

results: dict[str, dict] = {}


def category(name: str):
    if name not in results:
        results[name] = {'passed': 0, 'failed': 0, 'tests': []}
    return name


def record(cat: str, passed: bool, name: str, detail: str = ''):
    status = 'PASS' if passed else 'FAIL'
    results[cat]['passed' if passed else 'failed'] += 1
    results[cat]['tests'].append((status, name))
    msg = f'  {status}: {name}'
    if detail:
        msg += f' ({detail})'
    if passed:
        log.info(msg)
    else:
        log.error(msg)


def run_test(cat: str, fn):
    name = fn.__name__
    try:
        ok = fn()
        if ok is False:
            record(cat, False, name)
        else:
            record(cat, True, name)
    except AssertionError as e:
        record(cat, False, name, str(e))
    except Exception as e:  # pylint: disable=broad-except
        record(cat, False, name, f'exception: {e}')
        import traceback
        traceback.print_exc()


# ============================================================
# drop_companion_keys
# ============================================================

def test_drop_companion_keys_no_companions_pass_through():
    sd = {'net.blocks.0.weight': 1, 'net.blocks.1.weight': 2}
    out = nt.drop_companion_keys(sd, nt.DEFAULT_IGNORED_PREFIXES, 'Test')
    assert out == sd, 'files without companion keys must pass through unchanged'


def test_drop_companion_keys_filters_all_in_one_layout():
    sd = {
        'net.blocks.0.weight': 1,
        'net.llm_adapter.proj.weight': 2,
        'cond_stage_model.qwen3_06b.transformer.model.embed_tokens.weight': 3,
        'first_stage_model.decoder.conv1.weight': 4,
        'vae.decoder.conv_in.weight': 5,
        'text_encoders.clip_l.weight': 6,
    }
    out = nt.drop_companion_keys(sd, nt.DEFAULT_IGNORED_PREFIXES, 'Test')
    assert set(out.keys()) == {'net.blocks.0.weight', 'net.llm_adapter.proj.weight'}


def test_drop_companion_keys_raises_when_nothing_left():
    sd = {
        'cond_stage_model.te.weight': 1,
        'first_stage_model.decoder.weight': 2,
    }
    try:
        nt.drop_companion_keys(sd, nt.DEFAULT_IGNORED_PREFIXES, 'Test')
        raise AssertionError('expected ValueError')
    except ValueError as e:
        assert 'no transformer keys' in str(e)


def test_drop_companion_keys_empty_prefix_tuple_no_op():
    sd = {'cond_stage_model.te.weight': 1}
    out = nt.drop_companion_keys(sd, (), 'Test')
    assert out == sd, 'empty ignored_prefixes must disable filtering'


def test_drop_then_strip_all_in_one_layout():
    """Companion filtering must clear the way for normal prefix detection."""
    sd = {
        'net.blocks.0.weight': 1,
        'net.final_layer.weight': 2,
        'cond_stage_model.te.weight': 3,
        'first_stage_model.decoder.weight': 4,
    }
    out = nt.drop_companion_keys(sd, nt.DEFAULT_IGNORED_PREFIXES, 'Test')
    out, prefix = nt.strip_prefix(out, nt.DEFAULT_PREFIXES, 'Test')
    assert prefix == 'net.'
    assert set(out.keys()) == {'blocks.0.weight', 'final_layer.weight'}


# ============================================================
# strip_prefix
# ============================================================

def test_strip_prefix_bare_keys_pass_through():
    sd = {'layers.0.weight': 1, 'layers.0.bias': 2}
    out, prefix = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
    assert out == sd, 'bare keys must pass through unchanged'
    assert prefix == '', 'bare keys must report an empty prefix'


def test_strip_prefix_dominant_single_variant():
    sd = {f'model.diffusion_model.layers.{i}.weight': i for i in range(10)}
    out, prefix = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
    assert all(k.startswith('layers.') for k in out)
    assert len(out) == 10
    assert prefix == 'model.diffusion_model.'


def test_strip_prefix_picks_longest_match_first():
    """``model.diffusion_model.`` must beat ``diffusion_model.`` when both match."""
    sd = {
        'model.diffusion_model.layers.0.weight': 1,
        'model.diffusion_model.layers.1.weight': 2,
    }
    out, prefix = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
    # If shorter prefix matched, keys would start with 'model.'
    assert 'layers.0.weight' in out
    assert 'layers.1.weight' in out
    assert not any(k.startswith('model.') for k in out)
    assert prefix == 'model.diffusion_model.'


def test_strip_prefix_mixed_prefixes_raises():
    sd = {
        'model.diffusion_model.layers.0.weight': 1,
        'net.layers.0.weight': 2,
    }
    try:
        nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
        raise AssertionError('expected ValueError')
    except ValueError as e:
        assert 'mixed prefixes' in str(e)


def test_strip_prefix_net_variant():
    sd = {'net.layers.0.weight': 1, 'net.layers.1.bias': 2}
    out, prefix = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
    assert set(out.keys()) == {'layers.0.weight', 'layers.1.bias'}
    assert prefix == 'net.'


def test_strip_prefix_diffusion_model_variant():
    sd = {'diffusion_model.layers.0.weight': 1, 'diffusion_model.layers.1.bias': 2}
    out, prefix = nt.strip_prefix(sd, nt.DEFAULT_PREFIXES, 'Test')
    assert set(out.keys()) == {'layers.0.weight', 'layers.1.bias'}
    assert prefix == 'diffusion_model.'


def test_strip_prefix_custom_prefix_set():
    sd = {'lora_unet_blocks_0.weight': 1, 'lora_unet_blocks_1.weight': 2}
    out, prefix = nt.strip_prefix(sd, ('lora_unet_',), 'Test')
    assert set(out.keys()) == {'blocks_0.weight', 'blocks_1.weight'}
    assert prefix == 'lora_unet_'


# ============================================================
# partition_siblings
# ============================================================

def test_partition_siblings_empty_spec_returns_state_dict_unchanged():
    sd = {'a': 1, 'b': 2}
    transformer_sd, siblings = nt.partition_siblings(sd, {})
    assert transformer_sd == sd
    assert siblings == {}


def test_partition_siblings_no_matches_keeps_all_in_transformer():
    sd = {'layers.0.weight': 1, 'layers.1.weight': 2}
    siblings_spec = {'llm_adapter': nt.SiblingSpec(subfolder='llm_adapter', inline_prefix='llm_adapter.')}
    transformer_sd, siblings = nt.partition_siblings(sd, siblings_spec)
    assert transformer_sd == sd
    assert siblings == {'llm_adapter': {}}


def test_partition_siblings_single_sibling_split():
    sd = {
        'layers.0.weight': 'tx0',
        'layers.1.weight': 'tx1',
        'llm_adapter.input_proj.weight': 'ad0',
        'llm_adapter.output_proj.weight': 'ad1',
    }
    siblings_spec = {'llm_adapter': nt.SiblingSpec(subfolder='llm_adapter', inline_prefix='llm_adapter.')}
    transformer_sd, siblings = nt.partition_siblings(sd, siblings_spec)
    assert set(transformer_sd.keys()) == {'layers.0.weight', 'layers.1.weight'}
    assert set(siblings['llm_adapter'].keys()) == {'input_proj.weight', 'output_proj.weight'}
    assert siblings['llm_adapter']['input_proj.weight'] == 'ad0'


def test_partition_siblings_multiple_siblings():
    sd = {
        'layers.0.weight': 'tx',
        'sibling_a.x.weight': 'a0',
        'sibling_b.y.weight': 'b0',
        'sibling_b.z.weight': 'b1',
    }
    siblings_spec = {
        'sibling_a': nt.SiblingSpec(subfolder='a', inline_prefix='sibling_a.'),
        'sibling_b': nt.SiblingSpec(subfolder='b', inline_prefix='sibling_b.'),
    }
    transformer_sd, siblings = nt.partition_siblings(sd, siblings_spec)
    assert list(transformer_sd.keys()) == ['layers.0.weight']
    assert set(siblings['sibling_a'].keys()) == {'x.weight'}
    assert set(siblings['sibling_b'].keys()) == {'y.weight', 'z.weight'}


# ============================================================
# check_forbidden_markers
# ============================================================

def test_forbidden_markers_passes_when_absent():
    sd = {'layers.0.weight': 1}
    markers = (('legacy.marker.weight', 'old format'),)
    nt.check_forbidden_markers(sd, markers, 'Test', '/tmp/x.safetensors')
    # no exception = pass


def test_forbidden_markers_raises_when_present():
    sd = {'layers.0.weight': 1, 'legacy.marker.weight': 2}
    markers = (('legacy.marker.weight', 'old Cosmos 1.0 structure'),)
    try:
        nt.check_forbidden_markers(sd, markers, 'Test', '/tmp/x.safetensors')
        raise AssertionError('expected ValueError')
    except ValueError as e:
        msg = str(e)
        assert 'old Cosmos 1.0 structure' in msg
        assert 'legacy.marker.weight' in msg


def test_forbidden_markers_empty_tuple_no_op():
    sd = {'layers.0.weight': 1}
    nt.check_forbidden_markers(sd, (), 'Test', '/tmp/x.safetensors')


# ============================================================
# detect_comfy_quant / remap_comfy_quant
# ============================================================

def comfy_marker(fmt: str) -> torch.Tensor:
    import json
    return torch.tensor(list(json.dumps({'format': fmt}).encode()), dtype=torch.uint8)


def test_detect_comfy_no_markers_returns_none():
    sd = {'blocks.0.attn.wq.weight': torch.zeros(2), 'blocks.0.attn.wq.bias': torch.zeros(2)}
    assert nt.detect_comfy_quant(sd, 'Test') is None


def test_detect_comfy_valid_markers():
    sd = {
        'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8),
        'blocks.0.attn.wq.weight_scale': torch.tensor(0.5),
        'blocks.0.attn.wq.comfy_quant': comfy_marker('int8_tensorwise'),
        'blocks.0.mlp.up.weight': torch.zeros((4, 4), dtype=torch.int8),
        'blocks.0.mlp.up.weight_scale': torch.tensor(0.25),
        'blocks.0.mlp.up.comfy_quant': comfy_marker('int8_tensorwise'),
        'norm.weight': torch.zeros(4),
    }
    detected = nt.detect_comfy_quant(sd, 'Test')
    assert detected is not None
    marked, fmt = detected
    assert set(marked) == {'blocks.0.attn.wq', 'blocks.0.mlp.up'}
    assert marked['blocks.0.attn.wq'] == {'format': 'int8_tensorwise'}
    assert fmt == 'int8_tensorwise'


def test_detect_comfy_unsupported_format_raises():
    sd = {
        'blocks.0.attn.wq.weight': torch.zeros((4, 4)),
        'blocks.0.attn.wq.comfy_quant': comfy_marker('float8_e4m3fn_scaled'),
    }
    try:
        nt.detect_comfy_quant(sd, 'Test')
        raise AssertionError('expected OverrideArchMismatch')
    except nt.OverrideArchMismatch as e:
        assert 'float8_e4m3fn_scaled' in str(e)


def test_detect_comfy_malformed_marker_raises():
    sd = {
        'blocks.0.attn.wq.weight': torch.zeros((4, 4)),
        'blocks.0.attn.wq.comfy_quant': torch.tensor(list(b'not json'), dtype=torch.uint8),
    }
    try:
        nt.detect_comfy_quant(sd, 'Test')
        raise AssertionError('expected OverrideArchMismatch')
    except nt.OverrideArchMismatch as e:
        assert 'malformed' in str(e)


def test_detect_comfy_convrot_accepted():
    """ConvRot markers map onto SDNQ's Hadamard support; detection carries the
    per-layer fields through instead of rejecting."""
    import json
    payload = json.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 256, 'per_row': True})
    sd = {
        'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8),
        'blocks.0.attn.wq.comfy_quant': torch.tensor(list(payload.encode()), dtype=torch.uint8),
        'blocks.0.mlp.up.weight': torch.zeros((4, 4), dtype=torch.int8),
        'blocks.0.mlp.up.comfy_quant': comfy_marker('int8_tensorwise'),
    }
    marked, fmt = nt.detect_comfy_quant(sd, 'Test')
    assert fmt == 'int8_tensorwise'
    assert marked['blocks.0.attn.wq'].get('convrot') is True
    assert marked['blocks.0.attn.wq'].get('convrot_groupsize') == 256
    assert not marked['blocks.0.mlp.up'].get('convrot')


def test_detect_comfy_convrot_bad_groupsize_raises():
    """Regular Hadamards only exist for power-of-4 sizes; any other convrot
    group size cannot be a compatible rotation."""
    import json
    payload = json.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 8})
    sd = {
        'blocks.0.attn.wq.comfy_quant': torch.tensor(list(payload.encode()), dtype=torch.uint8),
    }
    try:
        nt.detect_comfy_quant(sd, 'Test')
        raise AssertionError('expected OverrideArchMismatch')
    except nt.OverrideArchMismatch as e:
        assert 'power of 4' in str(e)


def test_detect_comfy_marker_missing_format_field_raises():
    import json
    sd = {
        'blocks.0.attn.wq.comfy_quant': torch.tensor(list(json.dumps({'fmt': 'x'}).encode()), dtype=torch.uint8),
    }
    try:
        nt.detect_comfy_quant(sd, 'Test')
        raise AssertionError('expected OverrideArchMismatch')
    except nt.OverrideArchMismatch as e:
        assert 'malformed' in str(e)


def test_remap_comfy_renames_and_reshapes_scale():
    sd = {
        'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8),
        'blocks.0.attn.wq.weight_scale': torch.tensor(0.5),
        'blocks.0.attn.wq.comfy_quant': comfy_marker('int8_tensorwise'),
    }
    out = nt.remap_comfy_quant(sd, {'blocks.0.attn.wq'})
    assert set(out.keys()) == {'blocks.0.attn.wq.weight', 'blocks.0.attn.wq.scale'}
    assert out['blocks.0.attn.wq.scale'].shape == (1, 1)
    assert out['blocks.0.attn.wq.scale'].item() == 0.5
    assert out['blocks.0.attn.wq.weight'].dtype == torch.int8


def test_remap_comfy_passes_unmarked_keys_verbatim():
    unrelated_scale = torch.tensor(2.0)
    sd = {
        'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8),
        'blocks.0.attn.wq.weight_scale': torch.tensor(0.5),
        'blocks.0.attn.wq.comfy_quant': comfy_marker('int8_tensorwise'),
        'norm.weight': torch.ones(4),
        'other.weight_scale': unrelated_scale,
    }
    out = nt.remap_comfy_quant(sd, {'blocks.0.attn.wq'})
    assert 'norm.weight' in out
    assert out['other.weight_scale'] is unrelated_scale, 'weight_scale outside marked names must pass through untouched'


def test_remap_comfy_does_not_mutate_input():
    sd = {
        'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8),
        'blocks.0.attn.wq.weight_scale': torch.tensor(0.5),
        'blocks.0.attn.wq.comfy_quant': comfy_marker('int8_tensorwise'),
    }
    keys_before = set(sd.keys())
    nt.remap_comfy_quant(sd, {'blocks.0.attn.wq'})
    assert set(sd.keys()) == keys_before, 'input dict must not be mutated (read_state_dict caches it)'


def test_remap_comfy_drops_input_scale():
    """Optional activation-calibration sidecars are dropped for marked layers
    (SDNQ derives activation scales dynamically); unmarked ones pass through."""
    unrelated = torch.tensor(1.0)
    sd = {
        'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8),
        'blocks.0.attn.wq.weight_scale': torch.tensor(0.5),
        'blocks.0.attn.wq.input_scale': torch.tensor(0.125),
        'blocks.0.attn.wq.comfy_quant': comfy_marker('int8_tensorwise'),
        'other.input_scale': unrelated,
    }
    out = nt.remap_comfy_quant(sd, {'blocks.0.attn.wq'})
    assert 'blocks.0.attn.wq.input_scale' not in out
    assert out['other.input_scale'] is unrelated


# ============================================================
# header _quantization_metadata
# ============================================================

def quant_metadata(layers: dict) -> dict:
    import json
    return {'_quantization_metadata': json.dumps({'format_version': '1.0', 'layers': layers})}


def test_read_quant_metadata_absent_returns_none():
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        write_fixture({'a.weight': torch.zeros(2)}, fd, path)
        assert nt.read_quantization_metadata(path) is None
    finally:
        os.unlink(path)


def test_read_quant_metadata_parses_layers():
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        layers = {'blocks.0.attn.wq': {'format': 'int8_tensorwise'}}
        write_fixture({'blocks.0.attn.wq.weight': torch.zeros(2)}, fd, path, metadata=quant_metadata(layers))
        assert nt.read_quantization_metadata(path) == layers
    finally:
        os.unlink(path)


def test_read_quant_metadata_malformed_raises():
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        write_fixture({'a.weight': torch.zeros(2)}, fd, path, metadata={'_quantization_metadata': 'not json'})
        try:
            nt.read_quantization_metadata(path)
            raise AssertionError('expected OverrideArchMismatch')
        except nt.OverrideArchMismatch as e:
            assert 'malformed' in str(e)
    finally:
        os.unlink(path)


def test_read_quant_metadata_non_dict_layers_raises():
    import json
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        payload = json.dumps({'layers': {'blocks.0.attn.wq': 'int8_tensorwise'}})
        write_fixture({'a.weight': torch.zeros(2)}, fd, path, metadata={'_quantization_metadata': payload})
        try:
            nt.read_quantization_metadata(path)
            raise AssertionError('expected OverrideArchMismatch')
        except nt.OverrideArchMismatch as e:
            assert 'malformed' in str(e)
    finally:
        os.unlink(path)


def test_transcode_quant_metadata_synthesizes_markers():
    """Header entries become marker tensors for layers whose weight is present;
    entries without a matching weight (companion components) are skipped."""
    sd = {
        'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8),
        'blocks.0.attn.wq.weight_scale': torch.tensor(0.5),
    }
    layers = {
        'blocks.0.attn.wq': {'format': 'int8_tensorwise'},
        'text_encoder.mlp.up': {'format': 'int8_tensorwise'},
    }
    out, source = nt.transcode_quant_metadata(sd, layers)
    assert source == 'header'
    assert 'blocks.0.attn.wq.comfy_quant' in out
    assert 'text_encoder.mlp.up.comfy_quant' not in out
    assert 'blocks.0.attn.wq.comfy_quant' not in sd, 'input dict must not be mutated'
    marked, fmt = nt.detect_comfy_quant(out, 'Test')
    assert fmt == 'int8_tensorwise'
    assert marked['blocks.0.attn.wq'] == {'format': 'int8_tensorwise'}


def test_transcode_quant_metadata_header_wins():
    """When both forms are present the header entry overwrites the marker,
    matching the reference loader's precedence."""
    sd = {
        'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8),
        'blocks.0.attn.wq.comfy_quant': comfy_marker('float8_e4m3fn'),
    }
    layers = {'blocks.0.attn.wq': {'format': 'int8_tensorwise', 'full_precision_matrix_mult': True}}
    out, source = nt.transcode_quant_metadata(sd, layers)
    assert source == 'both'
    marked, fmt = nt.detect_comfy_quant(out, 'Test')
    assert fmt == 'int8_tensorwise'
    assert marked['blocks.0.attn.wq'].get('full_precision_matrix_mult') is True


def test_transcode_quant_metadata_no_matches_is_noop():
    sd = {'blocks.0.attn.wq.weight': torch.zeros((4, 4), dtype=torch.int8)}
    out, source = nt.transcode_quant_metadata(sd, {'unrelated.layer': {'format': 'int8_tensorwise'}})
    assert source == 'markers'
    assert out is sd


# ============================================================
# nvfp4: sdnq codec, scale unswizzle, detection
# ============================================================

E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0)


def to_blocked_reference(matrix: torch.Tensor) -> torch.Tensor:
    """cuBLAS block-scaling tile layout as written into nvfp4 containers;
    the loader's unswizzle must invert this exactly."""
    rows, cols = matrix.shape
    row_blocks = -(rows // -128)
    col_blocks = -(cols // -4)
    padded = torch.zeros((row_blocks * 128, col_blocks * 4), dtype=matrix.dtype)
    padded[:rows, :cols] = matrix
    blocks = padded.view(row_blocks, 128, col_blocks, 4).permute(0, 2, 1, 3)
    rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
    return rearranged.reshape(row_blocks * 128, col_blocks * 4)


def test_unswizzle_block_scales_roundtrip():
    for rows, groups in ((128, 4), (8, 3), (200, 10)):
        mat = torch.arange(rows * groups, dtype=torch.float32).reshape(rows, groups)
        back = nt.unswizzle_block_scales(to_blocked_reference(mat), rows, groups)
        assert torch.equal(back, mat), f'unswizzle mismatch for {(rows, groups)}'


def test_nvfp4_codec_ocp_table():
    """SDNQ's float4_e2m1fn decodes OCP FP4 E2M1 exactly, including the
    subnormal codes 1/9 as +/-0.5; nvfp4 containers adopt it directly."""
    from sdnq.packed_float import unpack_float
    packed = torch.tensor([(2 * j) | (((2 * j) + 1) << 4) for j in range(8)], dtype=torch.uint8)
    dec = unpack_float(packed, 'float4_e2m1fn', torch.Size([16]))
    for code in range(16):
        assert float(dec[code]) == E2M1_VALUES[code], f'code {code}: {float(dec[code])} != {E2M1_VALUES[code]}'


def test_nvfp4_pack_ocp_grid_roundtrip():
    """Every OCP grid value survives a pack/unpack round trip exactly
    (subnormals included), and off-grid values land inside the value set.
    Exact nearest-rounding near the grid midpoints is not asserted: the
    packer's staged rounding may resolve boundary values to either side."""
    from sdnq.packed_float import pack_float, unpack_float
    grid = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, 0.0]
    vals = torch.tensor(grid, dtype=torch.float32)
    out = unpack_float(pack_float(vals, 'float4_e2m1fn'), 'float4_e2m1fn', vals.shape)
    assert torch.equal(out, vals), f'grid values must roundtrip exactly: {out.tolist()}'
    off_grid = torch.tensor([0.2, 0.6, 0.9, 1.2, 2.4, 5.5, -0.4, -0.9, -1.7, -3.4, -5.9, 0.05, 0.99, -0.99], dtype=torch.float32)
    out = unpack_float(pack_float(off_grid, 'float4_e2m1fn'), 'float4_e2m1fn', off_grid.shape)
    for v, o in zip(off_grid.tolist(), out.tolist()):
        assert abs(o) in {0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0}, f'{v} -> {o} outside the OCP set'


def test_detect_comfy_nvfp4_marker_accepted():
    import json
    payload = json.dumps({'format': 'nvfp4', 'group_size': 16, 'orig_dtype': 'torch.float16', 'orig_shape': [32, 32]})
    sd = {
        'in_proj.weight': torch.zeros((32, 16), dtype=torch.uint8),
        'in_proj.comfy_quant': torch.tensor(list(payload.encode()), dtype=torch.uint8),
    }
    marked, fmt = nt.detect_comfy_quant(sd, 'Test')
    assert fmt == 'nvfp4'
    assert marked['in_proj']['orig_shape'] == [32, 32]


def test_detect_comfy_nvfp4_bad_group_size_raises():
    import json
    payload = json.dumps({'format': 'nvfp4', 'group_size': 32})
    sd = {'a.comfy_quant': torch.tensor(list(payload.encode()), dtype=torch.uint8)}
    try:
        nt.detect_comfy_quant(sd, 'Test')
        raise AssertionError('expected OverrideArchMismatch')
    except nt.OverrideArchMismatch as e:
        assert 'group_size' in str(e)


def test_detect_comfy_nvfp4_convrot_raises():
    import json
    payload = json.dumps({'format': 'nvfp4', 'convrot': True, 'convrot_groupsize': 16})
    sd = {'a.comfy_quant': torch.tensor(list(payload.encode()), dtype=torch.uint8)}
    try:
        nt.detect_comfy_quant(sd, 'Test')
        raise AssertionError('expected OverrideArchMismatch')
    except nt.OverrideArchMismatch as e:
        assert 'convrot' in str(e)


# ============================================================
# model_probe quant detection (header-only, mirrors the loader's coverage)
# ============================================================

def test_probe_detects_header_quant_metadata():
    import json
    from modules import model_probe
    header = {
        '__metadata__': {'_quantization_metadata': json.dumps({'layers': {'blocks.0.wq': {'format': 'float8_e4m3fn'}}})},
        'blocks.0.wq.weight': {'dtype': 'F8_E4M3', 'shape': [4, 4]},
        'blocks.0.wq.weight_scale': {'dtype': 'F32', 'shape': []},
    }
    quant = model_probe.analyze_header(header)['quant']
    assert quant == {'scheme': 'comfy_quant', 'format': 'float8_e4m3fn', 'marked_layers': 1, 'source': 'header'}


def test_probe_resolves_nvfp4_from_weight_dtype():
    from modules import model_probe
    header = {
        'model.diffusion_model.blocks.0.wq.weight': {'dtype': 'U8', 'shape': [32, 16]},
        'model.diffusion_model.blocks.0.wq.weight_scale': {'dtype': 'F8_E4M3', 'shape': [128, 4]},
        'model.diffusion_model.blocks.0.wq.weight_scale_2': {'dtype': 'F32', 'shape': []},
        'model.diffusion_model.blocks.0.wq.comfy_quant': {'dtype': 'U8', 'shape': [60]},
    }
    quant = model_probe.analyze_header(header)['quant']
    assert quant['format'] == 'nvfp4'
    assert quant['source'] == 'weight-dtype'


def test_probe_fp8_markers_despite_u8_noise():
    """The marked layer's own weight dtype decides the format; unrelated U8
    tensors must not outvote it."""
    from modules import model_probe
    header = {
        'blocks.0.wq.weight': {'dtype': 'F8_E4M3', 'shape': [8, 8]},
        'blocks.0.wq.weight_scale': {'dtype': 'F32', 'shape': []},
        'blocks.0.wq.comfy_quant': {'dtype': 'U8', 'shape': [27]},
        'extra.blob1': {'dtype': 'U8', 'shape': [64]},
        'extra.blob2': {'dtype': 'U8', 'shape': [64]},
    }
    quant = model_probe.analyze_header(header)['quant']
    assert quant['format'] == 'float8_e4m3fn'


def test_probe_int8_markers():
    from modules import model_probe
    header = {
        'blocks.0.wq.weight': {'dtype': 'I8', 'shape': [8, 8]},
        'blocks.0.wq.weight_scale': {'dtype': 'F32', 'shape': []},
        'blocks.0.wq.comfy_quant': {'dtype': 'U8', 'shape': [29]},
    }
    quant = model_probe.analyze_header(header)['quant']
    assert quant['format'] == 'int8_tensorwise'
    assert quant['marked_layers'] == 1


def test_probe_plain_file_no_quant():
    from modules import model_probe
    header = {'blocks.0.wq.weight': {'dtype': 'BF16', 'shape': [8, 8]}}
    assert model_probe.analyze_header(header)['quant']['scheme'] is None


# ============================================================
# is_noop_converter
# ============================================================

def test_noop_converter_identity_lambda():
    fn = lambda checkpoint, **kwargs: checkpoint  # pylint: disable=unnecessary-lambda-assignment
    assert nt.is_noop_converter(fn) is True


def test_noop_converter_real_function():
    def real(checkpoint, **kwargs):  # pylint: disable=unused-argument
        return {k.replace('a.', 'b.'): v for k, v in checkpoint.items()}
    assert nt.is_noop_converter(real) is False


def test_noop_converter_lambda_with_modification():
    fn = lambda checkpoint, **kwargs: {k: v.float() for k, v in checkpoint.items()}  # pylint: disable=unnecessary-lambda-assignment
    assert nt.is_noop_converter(fn) is False


# ============================================================
# validate_state_dict_load
# ============================================================

def test_validate_accepts_buffer_only_missing():
    nt.validate_state_dict_load(
        component_name='transformer',
        missing=['rope.freqs', 'pos_embedder.pos'],
        unexpected=[],
        acceptable_missing=('rope.', 'pos_embedder.'),
    )


def test_validate_rejects_unexpected():
    try:
        nt.validate_state_dict_load(
            component_name='transformer',
            missing=[],
            unexpected=['some.junk.weight'],
            acceptable_missing=(),
        )
        raise AssertionError('expected OverrideArchMismatch')
    except nt.OverrideArchMismatch as e:
        assert 'unexpected' in str(e)
        assert 'some.junk.weight' in str(e)


def test_validate_rejects_hard_missing():
    try:
        nt.validate_state_dict_load(
            component_name='transformer',
            missing=['layers.0.weight', 'rope.freqs'],
            unexpected=[],
            acceptable_missing=('rope.',),
        )
        raise AssertionError('expected OverrideArchMismatch')
    except nt.OverrideArchMismatch as e:
        msg = str(e)
        assert 'missing' in msg
        assert 'layers.0.weight' in msg
        # Buffer-only missing must not show up in the hard-missing list
        assert msg.count('rope.freqs') == 0


def test_validate_empty_passes():
    nt.validate_state_dict_load(
        component_name='transformer',
        missing=[],
        unexpected=[],
        acceptable_missing=(),
    )


# ============================================================
# make_default_spec
# ============================================================

class FakeTransformer:
    """Minimal stand-in for a diffusers transformer class."""


def test_make_default_spec_for_unknown_class():
    spec = nt.make_default_spec(FakeTransformer)
    assert spec.cls is FakeTransformer
    assert spec.subfolder == 'transformer'
    assert spec.prefixes == nt.DEFAULT_PREFIXES
    assert spec.converter is None  # no diffusers entry for FakeTransformer
    assert spec.siblings == {}
    assert spec.forbidden_markers == ()


def test_make_default_spec_picks_up_real_diffusers_converter():
    import diffusers
    spec = nt.make_default_spec(diffusers.FluxTransformer2DModel)
    assert spec.converter is not None
    assert spec.converter.__name__ == 'convert_flux_transformer_checkpoint_to_diffusers'


def test_make_default_spec_skips_qwen_image_noop():
    """QwenImageTransformer2DModel's diffusers entry is a no-op lambda; the
    default spec must NOT pick it up, leaving converter=None so the caller
    sees only their own (potentially absent) override."""
    import diffusers
    spec = nt.make_default_spec(diffusers.QwenImageTransformer2DModel)
    assert spec.converter is None


# ============================================================
# auto_pickup_converter
# ============================================================

def test_auto_pickup_returns_none_for_unknown_class():
    assert nt.auto_pickup_converter(FakeTransformer) is None


def test_auto_pickup_returns_real_diffusers_converter():
    import diffusers
    fn = nt.auto_pickup_converter(diffusers.FluxTransformer2DModel)
    assert fn is not None
    assert callable(fn)
    assert fn.__name__ == 'convert_flux_transformer_checkpoint_to_diffusers'


def test_auto_pickup_skips_noop_converter_qwen():
    """QwenImageTransformer2DModel registers a no-op lambda in diffusers;
    auto_pickup_converter must return None so the spec falls back to no
    converter (the user-registered spec can override with a real converter)."""
    import diffusers
    assert nt.auto_pickup_converter(diffusers.QwenImageTransformer2DModel) is None


# ============================================================
# TransformerSpec / SiblingSpec defaults
# ============================================================

def test_transformer_spec_defaults():
    spec = nt.TransformerSpec(cls=FakeTransformer)
    assert spec.subfolder == 'transformer'
    assert spec.prefixes == ('model.diffusion_model.', 'diffusion_model.', 'net.')
    assert spec.converter is None
    assert spec.siblings == {}
    assert spec.acceptable_missing == ('rope.', 'pos_embedder.', 'learnable_pos_embed.')
    assert spec.ignored_prefixes == ('cond_stage_model.', 'conditioner.', 'first_stage_model.', 'text_encoders.', 'vae.')
    assert spec.forbidden_markers == ()
    assert spec.infer_config is None


def test_sibling_spec_defaults():
    spec = nt.SiblingSpec(subfolder='llm_adapter', inline_prefix='llm_adapter.')
    assert spec.subfolder == 'llm_adapter'
    assert spec.inline_prefix == 'llm_adapter.'
    assert spec.acceptable_missing == ()


def test_transformer_spec_is_frozen():
    spec = nt.TransformerSpec(cls=FakeTransformer)
    try:
        spec.subfolder = 'changed'  # type: ignore[misc]
    except Exception as e:  # pylint: disable=broad-except
        assert 'FrozenInstanceError' in type(e).__name__ or 'frozen' in str(e).lower()
        return
    raise AssertionError('expected FrozenInstanceError')


# ============================================================
# Integration: end-to-end load() with a tiny mock module
# ============================================================
# We sidestep diffusers + hf_hub_download by patching:
#   - ``fetch_component_config`` to return a hand-rolled config dict
#   - ``model_quant.get_dit_args`` / ``model_quant.get_quant_type`` / quant
#     application to no-ops (we only want to test the load path itself).
# The mock cls is a torch.nn.Module subclass whose ``from_config`` constructs
# a fresh module of the expected shape; ``load_state_dict`` is the standard
# PyTorch method.

class MockMiniTransformer(torch.nn.Module):
    """Tiny stand-in: linear in -> linear out, plus a nested rope sub-module
    holding a buffer the trainer state dict won't carry. Nested mirrors how
    real DiTs structure rope / pos_embedder buffers."""

    @classmethod
    def from_config(cls, config: dict) -> 'MockMiniTransformer':
        return cls(dim=config['dim'])

    def __init__(self, dim: int):
        super().__init__()
        self.in_proj = torch.nn.Linear(dim, dim)
        self.out_proj = torch.nn.Linear(dim, dim)
        self.rope = torch.nn.Module()
        self.rope.register_buffer('freqs', torch.zeros(dim))


class MockKwargsTransformer(MockMiniTransformer):
    """Records the kwargs from_config received, so a test can assert the native
    path forwards caller kwargs to construction. Mirrors diffusers from_config,
    which accepts **kwargs."""

    last_kwargs: dict = {}

    @classmethod
    def from_config(cls, config: dict, **kwargs) -> 'MockKwargsTransformer':
        cls.last_kwargs = dict(kwargs)
        return cls(dim=config['dim'])


class MockLayeredTransformer(torch.nn.Module):
    """Depth comes from config, mirroring ``num_layers`` on real DiTs."""

    @classmethod
    def from_config(cls, config: dict) -> 'MockLayeredTransformer':
        return cls(dim=config['dim'], num_layers=config['num_layers'])

    def __init__(self, dim: int, num_layers: int):
        super().__init__()
        self.blocks = torch.nn.ModuleList([torch.nn.Linear(dim, dim) for _ in range(num_layers)])


def infer_num_layers(state_dict: dict) -> dict:
    indices = [int(k.split('.')[1]) for k in state_dict if k.startswith('blocks.')]
    return {'num_layers': max(indices) + 1} if indices else {}


def write_fixture(state_dict_keys: dict, fd: int, path: str, metadata: dict | None = None) -> str:
    os.close(fd)
    safetensors.torch.save_file(state_dict_keys, path, metadata=metadata)
    return path


def test_load_end_to_end_with_bfl_prefix_no_converter():
    """Exercise the full load pipeline: read .safetensors, strip prefix,
    no converter, instantiate via from_config, load weights, validate.
    """
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        # Save with model.diffusion_model. prefix; in_proj.* and out_proj.*
        # are the real weights the mock cls expects after the strip.
        raw = {
            'model.diffusion_model.in_proj.weight': torch.randn(dim, dim),
            'model.diffusion_model.in_proj.bias': torch.zeros(dim),
            'model.diffusion_model.out_proj.weight': torch.randn(dim, dim),
            'model.diffusion_model.out_proj.bias': torch.zeros(dim),
        }
        write_fixture(raw, fd, path)

        # Patch fetch_component_config to return our hand-rolled config.
        orig_fetch = nt.fetch_component_config
        nt.fetch_component_config = lambda repo, sub: {'dim': dim}

        # Patch quant helpers (we only care about the load path).
        from modules import model_quant
        orig_get_dit = model_quant.get_dit_args
        orig_get_qtype = model_quant.get_quant_type
        orig_do_post = model_quant.do_post_load_quant
        model_quant.get_dit_args = lambda *a, **k: ({}, {})
        model_quant.get_quant_type = lambda *a, **k: None
        model_quant.do_post_load_quant = lambda *a, **k: None

        try:
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            transformer, siblings = nt.load(
                local_file=path,
                repo_id='fake/repo',
                spec=spec,
                diffusers_cfg={},
            )
        finally:
            nt.fetch_component_config = orig_fetch
            model_quant.get_dit_args = orig_get_dit
            model_quant.get_quant_type = orig_get_qtype
            model_quant.do_post_load_quant = orig_do_post

        assert isinstance(transformer, MockMiniTransformer)
        assert transformer.in_proj.weight.shape == (dim, dim)
        # Weights from the fixture should match what was loaded.
        loaded_in_w = transformer.in_proj.weight.detach().cpu()
        fixture_in_w = raw['model.diffusion_model.in_proj.weight'].to(loaded_in_w.dtype)
        assert torch.allclose(loaded_in_w, fixture_in_w)
        assert siblings == {}
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_forwards_kwargs_to_from_config():
    """Caller **kwargs reach cls.from_config through the native load path
    rather than being dropped."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        raw = {
            'model.diffusion_model.in_proj.weight': torch.randn(dim, dim),
            'model.diffusion_model.in_proj.bias': torch.zeros(dim),
            'model.diffusion_model.out_proj.weight': torch.randn(dim, dim),
            'model.diffusion_model.out_proj.bias': torch.zeros(dim),
        }
        write_fixture(raw, fd, path)

        orig_fetch = nt.fetch_component_config
        nt.fetch_component_config = lambda repo, sub: {'dim': dim}
        from modules import model_quant
        orig_get_dit = model_quant.get_dit_args
        orig_get_qtype = model_quant.get_quant_type
        orig_do_post = model_quant.do_post_load_quant
        model_quant.get_dit_args = lambda *a, **k: ({}, {})
        model_quant.get_quant_type = lambda *a, **k: None
        model_quant.do_post_load_quant = lambda *a, **k: None

        MockKwargsTransformer.last_kwargs = {}
        try:
            spec = nt.TransformerSpec(cls=MockKwargsTransformer)
            transformer, _ = nt.load(
                local_file=path,
                repo_id='fake/repo',
                spec=spec,
                diffusers_cfg={},
                low_cpu_mem_usage=True,
            )
        finally:
            nt.fetch_component_config = orig_fetch
            model_quant.get_dit_args = orig_get_dit
            model_quant.get_quant_type = orig_get_qtype
            model_quant.do_post_load_quant = orig_do_post

        assert MockKwargsTransformer.last_kwargs == {'low_cpu_mem_usage': True}
        assert isinstance(transformer, MockKwargsTransformer)
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_infer_config_sizes_module_to_file():
    """A file deeper than the base config loads only with the hook; without it
    the extra block is an arch mismatch."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        file_layers = 3
        raw = {}
        for i in range(file_layers):
            raw[f'model.diffusion_model.blocks.{i}.weight'] = torch.randn(dim, dim)
            raw[f'model.diffusion_model.blocks.{i}.bias'] = torch.zeros(dim)
        write_fixture(raw, fd, path)

        orig_fetch = nt.fetch_component_config
        nt.fetch_component_config = lambda repo, sub: {'dim': dim, 'num_layers': 2}
        from modules import model_quant
        orig_get_dit = model_quant.get_dit_args
        orig_get_qtype = model_quant.get_quant_type
        orig_do_post = model_quant.do_post_load_quant
        model_quant.get_dit_args = lambda *a, **k: ({}, {})
        model_quant.get_quant_type = lambda *a, **k: None
        model_quant.do_post_load_quant = lambda *a, **k: None

        try:
            try:
                nt.load(local_file=path, repo_id='fake/repo', spec=nt.TransformerSpec(cls=MockLayeredTransformer), diffusers_cfg={})
                raise AssertionError('expected OverrideArchMismatch without infer_config')
            except nt.OverrideArchMismatch:
                pass
            spec = nt.TransformerSpec(cls=MockLayeredTransformer, infer_config=infer_num_layers)
            transformer, _ = nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={})
        finally:
            nt.fetch_component_config = orig_fetch
            model_quant.get_dit_args = orig_get_dit
            model_quant.get_quant_type = orig_get_qtype
            model_quant.do_post_load_quant = orig_do_post

        assert len(transformer.blocks) == file_layers
        loaded_w = transformer.blocks[2].weight.detach().cpu()
        assert torch.allclose(loaded_w, raw['model.diffusion_model.blocks.2.weight'].to(loaded_w.dtype))
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_end_to_end_with_sibling_partition():
    """Bundled-sibling case: file carries both transformer and sibling weights,
    sibling_classes supplies the runtime sibling class, partition routes each
    half into its target."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        sibling_dim = 4
        raw = {
            # Transformer half (after strip).
            'model.diffusion_model.in_proj.weight': torch.randn(dim, dim),
            'model.diffusion_model.in_proj.bias': torch.zeros(dim),
            'model.diffusion_model.out_proj.weight': torch.randn(dim, dim),
            'model.diffusion_model.out_proj.bias': torch.zeros(dim),
            # Sibling half (after strip + sibling partition).
            'model.diffusion_model.sibling.in_proj.weight': torch.randn(sibling_dim, sibling_dim),
            'model.diffusion_model.sibling.in_proj.bias': torch.zeros(sibling_dim),
            'model.diffusion_model.sibling.out_proj.weight': torch.randn(sibling_dim, sibling_dim),
            'model.diffusion_model.sibling.out_proj.bias': torch.zeros(sibling_dim),
        }
        write_fixture(raw, fd, path)

        orig_fetch = nt.fetch_component_config

        def patched_fetch(_repo, sub):
            return {'dim': dim if sub == 'transformer' else sibling_dim}

        nt.fetch_component_config = patched_fetch

        from modules import model_quant
        orig_get_dit = model_quant.get_dit_args
        orig_get_qtype = model_quant.get_quant_type
        orig_do_post = model_quant.do_post_load_quant
        model_quant.get_dit_args = lambda *a, **k: ({}, {})
        model_quant.get_quant_type = lambda *a, **k: None
        model_quant.do_post_load_quant = lambda *a, **k: None

        try:
            spec = nt.TransformerSpec(
                cls=MockMiniTransformer,
                siblings={
                    'sibling': nt.SiblingSpec(
                        subfolder='sibling',
                        inline_prefix='sibling.',
                        acceptable_missing=('rope.',),
                    ),
                },
            )
            transformer, siblings = nt.load(
                local_file=path,
                repo_id='fake/repo',
                spec=spec,
                diffusers_cfg={},
                sibling_classes={'sibling': MockMiniTransformer},
            )
        finally:
            nt.fetch_component_config = orig_fetch
            model_quant.get_dit_args = orig_get_dit
            model_quant.get_quant_type = orig_get_qtype
            model_quant.do_post_load_quant = orig_do_post

        assert isinstance(transformer, MockMiniTransformer)
        assert transformer.in_proj.weight.shape == (dim, dim)
        assert 'sibling' in siblings
        assert isinstance(siblings['sibling'], MockMiniTransformer)
        assert siblings['sibling'].in_proj.weight.shape == (sibling_dim, sibling_dim)
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_raises_on_missing_sibling_class():
    """Sibling keys present in file but caller forgot to supply the class."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        sibling_dim = 4
        raw = {
            'model.diffusion_model.in_proj.weight': torch.randn(dim, dim),
            'model.diffusion_model.in_proj.bias': torch.zeros(dim),
            'model.diffusion_model.out_proj.weight': torch.randn(dim, dim),
            'model.diffusion_model.out_proj.bias': torch.zeros(dim),
            'model.diffusion_model.sibling.in_proj.weight': torch.randn(sibling_dim, sibling_dim),
        }
        write_fixture(raw, fd, path)

        orig_fetch = nt.fetch_component_config
        nt.fetch_component_config = lambda repo, sub: {'dim': dim}

        from modules import model_quant
        orig_get_dit = model_quant.get_dit_args
        orig_get_qtype = model_quant.get_quant_type
        orig_do_post = model_quant.do_post_load_quant
        model_quant.get_dit_args = lambda *a, **k: ({}, {})
        model_quant.get_quant_type = lambda *a, **k: None
        model_quant.do_post_load_quant = lambda *a, **k: None

        try:
            spec = nt.TransformerSpec(
                cls=MockMiniTransformer,
                siblings={'sibling': nt.SiblingSpec(subfolder='s', inline_prefix='sibling.')},
            )
            raised = False
            try:
                nt.load(
                    local_file=path,
                    repo_id='fake/repo',
                    spec=spec,
                    diffusers_cfg={},
                    sibling_classes={},  # missing!
                )
            except ValueError as e:
                raised = True
                assert "'sibling'" in str(e)
                assert 'sibling_classes' in str(e)
            assert raised, 'expected ValueError'
        finally:
            nt.fetch_component_config = orig_fetch
            model_quant.get_dit_args = orig_get_dit
            model_quant.get_quant_type = orig_get_qtype
            model_quant.do_post_load_quant = orig_do_post
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_rejects_non_safetensors():
    spec = nt.TransformerSpec(cls=MockMiniTransformer)
    try:
        nt.load(
            local_file='/tmp/some.gguf',
            repo_id='fake/repo',
            spec=spec,
            diffusers_cfg={},
        )
        raise AssertionError('expected ValueError')
    except ValueError as e:
        assert '.safetensors' in str(e)


def crashing_converter(sd):
    """diffusers-style layer count that blows up when the block family is
    absent, mirroring convert_chroma_..._to_diffusers on a wrong-arch file."""
    return list(set(int(k.split('.')[1]) for k in sd if 'double_blocks.' in k))[-1]


def test_build_component_converter_crash_raises_mismatch():
    """A converter that crashes on wrong-arch keys is wrapped as
    OverrideArchMismatch (chaining the original), not the raw IndexError."""
    try:
        nt.build_component(
            component_name='transformer',
            state_dict={'blocks.0.self_attn.weight': torch.zeros(2)},
            config={'dim': 8},
            cls=MockMiniTransformer,
            converter=crashing_converter,
            acceptable_missing=(),
            quant_args={},
            quant_type=None,
        )
        raise AssertionError('expected OverrideArchMismatch')
    except nt.OverrideArchMismatch as e:
        assert 'MockMiniTransformer' in str(e)
        assert isinstance(e.__cause__, IndexError), 'original error must be chained'


def test_build_component_shape_mismatch_is_hard_error():
    """A tensor shape mismatch on otherwise-matching keys stays a hard
    RuntimeError (the native size-mismatch message), it is NOT converted to
    OverrideArchMismatch and so does not silently fall back to base."""
    # MockMiniTransformer(dim=8) expects (8, 8) projections; feed (4, 4).
    sd = {
        'in_proj.weight': torch.randn(4, 4), 'in_proj.bias': torch.zeros(4),
        'out_proj.weight': torch.randn(4, 4), 'out_proj.bias': torch.zeros(4),
    }
    orig_display = nt.errors.display
    nt.errors.display = lambda *a, **k: None  # silence the expected traceback dump
    try:
        nt.build_component(
            component_name='transformer', state_dict=sd, config={'dim': 8},
            cls=MockMiniTransformer, converter=None, acceptable_missing=(),
            quant_args={}, quant_type=None,
        )
        raise AssertionError('expected RuntimeError')
    except nt.OverrideArchMismatch:
        raise AssertionError('shape mismatch must not be OverrideArchMismatch') from None
    except RuntimeError as e:
        assert 'size mismatch' in str(e).lower()
    finally:
        nt.errors.display = orig_display


def test_load_converter_crash_raises_mismatch():
    """End-to-end: a crashing converter surfaces from load() as
    OverrideArchMismatch so load_transformer can drop the override."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        raw = {'model.diffusion_model.blocks.0.self_attn.weight': torch.zeros(8, 8)}
        write_fixture(raw, fd, path)
        orig_fetch = nt.fetch_component_config
        nt.fetch_component_config = lambda repo, sub: {'dim': 8}
        from modules import model_quant
        orig_get_dit = model_quant.get_dit_args
        orig_get_qtype = model_quant.get_quant_type
        model_quant.get_dit_args = lambda *a, **k: ({}, {})
        model_quant.get_quant_type = lambda *a, **k: None
        try:
            spec = nt.TransformerSpec(cls=MockMiniTransformer, converter=crashing_converter)
            raised = False
            try:
                nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={})
            except nt.OverrideArchMismatch as e:
                raised = True
                assert 'MockMiniTransformer' in str(e)
            assert raised, 'expected OverrideArchMismatch'
        finally:
            nt.fetch_component_config = orig_fetch
            model_quant.get_dit_args = orig_get_dit
            model_quant.get_quant_type = orig_get_qtype
    finally:
        if os.path.exists(path):
            os.unlink(path)


# ============================================================
# Ideogram 4 converter (fused community layout -> diffusers layout)
# ============================================================

def ideogram_converter():
    from pipelines.ideogram import convert_ideogram4_transformer_checkpoint
    return convert_ideogram4_transformer_checkpoint


def test_ideogram4_converter_splits_fused_qkv_weight():
    convert = ideogram_converter()
    fused = torch.arange(48, dtype=torch.float32).reshape(12, 4)
    out = convert({'layers.0.attention.qkv.weight': fused})
    assert set(out.keys()) == {f'layers.0.attention.{n}.weight' for n in ('to_q', 'to_k', 'to_v')}
    assert torch.equal(out['layers.0.attention.to_q.weight'], fused[0:4])
    assert torch.equal(out['layers.0.attention.to_k.weight'], fused[4:8])
    assert torch.equal(out['layers.0.attention.to_v.weight'], fused[8:12])


def test_ideogram4_converter_splits_rowwise_scale():
    convert = ideogram_converter()
    fused = torch.zeros((12, 4), dtype=torch.int8)
    scale = torch.arange(12, dtype=torch.float32).reshape(12, 1)
    out = convert({'layers.3.attention.qkv.weight': fused, 'layers.3.attention.qkv.weight_scale': scale})
    assert torch.equal(out['layers.3.attention.to_q.weight_scale'], scale[0:4])
    assert torch.equal(out['layers.3.attention.to_v.weight_scale'], scale[8:12])
    assert out['layers.3.attention.to_k.weight_scale'].shape == (4, 1)


def test_ideogram4_converter_scalar_scale_copied_not_sliced():
    convert = ideogram_converter()
    fused = torch.zeros((12, 4), dtype=torch.int8)
    scale = torch.tensor(0.5)
    out = convert({'layers.0.attention.qkv.weight': fused, 'layers.0.attention.qkv.weight_scale': scale})
    for name in ('to_q', 'to_k', 'to_v'):
        assert out[f'layers.0.attention.{name}.weight_scale'] is scale


def test_ideogram4_converter_duplicates_marker():
    convert = ideogram_converter()
    fused = torch.zeros((12, 4), dtype=torch.int8)
    marker = comfy_marker('int8_tensorwise')
    out = convert({'layers.0.attention.qkv.weight': fused, 'layers.0.attention.qkv.comfy_quant': marker})
    for name in ('to_q', 'to_k', 'to_v'):
        assert out[f'layers.0.attention.{name}.comfy_quant'] is marker


def test_ideogram4_converter_renames_o_with_sidecars():
    convert = ideogram_converter()
    sd = {
        'layers.5.attention.o.weight': torch.zeros((4, 4), dtype=torch.int8),
        'layers.5.attention.o.weight_scale': torch.tensor(1.0),
        'layers.5.attention.o.comfy_quant': comfy_marker('int8_tensorwise'),
    }
    out = convert(sd)
    assert set(out.keys()) == {
        'layers.5.attention.to_out.0.weight',
        'layers.5.attention.to_out.0.weight_scale',
        'layers.5.attention.to_out.0.comfy_quant',
    }


def test_ideogram4_converter_passthrough():
    convert = ideogram_converter()
    sd = {
        'input_proj.weight': torch.zeros(4),
        'llm_cond_proj.weight': torch.zeros(4),
        'layers.0.feed_forward.w1.weight': torch.zeros(4),
        'layers.0.attention.norm_q.weight': torch.zeros(4),
        'layers.0.adaln_modulation.bias': torch.zeros(4),
        'final_layer.linear.weight': torch.zeros(4),
    }
    out = convert(sd)
    assert set(out.keys()) == set(sd.keys())
    for k, v in sd.items():
        assert out[k] is v


def test_ideogram4_converter_defensive_bias_split():
    convert = ideogram_converter()
    fused = torch.zeros((12, 4), dtype=torch.float32)
    bias = torch.arange(12, dtype=torch.float32)
    out = convert({'layers.0.attention.qkv.weight': fused, 'layers.0.attention.qkv.bias': bias})
    assert torch.equal(out['layers.0.attention.to_k.bias'], bias[4:8])


def test_ideogram4_converter_does_not_mutate_input():
    convert = ideogram_converter()
    sd = {
        'layers.0.attention.qkv.weight': torch.zeros((12, 4)),
        'layers.0.attention.o.weight': torch.zeros((4, 4)),
        'input_proj.weight': torch.zeros(4),
    }
    keys_before = set(sd.keys())
    convert(sd)
    assert set(sd.keys()) == keys_before


# ============================================================
# Integration: comfy_quant pre-quantized load
# ============================================================
# Same patching strategy as the plain load tests, plus pinned SDNQ opts so the
# builder's settings reads are deterministic. dtype=float32 makes the dequant
# math bit-exact against a manual int8 * scale reference.

def comfy_fixture(dim: int, fmt: str = 'int8_tensorwise') -> dict:
    """comfy_quant export shape: in_proj carries a quantized weight + scalar
    fp32 scale + marker; out_proj and biases stay plain f16."""
    if fmt == 'float8_e4m3fn':
        weight = torch.randn(dim, dim).to(torch.float8_e4m3fn)
    else:
        weight = torch.randint(-128, 127, (dim, dim), dtype=torch.int8)
    return {
        'model.diffusion_model.in_proj.weight': weight,
        'model.diffusion_model.in_proj.weight_scale': torch.tensor(0.03125, dtype=torch.float32),
        'model.diffusion_model.in_proj.comfy_quant': comfy_marker(fmt),
        'model.diffusion_model.in_proj.bias': torch.randn(dim, dtype=torch.float16),
        'model.diffusion_model.out_proj.weight': torch.randn(dim, dim, dtype=torch.float16),
        'model.diffusion_model.out_proj.bias': torch.zeros(dim, dtype=torch.float16),
    }


class ComfyTestEnv:
    """Patches quant helpers + config fetch and pins the SDNQ opts the
    prequantized builder reads, restoring everything on exit.
    ``fp8_compile_supported`` pins the fp8 storage-remap decision so tests
    behave identically regardless of the host GPU and compile settings."""

    def __init__(self, dim: int, fp8_compile_supported: bool = True):
        self.dim = dim
        self.fp8_compile_supported = fp8_compile_supported

    def __enter__(self):
        from modules import model_quant, shared
        from sdnq import common as sdnq_common
        from sdnq import kernel_wrappers
        self.shared = shared
        self.sdnq_common = sdnq_common
        self.kernel_wrappers = kernel_wrappers
        self.orig_fetch = nt.fetch_component_config
        self.orig_get_dit = model_quant.get_dit_args
        self.orig_get_qtype = model_quant.get_quant_type
        self.orig_do_post = model_quant.do_post_load_quant
        self.orig_fp8_supported = kernel_wrappers.is_fp8_compile_supported
        self.orig_check_compile = sdnq_common.check_torch_compile
        self.model_quant = model_quant
        nt.fetch_component_config = lambda repo, sub: {'dim': self.dim}
        model_quant.get_dit_args = lambda *a, **k: ({}, {})
        model_quant.get_quant_type = lambda *a, **k: None
        model_quant.do_post_load_quant = lambda *a, **k: None
        kernel_wrappers.is_fp8_compile_supported = self.fp8_compile_supported
        sdnq_common.check_torch_compile = lambda: not self.fp8_compile_supported
        self.orig_opts = {
            'sdnq_quantize_matmul_mode': shared.opts.sdnq_quantize_matmul_mode,
            'sdnq_dequantize_fp32': shared.opts.sdnq_dequantize_fp32,
            'diffusers_offload_mode': shared.opts.diffusers_offload_mode,
        }
        shared.opts.data['sdnq_quantize_matmul_mode'] = 'disabled'
        shared.opts.data['sdnq_dequantize_fp32'] = True
        shared.opts.data['diffusers_offload_mode'] = 'model'  # force CPU placement
        return self

    def __exit__(self, *exc):
        nt.fetch_component_config = self.orig_fetch
        self.model_quant.get_dit_args = self.orig_get_dit
        self.model_quant.get_quant_type = self.orig_get_qtype
        self.model_quant.do_post_load_quant = self.orig_do_post
        self.kernel_wrappers.is_fp8_compile_supported = self.orig_fp8_supported
        self.sdnq_common.check_torch_compile = self.orig_check_compile
        for key, value in self.orig_opts.items():
            self.shared.opts.data[key] = value
        return False


def test_load_comfy_int8_end_to_end():
    """Full pipeline on a comfy_quant fixture: the marked linear becomes an
    SDNQ int8 layer holding the file's exact tensors, the unmarked linear
    loads plain, and the forward pass matches a dequantized reference."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        raw = comfy_fixture(dim)
        write_fixture(raw, fd, path)

        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            transformer, siblings = nt.load(
                local_file=path,
                repo_id='fake/repo',
                spec=spec,
                diffusers_cfg={},
                dtype=torch.float32,
            )

        assert siblings == {}
        in_proj = transformer.in_proj
        assert in_proj.__class__.__name__ == 'SDNQLinear', f'marked layer is {in_proj.__class__.__name__}'
        dq = in_proj.sdnq_dequantizer
        assert dq.weights_dtype == 'int8'
        assert dq.group_size == -1
        assert dq.use_quantized_matmul is False
        assert dq.original_shape == (dim, dim)

        # File tensors adopted bit-exact: int8 codes and fp32 scalar scale.
        assert in_proj.weight.dtype == torch.int8
        assert torch.equal(in_proj.weight.detach().cpu(), raw['model.diffusion_model.in_proj.weight'])
        assert in_proj.scale.dtype == torch.float32
        assert tuple(in_proj.scale.shape) == (1, 1)
        assert in_proj.scale.item() == raw['model.diffusion_model.in_proj.weight_scale'].item()
        assert in_proj.zero_point is None

        # Dequant matches comfy semantics exactly: fp = int8.float() * scale.
        expected = raw['model.diffusion_model.in_proj.weight'].float() * raw['model.diffusion_model.in_proj.weight_scale']
        dequantized = dq(in_proj.weight, in_proj.scale, zero_point=None, svd_up=None, svd_down=None)
        assert torch.equal(dequantized.detach().cpu(), expected)

        # Forward parity against a reference built from the dequantized weight.
        x = torch.randn(2, dim)
        out = in_proj(x)
        ref = torch.nn.functional.linear(x, expected, in_proj.bias.detach().cpu())
        assert torch.allclose(out.detach().cpu(), ref, atol=1e-6)

        # Unmarked layer loads as a plain Linear cast to the target dtype.
        assert transformer.out_proj.__class__ is torch.nn.Linear
        assert transformer.out_proj.weight.dtype == torch.float32
        assert torch.allclose(
            transformer.out_proj.weight.detach().cpu(),
            raw['model.diffusion_model.out_proj.weight'].float(),
        )

        # Marked as SDNQ-quantized so downstream never re-quantizes.
        assert getattr(transformer, 'quantization_config', None) is not None
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_comfy_fp8_end_to_end():
    """float8_e4m3fn variant: same container, fp8 storage. The marked linear
    must keep fp8 codes and dequantize as weight.float() * scale."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        raw = comfy_fixture(dim, fmt='float8_e4m3fn')
        write_fixture(raw, fd, path)

        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            transformer, _ = nt.load(
                local_file=path,
                repo_id='fake/repo',
                spec=spec,
                diffusers_cfg={},
                dtype=torch.float32,
            )

        in_proj = transformer.in_proj
        assert in_proj.__class__.__name__ == 'SDNQLinear'
        assert in_proj.sdnq_dequantizer.weights_dtype == 'float8_e4m3fn'
        assert in_proj.weight.dtype == torch.float8_e4m3fn
        assert in_proj.scale.dtype == torch.float32

        expected = raw['model.diffusion_model.in_proj.weight'].float() * raw['model.diffusion_model.in_proj.weight_scale']
        dequantized = in_proj.sdnq_dequantizer(in_proj.weight, in_proj.scale, zero_point=None, svd_up=None, svd_down=None)
        assert torch.equal(dequantized.detach().cpu(), expected)
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_comfy_fp8_remaps_storage_when_compile_unsupported():
    """When compiled dequant is on but the hardware cannot compile e4m3,
    fp8 weights are adopted through the uint8-backed codec: identical
    decoded values from storage the compiler can always touch."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        raw = comfy_fixture(dim, fmt='float8_e4m3fn')
        write_fixture(raw, fd, path)

        with ComfyTestEnv(dim, fp8_compile_supported=False):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            transformer, _ = nt.load(
                local_file=path,
                repo_id='fake/repo',
                spec=spec,
                diffusers_cfg={},
                dtype=torch.float32,
            )

        in_proj = transformer.in_proj
        assert in_proj.__class__.__name__ == 'SDNQLinear'
        assert in_proj.sdnq_dequantizer.weights_dtype == 'float8_e4m3fn_sdnq'
        assert in_proj.weight.dtype == torch.uint8
        expected = raw['model.diffusion_model.in_proj.weight'].float() * raw['model.diffusion_model.in_proj.weight_scale']
        dequantized = in_proj.sdnq_dequantizer(in_proj.weight, in_proj.scale, zero_point=None, svd_up=None, svd_down=None)
        assert torch.equal(dequantized.detach().cpu(), expected)
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_comfy_metadata_end_to_end():
    """Marker-less container: quant info only in the header
    _quantization_metadata, layer names carrying the file's tensor prefix.
    The load path must re-key through the prefix strip and land on the same
    prequantized build as the marker form."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        raw = comfy_fixture(dim)
        del raw['model.diffusion_model.in_proj.comfy_quant']
        layers = {'model.diffusion_model.in_proj': {'format': 'int8_tensorwise'}}
        write_fixture(raw, fd, path, metadata=quant_metadata(layers))

        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            transformer, _ = nt.load(
                local_file=path,
                repo_id='fake/repo',
                spec=spec,
                diffusers_cfg={},
                dtype=torch.float32,
            )

        in_proj = transformer.in_proj
        assert in_proj.__class__.__name__ == 'SDNQLinear'
        assert in_proj.weight.dtype == torch.int8
        assert torch.equal(in_proj.weight.detach().cpu(), raw['model.diffusion_model.in_proj.weight'])
        assert transformer.out_proj.__class__ is torch.nn.Linear
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_comfy_metadata_bare_names_end_to_end():
    """Marker-less container without a tensor prefix (official Comfy-Org
    exports): metadata names match the bare keys and need no re-keying."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        raw = {key[len('model.diffusion_model.'):]: value for key, value in comfy_fixture(dim).items()}
        del raw['in_proj.comfy_quant']
        raw['in_proj.input_scale'] = torch.tensor(0.125, dtype=torch.float32)
        layers = {'in_proj': {'format': 'int8_tensorwise'}}
        write_fixture(raw, fd, path, metadata=quant_metadata(layers))

        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            transformer, _ = nt.load(
                local_file=path,
                repo_id='fake/repo',
                spec=spec,
                diffusers_cfg={},
                dtype=torch.float32,
            )

        assert transformer.in_proj.__class__.__name__ == 'SDNQLinear'
        assert torch.equal(transformer.in_proj.weight.detach().cpu(), raw['in_proj.weight'])
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_comfy_full_precision_mm_excluded_from_matmul():
    """Layers flagged full_precision_matrix_mult land on the config's
    per-layer matmul exclusion list, which apply_sdnq_options_to_model
    honors when the user enables quantized matmul."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        import json
        raw = comfy_fixture(dim)
        payload = {'format': 'int8_tensorwise', 'full_precision_matrix_mult': True}
        raw['model.diffusion_model.in_proj.comfy_quant'] = torch.tensor(list(json.dumps(payload).encode()), dtype=torch.uint8)
        write_fixture(raw, fd, path)

        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            transformer, _ = nt.load(
                local_file=path,
                repo_id='fake/repo',
                spec=spec,
                diffusers_cfg={},
                dtype=torch.float32,
            )

        assert 'in_proj.weight' in transformer.quantization_config.modules_to_not_use_matmul
        assert transformer.in_proj.sdnq_dequantizer.use_quantized_matmul is False
    finally:
        if os.path.exists(path):
            os.unlink(path)


def nvfp4_fixture(dim: int, with_marker: bool = True) -> tuple[dict, torch.Tensor]:
    """A valid nvfp4 container for one marked linear, plus the reference
    dequantized weight per the format definition: E2M1 decode * unswizzled
    e4m3 block scale * fp32 global scale."""
    import json
    groups = dim // 16
    gen = torch.Generator().manual_seed(7)
    codes = torch.randint(0, 16, (dim, dim), generator=gen, dtype=torch.uint8)
    block_scales = (torch.rand(dim, groups, generator=gen) * 2 + 0.5).to(torch.float8_e4m3fn)
    global_scale = torch.tensor(0.0125, dtype=torch.float32)
    lut = torch.tensor(E2M1_VALUES, dtype=torch.float32)
    reference = lut[codes.long()].reshape(dim, groups, 16) * (block_scales.float() * global_scale).unsqueeze(-1)
    reference = reference.reshape(dim, dim)
    packed = torch.bitwise_or(torch.bitwise_left_shift(codes[:, 0::2], 4), codes[:, 1::2])  # even element in the high nibble
    sd = {
        'model.diffusion_model.in_proj.weight': packed,
        'model.diffusion_model.in_proj.weight_scale': to_blocked_reference(block_scales),
        'model.diffusion_model.in_proj.weight_scale_2': global_scale,
        'model.diffusion_model.in_proj.bias': torch.randn(dim, generator=gen).to(torch.float16),
        'model.diffusion_model.out_proj.weight': torch.randn(dim, dim, generator=gen).to(torch.float16),
        'model.diffusion_model.out_proj.bias': torch.zeros(dim, dtype=torch.float16),
    }
    if with_marker:
        marker = json.dumps({'format': 'nvfp4', 'group_size': 16, 'orig_dtype': 'torch.float16', 'orig_shape': [dim, dim]})
        sd['model.diffusion_model.in_proj.comfy_quant'] = torch.tensor(list(marker.encode()), dtype=torch.uint8)
    return sd, reference


def test_load_comfy_nvfp4_end_to_end():
    """Full pipeline on an nvfp4 container: packed uint8 weights survive with
    swapped nibble order, block scales unswizzle and fold with the global
    scale into fp32 grouped scales, and the dequantized weight matches the
    format reference exactly."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 32
        raw, reference = nvfp4_fixture(dim)
        write_fixture(raw, fd, path)

        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            transformer, _ = nt.load(
                local_file=path,
                repo_id='fake/repo',
                spec=spec,
                diffusers_cfg={},
                dtype=torch.float32,
            )

        in_proj = transformer.in_proj
        assert in_proj.__class__.__name__ == 'SDNQLinear'
        dq = in_proj.sdnq_dequantizer
        assert dq.weights_dtype == 'float4_e2m1fn'
        assert dq.group_size == 16
        assert tuple(dq.quantized_weight_shape) == (dim, dim // 16, 16)
        assert dq.re_quantize_for_matmul is True
        assert in_proj.weight.dtype == torch.uint8
        assert tuple(in_proj.weight.shape) == (dim, dim // 2)
        assert in_proj.scale.dtype == torch.float32
        assert tuple(in_proj.scale.shape) == (dim, dim // 16, 1)

        dequantized = dq(in_proj.weight, in_proj.scale, zero_point=None, svd_up=None, svd_down=None)
        assert torch.equal(dequantized.detach().cpu(), reference)

        x = torch.randn(2, dim)
        out = in_proj(x)
        ref = torch.nn.functional.linear(x, reference, in_proj.bias.detach().cpu())
        assert torch.allclose(out.detach().cpu(), ref, atol=1e-5)

        assert transformer.out_proj.__class__ is torch.nn.Linear
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_comfy_nvfp4_metadata_form():
    """Marker-less nvfp4 (header metadata only, no orig_shape/group_size):
    the format constant and module dimensions fill the gaps."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 32
        raw, reference = nvfp4_fixture(dim, with_marker=False)
        layers = {'model.diffusion_model.in_proj': {'format': 'nvfp4'}}
        write_fixture(raw, fd, path, metadata=quant_metadata(layers))

        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            transformer, _ = nt.load(
                local_file=path,
                repo_id='fake/repo',
                spec=spec,
                diffusers_cfg={},
                dtype=torch.float32,
            )

        in_proj = transformer.in_proj
        assert in_proj.__class__.__name__ == 'SDNQLinear'
        assert in_proj.sdnq_dequantizer.weights_dtype == 'float4_e2m1fn'
        dequantized = in_proj.sdnq_dequantizer(in_proj.weight, in_proj.scale, zero_point=None, svd_up=None, svd_down=None)
        assert torch.equal(dequantized.detach().cpu(), reference)
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_comfy_nvfp4_orig_shape_mismatch_raises():
    """A marker whose orig_shape disagrees with the model config is a wrong
    file for the class; reject so the base-repo fallback engages."""
    import json
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 32
        raw, _ = nvfp4_fixture(dim)
        marker = json.dumps({'format': 'nvfp4', 'group_size': 16, 'orig_shape': [dim, dim * 2]})
        raw['model.diffusion_model.in_proj.comfy_quant'] = torch.tensor(list(marker.encode()), dtype=torch.uint8)
        write_fixture(raw, fd, path)

        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            try:
                nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={}, dtype=torch.float32)
                raise AssertionError('expected OverrideArchMismatch')
            except nt.OverrideArchMismatch as e:
                assert 'orig_shape' in str(e)
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_detect_comfy_mixed_formats_raises():
    """One file mixing int8 and fp8 layers has no single SDNQ mapping; reject."""
    sd = {
        'a.weight': torch.zeros((4, 4), dtype=torch.int8),
        'a.comfy_quant': comfy_marker('int8_tensorwise'),
        'b.weight': torch.zeros((4, 4), dtype=torch.float8_e4m3fn),
        'b.comfy_quant': comfy_marker('float8_e4m3fn'),
    }
    try:
        nt.detect_comfy_quant(sd, 'Test')
        raise AssertionError('expected OverrideArchMismatch')
    except nt.OverrideArchMismatch as e:
        assert 'mixes formats' in str(e)


def test_load_comfy_marker_dtype_mismatch_raises():
    """A marker whose declared format does not match the stored weight dtype
    (mislabeled container) must be rejected, not silently misinterpreted."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        raw = comfy_fixture(dim)
        raw['model.diffusion_model.in_proj.weight'] = torch.randn(dim, dim, dtype=torch.float16)  # marker says int8
        write_fixture(raw, fd, path)

        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            raised = False
            try:
                nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={})
            except nt.OverrideArchMismatch as e:
                raised = True
                assert 'torch.int8' in str(e) and 'torch.float16' in str(e)
            assert raised, 'expected OverrideArchMismatch'
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_comfy_unsupported_format_raises_mismatch():
    """A comfy_quant file in a format sdnext cannot adopt must surface as
    OverrideArchMismatch so load_transformer falls back to the base repo."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        raw = comfy_fixture(dim)
        raw['model.diffusion_model.in_proj.comfy_quant'] = comfy_marker('float8_e4m3fn_scaled')
        write_fixture(raw, fd, path)

        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            raised = False
            try:
                nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={})
            except nt.OverrideArchMismatch as e:
                raised = True
                assert 'float8_e4m3fn_scaled' in str(e)
            assert raised, 'expected OverrideArchMismatch'
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_comfy_marker_for_unknown_module_raises_mismatch():
    """A marker naming a module the target class does not have means the file
    belongs to a different arch; must fall back, not crash."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        raw = comfy_fixture(dim)
        raw['model.diffusion_model.ghost.weight'] = torch.zeros((dim, dim), dtype=torch.int8)
        raw['model.diffusion_model.ghost.weight_scale'] = torch.tensor(1.0)
        raw['model.diffusion_model.ghost.comfy_quant'] = comfy_marker('int8_tensorwise')
        write_fixture(raw, fd, path)

        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            raised = False
            try:
                nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={})
            except nt.OverrideArchMismatch as e:
                raised = True
                assert 'ghost' in str(e)
            assert raised, 'expected OverrideArchMismatch'
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_transformer_syncs_loaded_unet():
    """A full model load that consumes the UNET dropdown override must mark it
    as loaded, so the queued sd_unet onchange callback does not trigger a
    second, redundant full model reload."""
    from modules import sd_unet, shared
    from pipelines import generic_transformer as gt

    fd, path = tempfile.mkstemp(suffix='.safetensors')
    dim = 8
    raw = {
        'model.diffusion_model.in_proj.weight': torch.randn(dim, dim),
        'model.diffusion_model.in_proj.bias': torch.zeros(dim),
        'model.diffusion_model.out_proj.weight': torch.randn(dim, dim),
        'model.diffusion_model.out_proj.bias': torch.zeros(dim),
    }
    write_fixture(raw, fd, path)

    orig_unet_opt = shared.opts.sd_unet
    orig_loaded = sd_unet.loaded_unet
    sd_unet.unet_dict['mock-unet'] = path
    shared.opts.data['sd_unet'] = 'mock-unet'
    sd_unet.loaded_unet = None
    try:
        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            transformer = gt.load_transformer('fake/repo', cls_name=MockMiniTransformer, native_spec=spec)
        assert transformer is not None
        assert sd_unet.loaded_unet == 'mock-unet', f'override consumed but loaded_unet={sd_unet.loaded_unet}'
    finally:
        sd_unet.unet_dict.pop('mock-unet', None)
        shared.opts.data['sd_unet'] = orig_unet_opt
        sd_unet.loaded_unet = orig_loaded
        if os.path.exists(path):
            os.unlink(path)


class MockFusedAttention(torch.nn.Module):
    """Split-attention module tree matching Ideogram4Transformer2DModel's
    naming (to_q/to_k/to_v + to_out ModuleList), fed by fused checkpoints."""

    def __init__(self, dim: int):
        super().__init__()
        self.to_q = torch.nn.Linear(dim, dim, bias=False)
        self.to_k = torch.nn.Linear(dim, dim, bias=False)
        self.to_v = torch.nn.Linear(dim, dim, bias=False)
        self.to_out = torch.nn.ModuleList([torch.nn.Linear(dim, dim, bias=False), torch.nn.Dropout(0.0)])


class MockFusedBlock(torch.nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.attention = MockFusedAttention(dim)


class MockFusedTransformer(torch.nn.Module):
    @classmethod
    def from_config(cls, config: dict, **kwargs) -> 'MockFusedTransformer':
        return cls(dim=config['dim'])

    def __init__(self, dim: int):
        super().__init__()
        self.layers = torch.nn.ModuleList([MockFusedBlock(dim)])
        self.input_proj = torch.nn.Linear(dim, dim)


def fused_spec():
    from pipelines.ideogram import convert_ideogram4_transformer_checkpoint
    return nt.TransformerSpec(
        cls=MockFusedTransformer,
        converter=convert_ideogram4_transformer_checkpoint,
        converter_handles_quant=True,
    )


def fused_fixture(dim: int, quantized: bool = True) -> dict:
    if quantized:
        qkv = torch.randint(-128, 127, (3 * dim, dim), dtype=torch.int8)
    else:
        qkv = torch.randn(3 * dim, dim, dtype=torch.float16)
    raw = {
        'model.diffusion_model.layers.0.attention.qkv.weight': qkv,
        'model.diffusion_model.layers.0.attention.o.weight': torch.randn(dim, dim, dtype=torch.float16),
        'model.diffusion_model.input_proj.weight': torch.randn(dim, dim, dtype=torch.float16),
        'model.diffusion_model.input_proj.bias': torch.zeros(dim, dtype=torch.float16),
    }
    if quantized:
        raw['model.diffusion_model.layers.0.attention.qkv.weight_scale'] = torch.rand(3 * dim, 1, dtype=torch.float32)
        raw['model.diffusion_model.layers.0.attention.qkv.comfy_quant'] = comfy_marker('int8_tensorwise')
    return raw


def test_load_fused_comfy_end_to_end():
    """Quant-aware converter + comfy adoption: fused int8 qkv with row-wise
    scales lands as three SDNQ linears holding the exact row slices."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        raw = fused_fixture(dim, quantized=True)
        write_fixture(raw, fd, path)
        with ComfyTestEnv(dim):
            transformer, _ = nt.load(local_file=path, repo_id='fake/repo', spec=fused_spec(), diffusers_cfg={}, dtype=torch.float32)

        attn = transformer.layers[0].attention
        fused_w = raw['model.diffusion_model.layers.0.attention.qkv.weight']
        fused_s = raw['model.diffusion_model.layers.0.attention.qkv.weight_scale']
        for i, name in enumerate(('to_q', 'to_k', 'to_v')):
            layer = getattr(attn, name)
            assert layer.__class__.__name__ == 'SDNQLinear', f'{name} is {layer.__class__.__name__}'
            assert layer.weight.dtype == torch.int8
            assert torch.equal(layer.weight.detach().cpu(), fused_w[i * dim:(i + 1) * dim])
            assert tuple(layer.scale.shape) == (dim, 1), f'{name} scale shape {layer.scale.shape}'
            assert torch.equal(layer.scale.detach().cpu(), fused_s[i * dim:(i + 1) * dim])
            expected = fused_w[i * dim:(i + 1) * dim].float() * fused_s[i * dim:(i + 1) * dim]
            dequantized = layer.sdnq_dequantizer(layer.weight, layer.scale, zero_point=None, svd_up=None, svd_down=None)
            assert torch.equal(dequantized.detach().cpu(), expected)
        # unmarked o.weight passes through the rename as a plain Linear
        assert attn.to_out[0].__class__ is torch.nn.Linear
        assert torch.allclose(attn.to_out[0].weight.detach().cpu(), raw['model.diffusion_model.layers.0.attention.o.weight'].float())
        assert transformer.input_proj.__class__ is torch.nn.Linear
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_fused_bf16_end_to_end():
    """Converter-first on a plain float fused file: no markers, standard load
    path, split weights land on the class-native linears."""
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        raw = fused_fixture(dim, quantized=False)
        write_fixture(raw, fd, path)
        with ComfyTestEnv(dim):
            transformer, _ = nt.load(local_file=path, repo_id='fake/repo', spec=fused_spec(), diffusers_cfg={}, dtype=torch.float32)

        attn = transformer.layers[0].attention
        fused_w = raw['model.diffusion_model.layers.0.attention.qkv.weight']
        for i, name in enumerate(('to_q', 'to_k', 'to_v')):
            layer = getattr(attn, name)
            assert layer.__class__ is torch.nn.Linear
            assert torch.allclose(layer.weight.detach().cpu(), fused_w[i * dim:(i + 1) * dim].float())
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_fused_convrot_end_to_end():
    """Fused qkv with a convrot marker: the converter duplicates the marker to
    the three split layers, each of which loads with hadamard geometry. Row
    slicing commutes with the in-features rotation, so slices stay lossless."""
    import json as json_mod
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim, group_size = 8, 4
        raw = fused_fixture(dim, quantized=True)
        payload = json_mod.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': group_size, 'per_row': True})
        raw['model.diffusion_model.layers.0.attention.qkv.comfy_quant'] = torch.tensor(list(payload.encode()), dtype=torch.uint8)
        raw['model.diffusion_model.layers.0.attention.qkv.weight_scale'] = torch.rand(3 * dim, 1, dtype=torch.float32)
        write_fixture(raw, fd, path)
        with ComfyTestEnv(dim):
            transformer, _ = nt.load(local_file=path, repo_id='fake/repo', spec=fused_spec(), diffusers_cfg={}, dtype=torch.float32)
        attn = transformer.layers[0].attention
        for name in ('to_q', 'to_k', 'to_v'):
            dq = getattr(attn, name).sdnq_dequantizer
            assert dq.use_hadamard is True, f'{name} lost the convrot flag'
            assert dq.hadamard_group_size == group_size
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_build_component_comfy_skips_converter_when_not_quant_aware():
    """With converter_handles_quant=False (the default for every other arch),
    a comfy state dict must reach the prequantized path without the converter
    ever running; a crashing converter proves it was not invoked."""
    dim = 8
    raw = comfy_fixture(dim)
    sd = {k[len('model.diffusion_model.'):]: v for k, v in raw.items()}
    with ComfyTestEnv(dim):
        component = nt.build_component(
            component_name='transformer',
            state_dict=sd,
            config={'dim': dim},
            cls=MockMiniTransformer,
            converter=crashing_converter,
            acceptable_missing=('rope.',),
            quant_args={},
            quant_type='SDNQConfig',
            dtype=torch.float32,
        )
    assert component.in_proj.__class__.__name__ == 'SDNQLinear'


def test_load_transformer_secondary_slot_syncs_tracker():
    """The secondary slot consumes sd_unet_secondary and syncs its own
    tracker without touching the primary slot."""
    from modules import sd_unet, shared
    from pipelines import generic_transformer as gt

    fd, path = tempfile.mkstemp(suffix='.safetensors')
    dim = 8
    raw = {
        'model.diffusion_model.in_proj.weight': torch.randn(dim, dim),
        'model.diffusion_model.in_proj.bias': torch.zeros(dim),
        'model.diffusion_model.out_proj.weight': torch.randn(dim, dim),
        'model.diffusion_model.out_proj.bias': torch.zeros(dim),
    }
    write_fixture(raw, fd, path)

    orig_primary_opt = shared.opts.sd_unet
    orig_secondary_opt = shared.opts.sd_unet_secondary
    orig_primary = sd_unet.loaded_unet
    orig_secondary = sd_unet.loaded_unet_secondary
    sd_unet.unet_dict['mock-unet-2'] = path
    shared.opts.data['sd_unet'] = 'Default'
    shared.opts.data['sd_unet_secondary'] = 'mock-unet-2'
    sd_unet.loaded_unet = None
    sd_unet.loaded_unet_secondary = None
    try:
        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            transformer = gt.load_transformer('fake/repo', cls_name=MockMiniTransformer, native_spec=spec, override_slot='secondary')
        assert transformer is not None
        assert sd_unet.loaded_unet_secondary == 'mock-unet-2'
        assert sd_unet.loaded_unet is None, 'primary tracker must stay untouched'
    finally:
        sd_unet.unet_dict.pop('mock-unet-2', None)
        shared.opts.data['sd_unet'] = orig_primary_opt
        shared.opts.data['sd_unet_secondary'] = orig_secondary_opt
        sd_unet.loaded_unet = orig_primary
        sd_unet.loaded_unet_secondary = orig_secondary
        if os.path.exists(path):
            os.unlink(path)


def regular_hadamard(n: int) -> torch.Tensor:
    """Regular Hadamard construction the ConvRot format and SDNQ share:
    4x4 base, Kronecker recursion, 1/sqrt(n) normalization. Symmetric and
    involutory, so rotation and inverse are the same matrix."""
    h4 = torch.tensor([[1., 1., 1., -1.], [1., 1., -1., 1.], [1., -1., 1., 1.], [-1., 1., 1., 1.]])
    h = h4
    size = 4
    while size < n:
        h = torch.kron(h, h4)
        size *= 4
    return h / (n ** 0.5)


def comfy_convrot_quantize(weight: torch.Tensor, group_size: int):
    """Reference ConvRot int8 quantizer (mirrors the ComfyUI runtime): rotate
    grouped in-features by the regular Hadamard, then row-wise symmetric int8."""
    h = regular_hadamard(group_size)
    out_f, in_f = weight.shape
    rotated = (weight.reshape(out_f, -1, group_size) @ h.T).reshape(out_f, in_f)
    scale = rotated.abs().amax(dim=-1, keepdim=True) / 127
    q = rotated.div(scale).round().clamp(-128, 127).to(torch.int8)
    return q, scale


def test_load_comfy_convrot_end_to_end():
    """ConvRot parity: a layer quantized with the reference ConvRot math
    loads as an SDNQ hadamard layer whose dequant matches the reference dequant
    and recovers the original weight within int8 quantization error. The plain
    second layer proves per-layer mixing within one file."""
    import json as json_mod
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim, group_size = 8, 4
        original = torch.randn(dim, dim)
        q, scale = comfy_convrot_quantize(original, group_size)
        marker = json_mod.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': group_size, 'per_row': True})
        raw = {
            'model.diffusion_model.in_proj.weight': q,
            'model.diffusion_model.in_proj.weight_scale': scale,
            'model.diffusion_model.in_proj.comfy_quant': torch.tensor(list(marker.encode()), dtype=torch.uint8),
            'model.diffusion_model.in_proj.bias': torch.zeros(dim, dtype=torch.float16),
            'model.diffusion_model.out_proj.weight': torch.randint(-128, 127, (dim, dim), dtype=torch.int8),
            'model.diffusion_model.out_proj.weight_scale': torch.tensor(0.03125, dtype=torch.float32),
            'model.diffusion_model.out_proj.comfy_quant': comfy_marker('int8_tensorwise'),
            'model.diffusion_model.out_proj.bias': torch.zeros(dim, dtype=torch.float16),
        }
        write_fixture(raw, fd, path)

        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            transformer, _ = nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={}, dtype=torch.float32)

        in_proj = transformer.in_proj
        dq = in_proj.sdnq_dequantizer
        assert dq.use_hadamard is True
        assert dq.hadamard_group_size == group_size
        assert tuple(in_proj.scale.shape) == (dim, 1)

        dequantized = dq(in_proj.weight, in_proj.scale, zero_point=None, svd_up=None, svd_down=None).detach().cpu()
        h = regular_hadamard(group_size)
        reference = ((q.float() * scale).reshape(dim, -1, group_size) @ h.T).reshape(dim, dim)
        assert torch.allclose(dequantized, reference, atol=1e-5), 'SDNQ hadamard dequant diverges from the ConvRot reference'
        assert torch.allclose(dequantized, original, atol=0.15), 'dequant does not recover the original weight'

        # plain layer in the same file stays rotation-free
        assert transformer.out_proj.sdnq_dequantizer.use_hadamard is False
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_load_comfy_convrot_nondivisible_falls_back():
    """A convrot group size that does not divide in_features cannot be undone;
    the file must fall back to the base repo."""
    import json as json_mod
    fd, path = tempfile.mkstemp(suffix='.safetensors')
    try:
        dim = 8
        raw = comfy_fixture(dim)
        marker = json_mod.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 16})
        raw['model.diffusion_model.in_proj.comfy_quant'] = torch.tensor(list(marker.encode()), dtype=torch.uint8)
        write_fixture(raw, fd, path)

        with ComfyTestEnv(dim):
            spec = nt.TransformerSpec(cls=MockMiniTransformer)
            raised = False
            try:
                nt.load(local_file=path, repo_id='fake/repo', spec=spec, diffusers_cfg={})
            except nt.OverrideArchMismatch as e:
                raised = True
                assert 'does not divide' in str(e)
            assert raised, 'expected OverrideArchMismatch'
    finally:
        if os.path.exists(path):
            os.unlink(path)


def test_build_component_comfy_preempts_sdnq_fresh_quant():
    """When SDNQ on-load quant settings are active (quant_type=SDNQConfig), a
    comfy_quant file must still take the pre-quantized path: fresh quant of
    int8 data would corrupt it. quant_args={} proves the comfy branch never
    reaches build_component_quantized (which requires a quantization_config)."""
    dim = 8
    raw = comfy_fixture(dim)
    sd = {k[len('model.diffusion_model.'):]: v for k, v in raw.items()}
    with ComfyTestEnv(dim):
        component = nt.build_component(
            component_name='transformer',
            state_dict=sd,
            config={'dim': dim},
            cls=MockMiniTransformer,
            converter=None,
            acceptable_missing=('rope.',),
            quant_args={},
            quant_type='SDNQConfig',
            dtype=torch.float32,
        )
    assert component.in_proj.__class__.__name__ == 'SDNQLinear'
    assert component.in_proj.weight.dtype == torch.int8


# ============================================================
# anima config inference (depth-expanded checkpoints)
# ============================================================

def anima_block_keys(count: int, prefix: str = 'blocks') -> dict:
    sd = {}
    for i in range(count):
        sd[f'{prefix}.{i}.self_attn.q_proj.weight'] = torch.zeros(2, 2)
        sd[f'{prefix}.{i}.mlp.layer2.weight'] = torch.zeros(2, 2)
    return sd


def test_anima_infer_config_counts_cosmos_blocks():
    """40-block file reports num_layers=40; adapter ``blocks.N`` keys and
    top-level tensors do not count."""
    from pipelines import anima
    sd = anima_block_keys(40)
    sd['llm_adapter.blocks.5.self_attn.q_proj.weight'] = torch.zeros(2, 2)
    sd['x_embedder.proj.1.weight'] = torch.zeros(2, 2)
    sd['final_layer.linear.weight'] = torch.zeros(2, 2)
    assert anima.infer_config(sd) == {'num_layers': 40}


def test_anima_infer_config_matches_base_depth():
    from pipelines import anima
    assert anima.infer_config(anima_block_keys(28)) == {'num_layers': 28}


def test_anima_infer_config_diffusers_layout():
    """Files already in diffusers key layout count ``transformer_blocks.N``."""
    from pipelines import anima
    assert anima.infer_config(anima_block_keys(40, prefix='transformer_blocks')) == {'num_layers': 40}


def test_anima_infer_config_comfy_sidecars_do_not_inflate():
    """Per-block quant sidecars share the block index and do not change the count."""
    from pipelines import anima
    sd = anima_block_keys(40)
    for i in range(40):
        sd[f'blocks.{i}.self_attn.q_proj.weight_scale'] = torch.zeros(2, 1)
        sd[f'blocks.{i}.self_attn.q_proj.comfy_quant'] = torch.zeros(8, dtype=torch.uint8)
    assert anima.infer_config(sd) == {'num_layers': 40}


def test_anima_infer_config_no_blocks_returns_empty():
    from pipelines import anima
    assert anima.infer_config({'x_embedder.proj.1.weight': torch.zeros(2, 2)}) == {}
    assert anima.infer_config({}) == {}


def test_anima_spec_registers_infer_config():
    from pipelines import anima
    assert anima.ANIMA_SPEC.infer_config is anima.infer_config


# ============================================================
# Run
# ============================================================

def run_all():
    log.warning('=== drop_companion_keys ===')
    cat = category('companion')
    for fn in [
        test_drop_companion_keys_no_companions_pass_through,
        test_drop_companion_keys_filters_all_in_one_layout,
        test_drop_companion_keys_raises_when_nothing_left,
        test_drop_companion_keys_empty_prefix_tuple_no_op,
        test_drop_then_strip_all_in_one_layout,
    ]:
        run_test(cat, fn)

    log.warning('=== strip_prefix ===')
    cat = category('strip')
    for fn in [
        test_strip_prefix_bare_keys_pass_through,
        test_strip_prefix_dominant_single_variant,
        test_strip_prefix_picks_longest_match_first,
        test_strip_prefix_mixed_prefixes_raises,
        test_strip_prefix_net_variant,
        test_strip_prefix_diffusion_model_variant,
        test_strip_prefix_custom_prefix_set,
    ]:
        run_test(cat, fn)

    log.warning('=== partition_siblings ===')
    cat = category('partition')
    for fn in [
        test_partition_siblings_empty_spec_returns_state_dict_unchanged,
        test_partition_siblings_no_matches_keeps_all_in_transformer,
        test_partition_siblings_single_sibling_split,
        test_partition_siblings_multiple_siblings,
    ]:
        run_test(cat, fn)

    log.warning('=== forbidden_markers ===')
    cat = category('forbidden')
    for fn in [
        test_forbidden_markers_passes_when_absent,
        test_forbidden_markers_raises_when_present,
        test_forbidden_markers_empty_tuple_no_op,
    ]:
        run_test(cat, fn)

    log.warning('=== comfy_quant detection / remap ===')
    cat = category('comfy')
    for fn in [
        test_detect_comfy_no_markers_returns_none,
        test_detect_comfy_valid_markers,
        test_detect_comfy_unsupported_format_raises,
        test_detect_comfy_malformed_marker_raises,
        test_detect_comfy_convrot_accepted,
        test_detect_comfy_convrot_bad_groupsize_raises,
        test_detect_comfy_marker_missing_format_field_raises,
        test_remap_comfy_renames_and_reshapes_scale,
        test_remap_comfy_passes_unmarked_keys_verbatim,
        test_remap_comfy_does_not_mutate_input,
        test_remap_comfy_drops_input_scale,
        test_read_quant_metadata_absent_returns_none,
        test_read_quant_metadata_parses_layers,
        test_read_quant_metadata_malformed_raises,
        test_read_quant_metadata_non_dict_layers_raises,
        test_transcode_quant_metadata_synthesizes_markers,
        test_transcode_quant_metadata_header_wins,
        test_transcode_quant_metadata_no_matches_is_noop,
    ]:
        run_test(cat, fn)

    log.warning('=== nvfp4 codec / unswizzle / load ===')
    cat = category('nvfp4')
    for fn in [
        test_unswizzle_block_scales_roundtrip,
        test_nvfp4_codec_ocp_table,
        test_nvfp4_pack_ocp_grid_roundtrip,
        test_detect_comfy_nvfp4_marker_accepted,
        test_detect_comfy_nvfp4_bad_group_size_raises,
        test_detect_comfy_nvfp4_convrot_raises,
        test_load_comfy_nvfp4_end_to_end,
        test_load_comfy_nvfp4_metadata_form,
        test_load_comfy_nvfp4_orig_shape_mismatch_raises,
    ]:
        run_test(cat, fn)

    log.warning('=== model_probe quant detection ===')
    cat = category('probe')
    for fn in [
        test_probe_detects_header_quant_metadata,
        test_probe_resolves_nvfp4_from_weight_dtype,
        test_probe_fp8_markers_despite_u8_noise,
        test_probe_int8_markers,
        test_probe_plain_file_no_quant,
    ]:
        run_test(cat, fn)

    log.warning('=== noop_converter detection ===')
    cat = category('noop')
    for fn in [
        test_noop_converter_identity_lambda,
        test_noop_converter_real_function,
        test_noop_converter_lambda_with_modification,
    ]:
        run_test(cat, fn)

    log.warning('=== validate_state_dict_load ===')
    cat = category('validate')
    for fn in [
        test_validate_accepts_buffer_only_missing,
        test_validate_rejects_unexpected,
        test_validate_rejects_hard_missing,
        test_validate_empty_passes,
    ]:
        run_test(cat, fn)

    log.warning('=== make_default_spec ===')
    cat = category('default_spec')
    for fn in [
        test_make_default_spec_for_unknown_class,
        test_make_default_spec_picks_up_real_diffusers_converter,
        test_make_default_spec_skips_qwen_image_noop,
    ]:
        run_test(cat, fn)

    log.warning('=== auto_pickup_converter ===')
    cat = category('autopickup')
    for fn in [
        test_auto_pickup_returns_none_for_unknown_class,
        test_auto_pickup_returns_real_diffusers_converter,
        test_auto_pickup_skips_noop_converter_qwen,
    ]:
        run_test(cat, fn)

    log.warning('=== TransformerSpec / SiblingSpec ===')
    cat = category('specs')
    for fn in [
        test_transformer_spec_defaults,
        test_sibling_spec_defaults,
        test_transformer_spec_is_frozen,
    ]:
        run_test(cat, fn)

    log.warning('=== end-to-end load ===')
    cat = category('load')
    for fn in [
        test_load_end_to_end_with_bfl_prefix_no_converter,
        test_load_forwards_kwargs_to_from_config,
        test_load_infer_config_sizes_module_to_file,
        test_load_end_to_end_with_sibling_partition,
        test_load_raises_on_missing_sibling_class,
        test_load_rejects_non_safetensors,
        test_build_component_converter_crash_raises_mismatch,
        test_build_component_shape_mismatch_is_hard_error,
        test_load_converter_crash_raises_mismatch,
    ]:
        run_test(cat, fn)

    log.warning('=== comfy_quant load ===')
    cat = category('comfy_load')
    for fn in [
        test_load_comfy_int8_end_to_end,
        test_load_comfy_fp8_end_to_end,
        test_load_comfy_fp8_remaps_storage_when_compile_unsupported,
        test_load_comfy_metadata_end_to_end,
        test_load_comfy_metadata_bare_names_end_to_end,
        test_load_comfy_full_precision_mm_excluded_from_matmul,
        test_detect_comfy_mixed_formats_raises,
        test_load_comfy_marker_dtype_mismatch_raises,
        test_load_comfy_unsupported_format_raises_mismatch,
        test_load_comfy_marker_for_unknown_module_raises_mismatch,
        test_load_comfy_convrot_end_to_end,
        test_load_comfy_convrot_nondivisible_falls_back,
        test_load_transformer_syncs_loaded_unet,
        test_load_transformer_secondary_slot_syncs_tracker,
        test_build_component_comfy_preempts_sdnq_fresh_quant,
        test_build_component_comfy_skips_converter_when_not_quant_aware,
    ]:
        run_test(cat, fn)

    log.warning('=== ideogram4 converter / fused load ===')
    cat = category('ideogram4')
    for fn in [
        test_ideogram4_converter_splits_fused_qkv_weight,
        test_ideogram4_converter_splits_rowwise_scale,
        test_ideogram4_converter_scalar_scale_copied_not_sliced,
        test_ideogram4_converter_duplicates_marker,
        test_ideogram4_converter_renames_o_with_sidecars,
        test_ideogram4_converter_passthrough,
        test_ideogram4_converter_defensive_bias_split,
        test_ideogram4_converter_does_not_mutate_input,
        test_load_fused_comfy_end_to_end,
        test_load_fused_bf16_end_to_end,
        test_load_fused_convrot_end_to_end,
    ]:
        run_test(cat, fn)

    log.warning('=== anima config inference ===')
    cat = category('anima')
    for fn in [
        test_anima_infer_config_counts_cosmos_blocks,
        test_anima_infer_config_matches_base_depth,
        test_anima_infer_config_diffusers_layout,
        test_anima_infer_config_comfy_sidecars_do_not_inflate,
        test_anima_infer_config_no_blocks_returns_empty,
        test_anima_spec_registers_infer_config,
    ]:
        run_test(cat, fn)

    log.warning('=== Results ===')
    total_passed = 0
    total_failed = 0
    for cat_name, info in results.items():
        ok = info['failed'] == 0
        status = 'PASS' if ok else 'FAIL'
        log.info(f"  {cat_name}: {info['passed']} passed, {info['failed']} failed [{status}]")
        total_passed += info['passed']
        total_failed += info['failed']
    log.warning(f'Total: {total_passed} passed, {total_failed} failed')
    return total_failed == 0


if __name__ == '__main__':
    import time
    t0 = time.time()
    ok = run_all()
    log.warning(f'Total time: {time.time() - t0:.2f}s')
    sys.exit(0 if ok else 1)