Support quantized diffusion model inputs

tests/test_adaptive_extraction_pipeline.py

@@ -122,6 +122,48 @@ def test_paired_stream_is_bounded_ordered_and_releases_once():
     assert all(stream.closed for stream in handler_a.streams + handler_b.streams)
 
 
+def test_paired_stream_decodes_quantized_logical_keys_per_work_unit():
+    streaming = _load_module("nodes.extraction_stream")
+    handler_a = _FakeHandler({
+        "layer1.code": torch.tensor(2),
+        "layer1.scale": torch.tensor(3),
+        "layer2.code": torch.tensor(4),
+        "layer2.scale": torch.tensor(5),
+    })
+    handler_b = _FakeHandler({"layer1": torch.tensor(1), "layer2": torch.tensor(2)})
+
+    class _Quantization:
+        quantized_keys = {"layer1", "layer2"}
+
+        def required_keys(self, key):
+            return (f"{key}.code", f"{key}.scale")
+
+        def decode(self, key, tensors, compute_dtype, device="cpu"):
+            assert set(tensors) == {f"{key}.code", f"{key}.scale"}
+            return tensors[f"{key}.code"] * tensors[f"{key}.scale"]
+
+    quantization = _Quantization()
+    observed = list(streaming.paired_async_tensors(
+        handler_a,
+        handler_b,
+        [("layer1", "layer1", "layer1"), ("layer2", "layer2", "layer2")],
+        pin_memory=False,
+        quantization_a=quantization,
+    ))
+
+    assert [(key, tensor_a.item(), tensor_b.item()) for key, tensor_a, tensor_b in observed] == [
+        ("layer1", 6, 1),
+        ("layer2", 20, 2),
+    ]
+    assert handler_a.calls[0][0] == (
+        "layer1.code", "layer1.scale", "layer2.code", "layer2.scale"
+    )
+    assert Counter(handler_a.marked) == Counter({
+        "layer1.code": 1, "layer1.scale": 1, "layer2.code": 1, "layer2.scale": 1,
+    })
+    assert all(stream.closed for stream in handler_a.streams + handler_b.streams)
+
+
 def test_cuda_oom_compute_retries_on_cpu(monkeypatch):
     streaming = _load_module("nodes.extraction_stream")
     monkeypatch.setattr(streaming, "release_cuda_oom", lambda: None)

tests/test_lora_model_analysis.py

@@ -1,6 +1,7 @@
 import csv
 import importlib
 from io import StringIO
+import json
 import logging
 from pathlib import Path
 import sys
