diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index 01385519c..ab3a0e42b 100644 --- a/pipelines/native_transformer.py +++ b/pipelines/native_transformer.py @@ -71,7 +71,10 @@ DEFAULT_IGNORED_PREFIXES: tuple[str, ...] = ( "vae.", ) COMFY_QUANT_MARKER = ".comfy_quant" -COMFY_SUPPORTED_FORMATS: tuple[str, ...] = ("int8_tensorwise",) +COMFY_QUANT_FORMATS: dict[str, str] = { # comfy_quant format string -> SDNQ weights_dtype + "int8_tensorwise": "int8", + "float8_e4m3fn": "float8_e4m3fn", +} class OverrideArchMismatch(Exception): @@ -426,7 +429,8 @@ def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[set[str], str] 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. + markers, unsupported formats, or a file mixing formats across layers, 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: @@ -443,18 +447,31 @@ def detect_comfy_quant(state_dict: dict, type_name: str) -> tuple[set[str], str] f"Load model: type={type_name} native_transformer comfy_quant marker " f"for {name!r} is malformed ({type(e).__name__}: {e})" ) from e + if meta.get("convrot"): + # ConvRot-flagged layers may require an inverse rotation at + # runtime; adopting them as plain tensors would corrupt outputs. + log.error(f'Load model: type={type_name} quant=comfy format={fmt} convrot=true not supported') + raise OverrideArchMismatch( + f"Load model: type={type_name} native_transformer comfy_quant layer " + f"{name!r} carries a convrot rotation flag, not supported" + ) marked.add(name) formats.add(fmt) - unsupported = sorted(formats - set(COMFY_SUPPORTED_FORMATS)) + unsupported = sorted(formats - set(COMFY_QUANT_FORMATS)) if unsupported: log.error( f'Load model: type={type_name} quant=comfy format={",".join(unsupported)} not supported ' - f'(supported: {",".join(COMFY_SUPPORTED_FORMATS)})' + f'(supported: {",".join(COMFY_QUANT_FORMATS)})' ) raise OverrideArchMismatch( f"Load model: type={type_name} native_transformer comfy_quant format " f"{', '.join(unsupported)} not supported" ) + if len(formats) > 1: + raise OverrideArchMismatch( + f"Load model: type={type_name} native_transformer comfy_quant mixes formats " + f"across layers ({', '.join(sorted(formats))})" + ) return marked, next(iter(formats)) @@ -622,24 +639,27 @@ def build_component_prequantized( config: dict, cls: type, marked_names: set[str], + comfy_format: 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. + the marked layers onto SDNQ quantized 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). + The supported comfy formats (``int8_tensorwise``, ``float8_e4m3fn``) are + strict subsets of SDNQ's symmetric quantization for the corresponding + weights dtype: same storage layout (unpacked 8-bit ``[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 quantized 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 @@ -653,18 +673,33 @@ def build_component_prequantized( 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.common import dtype_dict 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 + weights_dtype = COMFY_QUANT_FORMATS[comfy_format] + matmul_dtype = "int8" if dtype_dict[weights_dtype]["is_integer"] else "float8_e4m3fn" + storage_dtype = dtype_dict[weights_dtype]["storage_dtype"] target_dtype = dtype if dtype is not None else devices.dtype sd = remap_comfy_quant(state_dict, marked_names) + # Civitai relabels these containers freely; trust the marker only as far + # as the stored tensors actually match it. + for name in marked_names: + weight = sd.get(f"{name}.weight") + if weight is None or weight.dtype != storage_dtype: + found = weight.dtype if weight is not None else "missing" + raise OverrideArchMismatch( + f"Load model: transformer=native {component_name} comfy_quant format " + f"{comfy_format} expects {storage_dtype} weights but {name!r} has {found}" + ) + quantization_config = SDNQConfig( - weights_dtype="int8", - quantized_matmul_dtype="int8", + weights_dtype=weights_dtype, + quantized_matmul_dtype=matmul_dtype, group_size=-1, use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul, dequantize_fp32=shared.opts.sdnq_dequantize_fp32, @@ -677,7 +712,7 @@ def build_component_prequantized( with init_empty_weights(include_buffers=False): component = cls.from_config(config, **kwargs) - dequant_forward = get_forward_func("Linear", "int8", False) + dequant_forward = get_forward_func("Linear", matmul_dtype, False) for name in sorted(marked_names): try: parent, child = get_module_from_name(component, name) @@ -698,8 +733,8 @@ def build_component_prequantized( 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", + weights_dtype=weights_dtype, + quantized_matmul_dtype=matmul_dtype, hadamard_group_size=256, group_size=-1, svd_rank=32, @@ -818,6 +853,7 @@ def build_component( config=config, cls=cls, marked_names=marked_names, + comfy_format=comfy_format, dtype=dtype, acceptable_missing=acceptable_missing, zero_init_missing=zero_init_missing, diff --git a/test/test-krea2-transformer.py b/test/test-krea2-transformer.py index 11cb7b5f6..d4560f9a0 100644 --- a/test/test-krea2-transformer.py +++ b/test/test-krea2-transformer.py @@ -187,9 +187,11 @@ def run_comfy_quant_real_file(): 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") + storage_dtypes = {m.weight.dtype for m in sdnq_layers} + print(f"comfy_quant real file: {len(sdnq_layers)} SDNQ layers, storage {storage_dtypes}") 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 storage_dtypes <= {torch.int8, torch.float8_e4m3fn}, f"unexpected storage dtypes: {storage_dtypes}" + assert len(storage_dtypes) == 1, "adopted weights must share one storage dtype" assert transformer.blocks[0].attn.wq.__class__.__name__ == "SDNQLinear" assert getattr(transformer, "quantization_config", None) is not None diff --git a/test/test-native-transformer.py b/test/test-native-transformer.py index 11e907d4b..0d3ab10af 100644 --- a/test/test-native-transformer.py +++ b/test/test-native-transformer.py @@ -351,6 +351,22 @@ def test_detect_comfy_malformed_marker_raises(): assert 'malformed' in str(e) +def test_detect_comfy_convrot_flag_raises(): + """Markers carrying convrot=true may need runtime inverse rotation; adopting + them as plain tensors would corrupt outputs, so they must be rejected.""" + import json + payload = json.dumps({'format': 'int8_tensorwise', 'convrot': True, 'convrot_groupsize': 256}) + 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), + } + try: + nt.detect_comfy_quant(sd, 'Test') + raise AssertionError('expected OverrideArchMismatch') + except nt.OverrideArchMismatch as e: + assert 'convrot' in str(e) + + def test_detect_comfy_marker_missing_format_field_raises(): import json sd = { @@ -945,13 +961,17 @@ def test_load_converter_crash_raises_mismatch(): # 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.""" +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': torch.randint(-128, 127, (dim, dim), dtype=torch.int8), + '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('int8_tensorwise'), + '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), @@ -1060,6 +1080,78 @@ def test_load_comfy_int8_end_to_end(): 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_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.""" @@ -1189,6 +1281,7 @@ def run_all(): test_detect_comfy_valid_markers, test_detect_comfy_unsupported_format_raises, test_detect_comfy_malformed_marker_raises, + test_detect_comfy_convrot_flag_raises, test_detect_comfy_marker_missing_format_field_raises, test_remap_comfy_renames_and_reshapes_scale, test_remap_comfy_passes_unmarked_keys_verbatim, @@ -1260,6 +1353,9 @@ def run_all(): cat = category('comfy_load') for fn in [ test_load_comfy_int8_end_to_end, + test_load_comfy_fp8_end_to_end, + 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_build_component_comfy_preempts_sdnq_fresh_quant,