feat(model): extend comfy_quant loading to float8_e4m3fn

Parameterize the pre-quantized comfy_quant path by format so fp8
tensorwise checkpoints load alongside int8; both map onto SDNQ's
symmetric dequant for the matching weights dtype. Containers are
mislabeled in the wild, so the stored weight dtype is validated
against the marker, mixed-format files are rejected, and markers
carrying a convrot rotation flag fall back to the base repo since
they may require runtime inverse rotation.

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 ``<name>.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,

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
 

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,