From 9ebe34565c4be35c46fdfccb3fb565c9dedf3f84 Mon Sep 17 00:00:00 2001
From: silveroxides <ishimarukaito@gmail.com>
Date: Sat, 31 Oct 2026 23:16:02 +0100
Subject: [PATCH] Remove arbitrary low-bit tensor count cap and honor discard
 before preservation

---
 nodes/merger.py                       |  8 +++---
 nodes/quantization_guard.py           | 13 ----------
 tests/test_lora_quantization_guard.py | 37 +++++++++++++++------------
 3 files changed, 25 insertions(+), 33 deletions(-)

diff --git a/nodes/merger.py b/nodes/merger.py
index 5c67375..af9afe5 100644
--- a/nodes/merger.py
+++ b/nodes/merger.py
@@ -312,15 +312,15 @@ class MergerLogic:
                     discarded_keys += 1
                     pbar.update(1)
                     continue
-                if preserve_low_bit:
-                    write_preserved_tensor(writer, key, primary_handler, force_raw=True)
-                    pbar.update(1)
-                    continue
                 # Check discard patterns first - skip entirely
                 if _matches_any_pattern(key, discard_patterns, glob_mode=glob_mode):
                     discarded_keys += 1
                     pbar.update(1)
                     continue
+                if preserve_low_bit:
+                    write_preserved_tensor(writer, key, primary_handler, force_raw=True)
+                    pbar.update(1)
+                    continue
 
                 for name, mapping in alpha_normalization.items():
                     alpha_info = mapping.get(key)
diff --git a/nodes/quantization_guard.py b/nodes/quantization_guard.py
index 1e0d208..eb76259 100644
--- a/nodes/quantization_guard.py
+++ b/nodes/quantization_guard.py
@@ -7,9 +7,6 @@ import struct
 import torch
 
 
-MAX_ISOLATED_LOW_BIT_TENSORS = 2
-
-
 def _low_bit_dtypes():
     names = (
         "int8",
@@ -135,11 +132,6 @@ class DiffusionQuantization:
             key for key in self.data_keys
             if key not in self.quantized_keys and _is_low_bit(handler, key)
         }
-        if len(self.isolated_low_bit_keys) > MAX_ISOLATED_LOW_BIT_TENSORS:
-            raise ValueError(
-                f"[{label}] {model_path} contains {len(self.isolated_low_bit_keys)} low-bit tensors; "
-                "casted or quantized inputs are unsupported"
-            )
         if self.isolated_low_bit_keys:
             logging.warning(
                 "[%s] %s contains %d isolated low-bit tensor(s); preserving/excluding unchanged: %s",
@@ -413,11 +405,6 @@ def inspect_low_bit_input(handler, label: str, operation: str) -> set[str]:
         )
 
     low_bit_keys = {key for key in keys if _is_low_bit(handler, key)}
-    if len(low_bit_keys) > MAX_ISOLATED_LOW_BIT_TENSORS:
-        raise ValueError(
-            f"[{operation}] {label} contains {len(low_bit_keys)} low-bit tensors; "
-            "casted or quantized inputs are unsupported"
-        )
     if low_bit_keys:
         logging.warning(
             "[%s] %s contains %d isolated low-bit tensor(s); "
diff --git a/tests/test_lora_quantization_guard.py b/tests/test_lora_quantization_guard.py
index 8cbc056..2e89d36 100644
--- a/tests/test_lora_quantization_guard.py
+++ b/tests/test_lora_quantization_guard.py
@@ -85,17 +85,16 @@ def test_int8_uint8_are_detected_and_warn_once(tmp_path, modules, caplog, dtype)
     assert "low" in records[0].message
 
 
-def test_exact_two_tensor_threshold(tmp_path, modules):
-    for count in range(4):
+def test_multiple_low_bit_config_tensors_allowed(tmp_path, modules):
+    for count in (0, 1, 2, 52):
         path = tmp_path / f"threshold_{count}.safetensors"
-        tensors = {f"low_{i}": torch.ones(2, dtype=torch.uint8) for i in range(count)}
+        tensors = {
+            f"diffusion_model.blocks.{i}.attn.comfy_attention.config": torch.ones(2, dtype=torch.uint8)
+            for i in range(count)
+        }
         tensors["normal"] = torch.ones(1)
         save_file(tensors, str(path))
-        if count <= 2:
-            assert len(_inspect_file(path, modules)) == count
-        else:
-            with pytest.raises(ValueError, match="3 low-bit tensors.*casted or quantized"):
-                _inspect_file(path, modules)
+        assert len(_inspect_file(path, modules)) == count
 
 
 @pytest.mark.parametrize(
@@ -315,18 +314,24 @@ def test_multi_variants_preserve_earliest_complete_affected_layer(
     assert "diffusion_model.normal.diff" in result
 
 
-def test_hard_error_happens_before_resize_output_creation(monkeypatch, tmp_path, modules):
+def test_resize_preserves_multiple_low_bit_config_tensors(monkeypatch, tmp_path, modules):
     resize = modules["resize"]
     _patch_output(monkeypatch, resize, tmp_path)
     source = tmp_path / "three_low.safetensors"
-    save_file({f"low_{i}": torch.ones(1, dtype=torch.uint8) for i in range(3)}, str(source))
-    output = tmp_path / "must_not_exist.safetensors"
+    tensors = {
+        f"diffusion_model.blocks.{i}.attn.comfy_attention.config": torch.ones(1, dtype=torch.uint8)
+        for i in range(3)
+    }
+    tensors["diffusion_model.normal.lora_down.weight"] = torch.ones((2, 2))
+    tensors["diffusion_model.normal.lora_up.weight"] = torch.ones((2, 2))
+    save_file(tensors, str(source))
 
-    with pytest.raises(ValueError, match="3 low-bit tensors"):
-        resize.resize_lora_file(
-            str(source), 1, None, None, "cpu", torch.float16, output.stem, verbose=False
-        )
-    assert not output.exists()
+    output = resize.resize_lora_file(
+        str(source), 1, None, None, "cpu", torch.float16, "three_low_out", verbose=False
+    )
+    result = load_file(output)
+    for i in range(3):
+        assert result[f"diffusion_model.blocks.{i}.attn.comfy_attention.config"].dtype == torch.uint8
 
 
 def test_recognized_quantization_creates_no_resize_output(monkeypatch, tmp_path, modules):