@@ -67,6 +68,21 @@ def test_reports_actual_applied_delta_without_saving(monkeypatch, tmp_path, anal
     assert set(tmp_path.iterdir()) == set(paths.values())
 
 
+def test_quantized_base_is_dequantized_and_sidecars_are_not_analyzed(monkeypatch, tmp_path, analysis):
+    prefix = "model.diffusion_model.layer"
+    quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8)
+    setup(monkeypatch, tmp_path, analysis, pair(), {
+        f"{prefix}.weight": torch.ones((2, 2), dtype=torch.int8),
+        f"{prefix}.weight_scale": torch.tensor(2.0),
+        f"{prefix}.comfy_quant": quant,
+    })
+    comparison, _, metrics_csv, _ = analysis.analyze_lora_on_model("adapter", "base", params())
+    rows = list(csv.DictReader(StringIO(metrics_csv)))
+    assert [row["key"] for row in rows] == [f"{prefix}.weight"]
+    assert float(rows[0]["mae"]) == pytest.approx(2.25)
+    assert "weight_scale" not in comparison
+
+
 @pytest.mark.parametrize("suffixes", [(".lora_down.weight", ".lora_up.weight"), (".lora_A.default.weight", ".lora_B.default.weight"), (".lora_A", ".lora_B")])
 def test_pair_formats_and_flattened_reference_mapping(monkeypatch, tmp_path, analysis, suffixes):
     setup(monkeypatch, tmp_path, analysis, {

tests/test_lora_quantization_guard.py

@@ -436,3 +436,152 @@ def test_merge_to_model_preserves_low_bit_base_and_skips_low_bit_adapter_layer(
     assert "diffusion_model.guarded.lora_down.weight" in report
     assert "torch.int8" in report
     assert "model.diffusion_model.base_quant.weight" in report
+
+
+def test_merge_to_model_dequantizes_base_and_omits_sidecars(monkeypatch, tmp_path, modules):
+    resize = modules["resize"]
+    _patch_output(monkeypatch, resize, tmp_path)
+    base = tmp_path / "quantized_base.safetensors"
+    adapter = tmp_path / "quantized_adapter.safetensors"
+    quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8)
+    prefix = "model.diffusion_model.layer"
+    save_file({
+        f"{prefix}.weight": torch.ones((2, 2), dtype=torch.int8),
+        f"{prefix}.weight_scale": torch.tensor(2.0),
+        f"{prefix}.comfy_quant": quant,
+    }, str(base))
+    save_file({"diffusion_model.layer.diff": torch.ones((2, 2))}, str(adapter))
+
+    output = resize.merge_loras_to_model(
+        [str(adapter)], [1.0], str(base), "cpu", torch.float16,
+        "dequantized_merge", verbose=False,
+    )
+    result = load_file(output)
+    assert list(result) == [f"{prefix}.weight"]
+    torch.testing.assert_close(result[f"{prefix}.weight"].float(), torch.full((2, 2), 3.0))
+
+
+def test_diffusion_decoder_matches_core_quantized_tensor(tmp_path, modules):
+    from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class
+
+    path = tmp_path / "core_int8.safetensors"
+    key = "layer.weight"
+    qdata = torch.tensor([[1, -2], [3, 4]], dtype=torch.int8)
+    scale = torch.tensor(0.5)
+    save_file({
+        key: qdata,
+        "layer.weight_scale": scale,
+        "layer.comfy_quant": torch.tensor(
+            list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8
+        ),
+    }, str(path))
+    handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
+    try:
+        decoder = modules["guard"].DiffusionQuantization(handler, str(path), "Core comparison")
+        keys = decoder.required_keys(key)
+        assert set(keys) == {key, "layer.weight_scale", "layer.comfy_quant"}
+        assert decoder.data_keys == (key,)
+        stream = handler.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False)
+        tensors = {}
+        try:
+            for expected in keys:
+                batch = next(stream)
+                assert len(batch) == 1 and batch[0][0] == expected
+                tensors[expected] = batch[0][1]
+            actual = decoder.decode(key, tensors, torch.float32)
+        finally:
+            for source_key in tensors:
+                handler.mark_processed(source_key)
+            stream.close()
+
+        layout_name = QUANT_ALGOS["int8_tensorwise"]["comfy_tensor_layout"]
+        layout = get_layout_class(layout_name)
+        expected = QuantizedTensor(
+            qdata, layout_name,
+            layout.Params(scale=scale, orig_dtype=torch.float32, orig_shape=(2, 2)),
+        ).dequantize()
+        torch.testing.assert_close(actual, expected)
+        torch.testing.assert_close(actual, qdata.float() * scale)
+    finally:
+        handler.__exit__(None, None, None)
+
+
+def test_padded_shape_comes_from_core_meta_architecture(monkeypatch, modules):
+    import comfy.model_detection
+
+    class HeaderOnly:
+        def keys(self):
+            return ["model.layer.weight", "model.layer.weight_scale", "model.layer.weight_scale_2"]
+
+        def metadata(self):
+            return {"_quantization_metadata": json.dumps({"layers": {"model.layer": {"format": "nvfp4"}}})}
+
+        def get_shape(self, key):
+            return {"model.layer.weight": (16, 8), "model.layer.weight_scale": (16, 1), "model.layer.weight_scale_2": ()}[key]
+
+        def get_dtype(self, key):
+            return torch.uint8 if key == "model.layer.weight" else torch.float32
+
+    class Config:
+        def get_model(self, state, prefix, device):
+            assert device == torch.device("meta")
+            assert all(tensor.device.type == "meta" for tensor in state.values())
+            model = torch.nn.Module()
+            model.diffusion_model = torch.nn.Module()
+            model.diffusion_model.layer = torch.nn.Module()
+            model.diffusion_model.layer._orig_shape = (3, 5)
+            return model
+
+    monkeypatch.setattr(comfy.model_detection, "model_config_from_unet", lambda *args, **kwargs: Config())
+    decoder = modules["guard"].DiffusionQuantization(HeaderOnly(), "header-only", "Shape test")
+    assert decoder.logical_shape("model.layer.weight") == (3, 5)
+
+
+def test_nvfp4_stream_decode_matches_core_full_dequantization(monkeypatch, tmp_path, modules):
+    import comfy.model_detection
+    from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor, get_layout_class
+
+    layout_name = QUANT_ALGOS["nvfp4"]["comfy_tensor_layout"]
+    layout = get_layout_class(layout_name)
+    original = torch.arange(15, dtype=torch.float16).reshape(3, 5) / 10
+    qdata, params = layout.quantize(original)
+    path = tmp_path / "nvfp4.safetensors"
+    metadata = {"_quantization_metadata": json.dumps({"layers": {"model.layer": {"format": "nvfp4"}}})}
+    save_file({
+        "model.layer.weight": qdata,
+        "model.layer.weight_scale": params.block_scale,
+        "model.layer.weight_scale_2": params.scale,
+    }, str(path), metadata=metadata)
+
+    class Config:
+        def get_model(self, state, prefix, device):
+            model = torch.nn.Module()
+            model.diffusion_model = torch.nn.Module()
+            model.diffusion_model.layer = torch.nn.Module()
+            model.diffusion_model.layer._orig_shape = tuple(original.shape)
+            return model
+
+    monkeypatch.setattr(comfy.model_detection, "model_config_from_unet", lambda *args, **kwargs: Config())
+    handler = modules["resize"].MemoryEfficientSafeOpen(str(path), low_memory=True)
+    try:
+        decoder = modules["guard"].DiffusionQuantization(handler, str(path), "NVFP4 test")
+        key = "model.layer.weight"
+        assert decoder.logical_shape(key) == tuple(original.shape)
+        keys = decoder.required_keys(key)
+        stream = handler.async_stream(keys, batch_size=1, prefetch_batches=1, pin_memory=False)
+        tensors = {}
+        try:
+            for source_key in keys:
+                batch = next(stream)
+                assert len(batch) == 1 and batch[0][0] == source_key
+                tensors[source_key] = batch[0][1]
+            actual = decoder.decode(key, tensors, torch.float16)
+        finally:
+            for source_key in tensors:
+                handler.mark_processed(source_key)
+            stream.close()
+        expected = QuantizedTensor(qdata, layout_name, params).dequantize()
+        torch.testing.assert_close(actual, expected)
+        assert actual.shape == original.shape
+    finally:
+        handler.__exit__(None, None, None)

tests/test_merger_filters.py

@@ -1,10 +1,12 @@
 import importlib
+import json
 import sys
 import types
 from pathlib import Path
 
 import pytest
 import torch
+from safetensors import safe_open
 from safetensors.torch import load_file, save_file
 
 
@@ -130,3 +132,47 @@ def test_lora_factor_group_uses_one_normalized_layer_name(monkeypatch, tmp_path,
     assert "diffusion_model.layer.lora_A.weight" in result
     assert "diffusion_model.layer.lora_B.weight" in result
     assert not any(key.endswith(".alpha") for key in result)
+
+
+def test_diffusion_merge_decodes_tensorwise_quantization_and_drops_sidecars(
+    monkeypatch, tmp_path, modules,
+):
+    merger, operations = modules
+    paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
+    quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8)
+    save_file({
+        "model.diffusion_model.quant.weight": torch.tensor([1, 2], dtype=torch.int8),
+        "model.diffusion_model.quant.weight_scale": torch.tensor(2.0),
+        "model.diffusion_model.quant.comfy_quant": quant,
+        "dense.weight": torch.tensor([2.0, 4.0]),
+    }, paths["a"])
+    save_file({
+        "model.diffusion_model.quant.weight": torch.tensor([3, 6], dtype=torch.int8),
+        "model.diffusion_model.quant.weight_scale": torch.tensor(2.0),
+        "model.diffusion_model.quant.comfy_quant": quant,
+        "dense.weight": torch.tensor([6.0, 8.0]),
+    }, paths["b"])
+
+    monkeypatch.setattr(merger.folder_paths, "get_full_path", lambda _kind, name: paths.get(name))
+    monkeypatch.setattr(merger.folder_paths, "get_folder_paths", lambda _kind: [str(tmp_path)])
+    monkeypatch.setattr(merger.folder_paths, "models_dir", str(tmp_path))
+    monkeypatch.setattr(merger, "prepare_for_large_operation", lambda *args: None)
+    monkeypatch.setattr(merger, "cleanup_after_operation", lambda: None)
+    params = _params("quantized_merge")
+    params.update({"model_a": "a", "model_b": "b"})
+
+    merger.MergerLogic.execute_merge(
+        {"model_a": "a", "model_b": "b"}, "Weight-Sum",
+        operations.TWO_MODEL_MODES, params, "diffusion_models",
+    )
+
+    output_path = tmp_path / "diffusion_models" / "quantized_merge.safetensors"
+    result = load_file(str(output_path))
+    key = "model.diffusion_model.quant.weight"
+    assert set(result) == {key, "dense.weight"}
+    assert result[key].tolist() == pytest.approx([4.0, 8.0])
+    assert result[key].dtype == torch.float32
+    assert result["dense.weight"].tolist() == pytest.approx([4.0, 6.0])
+    with safe_open(str(output_path), framework="pt", device="cpu") as output:
+        metadata = output.metadata()
+    assert metadata in (None, {})

tests/test_model_analysis.py

@@ -1,6 +1,7 @@
 import importlib
 import csv
 from io import StringIO
+import json
 import math
 import sys
 import types
@@ -331,6 +332,51 @@ def test_include_filter_limits_analysis_inventory(
     assert sorted(streamed) == sorted(expected_keys * 2)
 
 
+def test_diffusion_analysis_decodes_int8_tensorwise_and_skips_sidecars(
+    monkeypatch, tmp_path, analysis,
+):
+    prefix = "model.diffusion_model.layer"
+    quant = torch.tensor(list(json.dumps({"format": "int8_tensorwise"}).encode()), dtype=torch.uint8)
+    paths = {name: str(tmp_path / f"{name}.safetensors") for name in ("a", "b")}
+    save_file({
+        f"{prefix}.weight": torch.tensor([1, 2], dtype=torch.int8),
+        f"{prefix}.weight_scale": torch.tensor(2.0),
+        f"{prefix}.comfy_quant": quant,
+        "dense.weight": torch.tensor([2.0, 4.0]),
+    }, paths["a"])
+    save_file({
+        f"{prefix}.weight": torch.tensor([3, 6], dtype=torch.int8),
+        f"{prefix}.weight_scale": torch.tensor(2.0),
+        f"{prefix}.comfy_quant": quant,
+        "dense.weight": torch.tensor([4.0, 8.0]),
+    }, paths["b"])
+    monkeypatch.setattr(analysis.folder_paths, "get_full_path", lambda _, name: paths[name])
+    monkeypatch.setattr(analysis, "prepare_for_large_operation", lambda *args, **kwargs: None)
+    monkeypatch.setattr(analysis, "cleanup_after_operation", lambda: None)
+    streamed = []
+    original_stream = analysis.MemoryEfficientSafeOpen.async_stream
+
+    def track_stream(handler, keys, *args, **kwargs):
+        streamed.append((handler.filename, tuple(keys)))
+        return original_stream(handler, keys, *args, **kwargs)
+
+    monkeypatch.setattr(analysis.MemoryEfficientSafeOpen, "async_stream", track_stream)
+    _, _, metrics_csv, _ = analysis.ModelAnalysisLogic.execute(
+        "a", "b", "diffusion_models", _params(),
+    )
+
+    rows = {row["key"]: row for row in csv.DictReader(StringIO(metrics_csv))}
+    assert set(rows) == {f"{prefix}.weight", "dense.weight"}
+    assert float(rows[f"{prefix}.weight"]["MAE"]) == pytest.approx(6.0)
+    assert float(rows["dense.weight"]["MAE"]) == pytest.approx(3.0)
+    for path, keys in streamed:
+        if path.endswith("a.safetensors") or path.endswith("b.safetensors"):
+            if f"{prefix}.weight" in keys:
+                assert f"{prefix}.weight_scale" in keys
+                assert f"{prefix}.comfy_quant" in keys
+    assert sum(f"{prefix}.weight" in keys or "dense.weight" in keys for _, keys in streamed) == 4
+
+
 def test_embedding_similarity_alignment_is_explicit_opt_in(analysis):
     a = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
     b = torch.tensor([[0.0, 1.0], [1.0, 0.0]])