From 68c1ea8e78aa5fda86fc621918c922d7f2cec808 Mon Sep 17 00:00:00 2001 From: CalamitousFelicitousness Date: Fri, 10 Jul 2026 01:54:08 +0100 Subject: [PATCH] feat(model): load comfyui comfy_quant int8 checkpoints via sdnq Detect comfy_quant markers in native single-file transformer loads and adopt the pre-quantized tensors as SDNQ int8 layers instead of rejecting them. ComfyUI int8_tensorwise is a strict subset of SDNQ symmetric int8, so weights and scales are taken bit-exact with no dequantize-requantize round trip; quantized matmul and scale-dtype settings apply as usual. - detect_comfy_quant/remap_comfy_quant helpers plus a prequantized builder in native_transformer; file markers dictate the layer set - unsupported formats and wrong-arch markers fall back to the base repo - offline unit tests and an opt-in real-file krea2 test --- pipelines/native_transformer.py | 252 +++++++++++++++++++++++- test/test-krea2-transformer.py | 61 +++++- test/test-native-transformer.py | 327 ++++++++++++++++++++++++++++++++ 3 files changed, 632 insertions(+), 8 deletions(-) diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index ec312594c..01385519c 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -41,6 +41,7 @@ loading the adapter from the base repo. """ import os +import json import time from dataclasses import dataclass, field from typing import Callable @@ -69,6 +70,8 @@ DEFAULT_IGNORED_PREFIXES: tuple[str, ...] = ( "text_encoders.", "vae.", ) +COMFY_QUANT_MARKER = ".comfy_quant" +COMFY_SUPPORTED_FORMATS: tuple[str, ...] = ("int8_tensorwise",) class OverrideArchMismatch(Exception): @@ -415,6 +418,66 @@ def check_forbidden_markers( ) +def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[set[str], str] | None: + """Detect ComfyUI ``comfy_quant`` pre-quantized layers in a state dict. + + ComfyUI's quantized checkpoints mark each quantized layer with a + ``.comfy_quant`` uint8 tensor whose bytes are a JSON object naming + the storage format, alongside a ``.weight_scale`` tensor. Returns + the set of marked module names and the format string, or ``None`` when no + markers are present. Raises :class:`OverrideArchMismatch` for malformed + markers or unsupported formats so the caller's base-repo fallback engages. + """ + marker_keys = [key for key in state_dict if key.endswith(COMFY_QUANT_MARKER)] + if not marker_keys: + return None + marked: set[str] = set() + formats: set[str] = set() + for key in marker_keys: + name = key[: -len(COMFY_QUANT_MARKER)] + try: + meta = json.loads(state_dict[key].cpu().numpy().tobytes()) + fmt = meta["format"] + except Exception as e: + raise OverrideArchMismatch( + f"Load model: type={type_name} native_transformer comfy_quant marker " + f"for {name!r} is malformed ({type(e).__name__}: {e})" + ) from e + marked.add(name) + formats.add(fmt) + unsupported = sorted(formats - set(COMFY_SUPPORTED_FORMATS)) + if unsupported: + log.error( + f'Load model: type={type_name} quant=comfy format={",".join(unsupported)} not supported ' + f'(supported: {",".join(COMFY_SUPPORTED_FORMATS)})' + ) + raise OverrideArchMismatch( + f"Load model: type={type_name} native_transformer comfy_quant format " + f"{', '.join(unsupported)} not supported" + ) + return marked, next(iter(formats)) + + +def remap_comfy_quant(state_dict: dict, marked_names: set[str]) -> dict: + """Translate comfy_quant tensor naming to SDNQ naming. + + Renames ``.weight_scale`` to ``.scale`` (reshaped to 2-D, + ``[]`` becomes ``[1, 1]``, since SDNQ transposes scales in place) and + drops the ``.comfy_quant`` markers. Returns a new dict; the input + (which may be the cached state dict) is not mutated. + """ + scale_suffix = ".weight_scale" + remapped: dict = {} + for key, value in state_dict.items(): + if key.endswith(COMFY_QUANT_MARKER) and key[: -len(COMFY_QUANT_MARKER)] in marked_names: + continue + if key.endswith(scale_suffix) and key[: -len(scale_suffix)] in marked_names: + remapped[f"{key[: -len(scale_suffix)]}.scale"] = value.reshape(-1, 1) + continue + remapped[key] = value + return remapped + + def partition_siblings( state_dict: dict, siblings: dict[str, SiblingSpec], @@ -552,6 +615,159 @@ def build_component_quantized( return component +def build_component_prequantized( + *, + component_name: str, + state_dict: dict, + config: dict, + cls: type, + marked_names: set[str], + dtype, + acceptable_missing: tuple[str, ...], + zero_init_missing: tuple[str, ...] = (), + **kwargs, +) -> object: + """Build a component from a comfy_quant pre-quantized state dict, mapping + the marked layers onto SDNQ int8 layers without dequantizing. + + ComfyUI ``int8_tensorwise`` is a strict subset of SDNQ symmetric int8: + same storage layout (unpacked int8 ``[out, in]``), same dequant math + (``weight * scale``, no zero point), so the file's tensors are adopted + bit-exact. The model is built under ``init_empty_weights`` and each marked + Linear is swapped for an SDNQ wrapper with per-tensor dequant geometry + (``group_size=-1``, scalar scale); the file dictates which layers are + quantized, independent of the user's quantization settings. Weights load + through ``SDNQQuantizer``'s pre-quantized path, which preserves the int8 + weight dtype and the fp32 scale (unlike :func:`build_component_quantized`, + floating-point SDNQ params are deliberately not cast to the target dtype). + + Layers are assembled in canonical dequant layout first; + ``apply_sdnq_options_to_model`` then enables quantized matmul per the + user's settings (transposing eligible layers), matching the order used by + ``modules.sdnq.loader.load_sdnq_model``. + + Caller is responsible for prefix stripping; the state dict arrives with + its comfy marker keys intact and is remapped here. + """ + import rich.progress as rp + from accelerate import init_empty_weights + from accelerate.utils import set_module_tensor_to_device + from diffusers.utils import get_module_from_name + from modules.sdnq.quantizer import SDNQConfig, SDNQQuantizer + from modules.sdnq.dequantizer import SDNQDequantizer + from modules.sdnq.layers import get_sdnq_wrapper_class + from modules.sdnq.forward import get_forward_func + from modules.sdnq.loader import apply_sdnq_options_to_model + + target_dtype = dtype if dtype is not None else devices.dtype + sd = remap_comfy_quant(state_dict, marked_names) + + quantization_config = SDNQConfig( + weights_dtype="int8", + quantized_matmul_dtype="int8", + group_size=-1, + use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul, + dequantize_fp32=shared.opts.sdnq_dequantize_fp32, + add_skip_keys=False, + modules_to_not_convert=[], + ) + quantizer = SDNQQuantizer(quantization_config, pre_quantized=True) + quantizer.torch_dtype = target_dtype + + with init_empty_weights(include_buffers=False): + component = cls.from_config(config, **kwargs) + + dequant_forward = get_forward_func("Linear", "int8", False) + for name in sorted(marked_names): + try: + parent, child = get_module_from_name(component, name) + linear = getattr(parent, child) + except (AttributeError, ValueError) as e: + raise OverrideArchMismatch( + f"Load model: transformer=native {component_name} comfy_quant marked " + f"module {name!r} not found in {cls.__name__}" + ) from e + if not isinstance(linear, torch.nn.Linear): + raise OverrideArchMismatch( + f"Load model: transformer=native {component_name} comfy_quant marked " + f"module {name!r} is {linear.__class__.__name__}, expected Linear" + ) + linear.sdnq_dequantizer = SDNQDequantizer( + result_dtype=target_dtype, + result_shape=None, + original_shape=torch.Size((linear.out_features, linear.in_features)), + original_stride=(linear.in_features, 1), + quantized_weight_shape=torch.Size((linear.out_features, linear.in_features)), + weights_dtype="int8", + quantized_matmul_dtype="int8", + hadamard_group_size=256, + group_size=-1, + svd_rank=32, + svd_steps=8, + use_quantized_matmul=False, + re_quantize_for_matmul=False, + use_stochastic_rounding=False, + use_hadamard=False, + layer_class_name="Linear", + ) + wrapped = get_sdnq_wrapper_class(linear, dequant_forward) + wrapped.scale = torch.nn.Parameter(torch.empty((1, 1), dtype=torch.float32, device="meta"), requires_grad=False) + wrapped.zero_point = None + wrapped.svd_up = None + wrapped.svd_down = None + setattr(parent, child, wrapped) + + target_device = ( + devices.cpu if shared.opts.diffusers_offload_mode != "none" + else devices.device + ) + + expected_keys = set(component.state_dict().keys()) + loaded_keys: set[str] = set() + unexpected: list[str] = [] + total = len(sd) + + pbar = rp.Progress( + rp.TextColumn(f'[cyan]Load {component_name}:'), + rp.BarColumn(), + rp.MofNCompleteColumn(), + rp.TaskProgressColumn(), + rp.TimeRemainingColumn(), + rp.TimeElapsedColumn(), + rp.TextColumn('[cyan]{task.description}'), + console=console, + ) + with pbar: + task = pbar.add_task(total=total, description=cls.__name__) + for name, value in sd.items(): + if name in expected_keys: + if quantizer.check_if_quantized_param(component, value, name): + quantizer.create_quantized_param(component, value, name, target_device, dtype=target_dtype) + else: + if torch.is_floating_point(value): + value = value.to(target_dtype) + set_module_tensor_to_device(component, name, target_device, value=value, dtype=target_dtype) + loaded_keys.add(name) + else: + unexpected.append(name) + pbar.update(task, advance=1) + + missing = sorted(expected_keys - loaded_keys) + missing = materialize_zero_init( + component, missing, zero_init_missing, device=target_device, dtype=target_dtype + ) + validate_state_dict_load(component_name, missing, unexpected, acceptable_missing) + + component = quantizer._process_model_after_weight_loading(component) # pylint: disable=protected-access + component = apply_sdnq_options_to_model( + component, + dtype=target_dtype, + dequantize_fp32=shared.opts.sdnq_dequantize_fp32, + use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul, + ) + return component + + def build_component( *, component_name: str, @@ -571,11 +787,16 @@ def build_component( """Convert (if needed), instantiate, load weights, dtype-cast, quantize, and offload-place a single component. Raises on any hard failure. - For the transformer component under SDNQ, the per-tensor pre-mode path - in :func:`build_component_quantized` is used so quantization is applied - in flight (one layer's worth of bf16 in memory at a time). All other - cases (siblings, non-quantized loads, NVIDIAModelOptConfig, layerwise - quant) go through the standard load_state_dict + post-quantize path. + Transformer state dicts carrying ComfyUI ``comfy_quant`` markers are + dispatched to :func:`build_component_prequantized`, which adopts the + file's int8 tensors as SDNQ layers regardless of quantization settings + (the converter and ``quant_args`` are bypassed: the file is already + quantized). For the transformer component under SDNQ, the per-tensor + pre-mode path in :func:`build_component_quantized` is used so + quantization is applied in flight (one layer's worth of bf16 in memory + at a time). All other cases (siblings, non-quantized loads, + NVIDIAModelOptConfig, layerwise quant) go through the standard + load_state_dict + post-quantize path. ``dtype`` overrides ``devices.dtype`` when supplied; otherwise the global default is used. ``modules_to_not_convert`` and ``modules_dtype_dict`` @@ -584,6 +805,27 @@ def build_component( reach ``cls.from_config`` for both construction paths. """ try: + comfy_quant = detect_comfy_quant(state_dict, cls.__name__) if component_name == "transformer" else None + if comfy_quant is not None: + marked_names, comfy_format = comfy_quant + log.info( + f'Load model: transformer=native {component_name} quant=comfy ' + f'format={comfy_format} layers={len(marked_names)} keys={len(state_dict)} cls={cls.__name__}' + ) + component = build_component_prequantized( + component_name=component_name, + state_dict=state_dict, + config=config, + cls=cls, + marked_names=marked_names, + dtype=dtype, + acceptable_missing=acceptable_missing, + zero_init_missing=zero_init_missing, + **kwargs, + ) + devices.torch_gc() + return component + if converter is not None: log.debug(f'Load model: transformer=native {component_name} converter={converter.__name__} keys={len(state_dict)}') try: diff --git a/test/test-krea2-transformer.py b/test/test-krea2-transformer.py index b46341f71..11cb7b5f6 100644 --- a/test/test-krea2-transformer.py +++ b/test/test-krea2-transformer.py @@ -8,8 +8,12 @@ - Zero-init regression: a checkpoint that omits the dormant last.up/last.down residual branch, after materialize_zero_init, produces the same output as the base whose up is zeroed. Port only, so it runs without the reference. +- comfy_quant real-file (opt-in): loads an actual ComfyUI int8_tensorwise Krea2 single file + through the native loader and verifies SDNQ adoption plus a finite tiny forward. Enabled by + setting $KREA2_COMFY_FILE to the .safetensors path; needs network for the base repo config + ($KREA2_COMFY_REPO, default CalamitousFelicitousness/Krea-2-Base-Diffusers). -No server, no checkpoint. +No server, no checkpoint (except the opt-in comfy_quant test). """ import importlib.util @@ -89,8 +93,8 @@ def run_parity(mmdit, port): print("PARITY OK") -def load_materialize_zero_init(): - """Import the real loader helper. native_transformer pulls in modules.shared, which needs +def bootstrap_repo(): + """Make repo modules importable. native_transformer pulls in modules.shared, which needs cmd_args parsed first, so bootstrap it the same way the native-transformer suite does.""" repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) if repo not in sys.path: @@ -106,6 +110,10 @@ def load_materialize_zero_init(): sys.argv = orig_argv installer.add_args(modules.cmd_args.parser) modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([]) + + +def load_materialize_zero_init(): + bootstrap_repo() from pipelines.native_transformer import materialize_zero_init return materialize_zero_init @@ -159,12 +167,59 @@ def run_zero_init_regression(port): print("ZERO-INIT OK") +def run_comfy_quant_real_file(): + """Opt-in end-to-end check against a real ComfyUI int8_tensorwise Krea2 file: the native + loader must adopt every marked linear as an SDNQ int8 layer and produce a finite output on + a tiny forward. $KREA2_COMFY_LAYERS overrides the expected layer count (default 224).""" + path = os.environ.get("KREA2_COMFY_FILE") + if not path: + print("COMFY REAL-FILE SKIPPED (set KREA2_COMFY_FILE to enable)") + return + assert os.path.exists(path), f"KREA2_COMFY_FILE not found: {path}" + expected_layers = int(os.environ.get("KREA2_COMFY_LAYERS", "224")) + repo_id = os.environ.get("KREA2_COMFY_REPO", "CalamitousFelicitousness/Krea-2-Base-Diffusers") + + bootstrap_repo() + from pipelines import native_transformer as nt + from pipelines.krea2 import KREA2_SPEC + + transformer, siblings = nt.load(local_file=path, repo_id=repo_id, spec=KREA2_SPEC, diffusers_cfg={}) + assert siblings == {} + + sdnq_layers = [m for m in transformer.modules() if m.__class__.__name__ == "SDNQLinear"] + print(f"comfy_quant real file: {len(sdnq_layers)} SDNQ layers") + assert len(sdnq_layers) == expected_layers, f"expected {expected_layers} SDNQ layers, got {len(sdnq_layers)}" + assert all(m.weight.dtype == torch.int8 for m in sdnq_layers), "all adopted weights must stay int8" + assert transformer.blocks[0].attn.wq.__class__.__name__ == "SDNQLinear" + assert getattr(transformer, "quantization_config", None) is not None + + cfg = transformer.config + param = next(p for p in transformer.parameters() if p.is_floating_point()) + device, dtype = param.device, param.dtype + batch, txtlen, imglen = 1, 3, 4 + seq = txtlen + imglen + gen = torch.Generator().manual_seed(1) + img = torch.randn(batch, imglen, cfg.channels * cfg.patch ** 2, generator=gen).to(device=device, dtype=dtype) + context = torch.randn(batch, txtlen, cfg.txtlayers, cfg.txtdim, generator=gen).to(device=device, dtype=dtype) + timestep = torch.rand(batch, generator=gen).to(device=device, dtype=dtype) + pos = torch.randint(0, 16, (batch, seq, 3), generator=gen).float().to(device=device) + mask = torch.ones(batch, seq, dtype=torch.bool, device=device) + with torch.no_grad(): + out = transformer( + hidden_states=img, encoder_hidden_states=context, timestep=timestep, + position_ids=pos, attention_mask=mask, return_dict=False, + )[0] + assert torch.isfinite(out).all(), "forward produced non-finite values" + print("COMFY REAL-FILE OK") + + def main(): port = load_port() # Parity runs first in pristine torch state; the regression imports modules afterwards. mmdit = load_reference() run_parity(mmdit, port) run_zero_init_regression(port) + run_comfy_quant_real_file() if __name__ == "__main__": diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index 3bcaaa432..11e907d4b 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -8,6 +8,8 @@ Covers the pure helpers that own per-arch knob handling: - ``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 ComfyUI 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 @@ -294,6 +296,111 @@ def test_forbidden_markers_empty_tuple_no_op(): 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 marked == {'blocks.0.attn.wq', 'blocks.0.mlp.up'} + 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_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)' + + # ============================================================ # is_noop_converter # ============================================================ @@ -831,6 +938,202 @@ def test_load_converter_crash_raises_mismatch(): os.unlink(path) +# ============================================================ +# 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) -> dict: + """comfy_quant int8_tensorwise export shape: in_proj carries int8 weight + + scalar fp32 scale + marker; out_proj and biases stay plain f16.""" + return { + 'model.diffusion_model.in_proj.weight': torch.randint(-128, 127, (dim, dim), dtype=torch.int8), + 'model.diffusion_model.in_proj.weight_scale': torch.tensor(0.03125, dtype=torch.float32), + 'model.diffusion_model.in_proj.comfy_quant': comfy_marker('int8_tensorwise'), + '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.""" + + def __init__(self, dim: int): + self.dim = dim + + def __enter__(self): + from modules import model_quant, shared + self.shared = shared + 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.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 + self.orig_opts = { + 'sdnq_use_quantized_matmul': shared.opts.sdnq_use_quantized_matmul, + 'sdnq_dequantize_fp32': shared.opts.sdnq_dequantize_fp32, + 'diffusers_offload_mode': shared.opts.diffusers_offload_mode, + } + shared.opts.data['sdnq_use_quantized_matmul'] = False + 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 + 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_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_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 + + # ============================================================ # Run # ============================================================ @@ -879,6 +1182,20 @@ def run_all(): ]: 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_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, + ]: + run_test(cat, fn) + log.warning('=== noop_converter detection ===') cat = category('noop') for fn in [ @@ -939,6 +1256,16 @@ def run_all(): ]: 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_unsupported_format_raises_mismatch, + test_load_comfy_marker_for_unknown_module_raises_mismatch, + test_build_component_comfy_preempts_sdnq_fresh_quant, + ]: + run_test(cat, fn) + log.warning('=== Results ===') total_passed = 0 total_failed = 0