From 20467f61ee89bb7c9162d45f0e9ba1c9caffb9ed Mon Sep 17 00:00:00 2001 From: Alan Li Date: Sat, 25 Jul 2026 21:49:05 -0400 Subject: [PATCH 1/3] Fix convert_old_quants ignoring model_prefix in _quantization_metadata branch (#11864) The new-format branch of convert_old_quants() writes "{layer_key}.comfy_quant" markers straight from the checkpoint's _quantization_metadata JSON, ignoring model_prefix entirely. Checkpoint metadata stores layer keys either with the full diffusion-model prefix or already stripped of it, and convert_old_quants() is invoked before and/or after that prefix is stripped from state_dict (comfy/sd.py calls it up to twice around the strip). A fixed assumption about which convention is in play silently mismatches the other, so affected layers never get wrapped in a QuantizedTensor and fall back to plain-dtype storage (observed as VRAM blowup / manual-cast warnings on NVFP4 checkpoints such as LTX-Video 2.3). PR #13328 tried to fix this by reordering comfy/sd.py's calls around the prefix strip, but that only swaps which convention works and was closed by the maintainer for risking regressions on checkpoints using the other convention. This instead makes convert_old_quants() match each layer key against the state_dict's real key first (zero behavior change when metadata is already aligned with state_dict), then try adding/stripping model_prefix, and only fall back to today's blind write when neither matches. No call sites in comfy/sd.py are touched or reordered. The write is now also idempotent, guarding against the double-call case rewriting a resolved key. Adds tests-unit/comfy_quant/test_convert_old_quants_prefix.py covering both metadata conventions, the legacy scaled_fp8 branch (unaffected), the exact two-call dance from load_diffusion_model_state_dict, and a documented (xfail) residual gap: when a checkpoint's real weight keys carry no prefix at all but its metadata keys do, comfy/sd.py's model_prefix="" call-site argument still lets model_detection.unet_prefix_from_state_dict get poisoned by the resulting spurious marker keys. That is pre-existing (reproduces identically without this patch) and out of scope here since fixing it needs a comfy/sd.py or model_detection.py change. Co-Authored-By: Claude Fable 5 --- comfy/utils.py | 26 +- .../test_convert_old_quants_prefix.py | 278 ++++++++++++++++++ 2 files changed, 303 insertions(+), 1 deletion(-) create mode 100644 tests-unit/comfy_quant/test_convert_old_quants_prefix.py diff --git a/comfy/utils.py b/comfy/utils.py index 61c2a22dd..3df92aa59 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -1400,6 +1400,27 @@ def detect_layer_quantization(state_dict, prefix): return {"mixed_ops": True} return None +def _resolve_quant_metadata_layer_key(state_dict, layer_key, model_prefix): + # _quantization_metadata layer keys come from two different checkpoint + # export conventions in the wild: some store them with the full + # diffusion-model prefix (e.g. "model.diffusion_model.proj_in"), others + # already strip it (e.g. "proj_in"). convert_old_quants() can also be + # called before or after model_prefix is stripped from state_dict + # (comfy/sd.py calls it up to twice around that strip, see #13328), so a + # single fixed convention silently mismatches one case or the other and + # the layer never gets wrapped in a QuantizedTensor (see #11864). Match + # against state_dict's real keys instead of assuming a convention. + if "{}.weight".format(layer_key) in state_dict: + return layer_key + if model_prefix: + if layer_key.startswith(model_prefix): + candidate = layer_key[len(model_prefix):] + else: + candidate = "{}{}".format(model_prefix, layer_key) + if "{}.weight".format(candidate) in state_dict: + return candidate + return layer_key + def convert_old_quants(state_dict, model_prefix="", metadata={}): if metadata is None: metadata = {} @@ -1456,7 +1477,10 @@ def convert_old_quants(state_dict, model_prefix="", metadata={}): if quant_metadata is not None: layers = quant_metadata["layers"] for k, v in layers.items(): - state_dict["{}.comfy_quant".format(k)] = torch.tensor(list(json.dumps(v).encode('utf-8')), dtype=torch.uint8) + resolved_key = _resolve_quant_metadata_layer_key(state_dict, k, model_prefix) + marker_key = "{}.comfy_quant".format(resolved_key) + if marker_key not in state_dict: # idempotent: convert_old_quants may run twice on the same checkpoint + state_dict[marker_key] = torch.tensor(list(json.dumps(v).encode('utf-8')), dtype=torch.uint8) return state_dict, metadata diff --git a/tests-unit/comfy_quant/test_convert_old_quants_prefix.py b/tests-unit/comfy_quant/test_convert_old_quants_prefix.py new file mode 100644 index 000000000..7645db68f --- /dev/null +++ b/tests-unit/comfy_quant/test_convert_old_quants_prefix.py @@ -0,0 +1,278 @@ +import unittest +import torch +import sys +import os +import json + +# Add comfy to path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) + +def has_gpu(): + return torch.cuda.is_available() + +from comfy.cli_args import args +if not has_gpu(): + args.cpu = True + +from comfy import ops +from comfy.quant_ops import QuantizedTensor +import comfy.utils +import comfy.model_detection as model_detection + + +def marker_json(state_dict, key): + """Decode a `.comfy_quant` marker tensor back into its layer_conf dict.""" + return json.loads(state_dict[key].numpy().tobytes()) + + +class SimpleModel(torch.nn.Module): + """Mirrors tests-unit/comfy_quant/test_mixed_precision.py::SimpleModel.""" + + def __init__(self, operations=ops.disable_weight_init): + super().__init__() + self.layer1 = operations.Linear(10, 20, device="cpu", dtype=torch.bfloat16) + + def forward(self, x): + return self.layer1(x) + + +class TestConvertOldQuantsPrefixAware(unittest.TestCase): + """Regression tests for GitHub #11864 / #13328: convert_old_quants()'s + new-format (_quantization_metadata) branch must match the layer key + convention actually used by the state_dict it is given, regardless of + whether that convention is prefixed or already stripped, and regardless + of whether model_prefix is stripped/added/empty at the call site. + """ + + # ---- scenario 1: metadata key and sd key both carry the prefix (aligned) ---- + def test_scenario1_prefixed_metadata_matches_prefixed_sd(self): + layer_quant_config = {"model.diffusion_model.proj_in": {"format": "float8_e4m3fn"}} + state_dict = { + "model.diffusion_model.proj_in.weight": torch.randn(4, 4, dtype=torch.float32).to(torch.float8_e4m3fn), + "model.diffusion_model.proj_in.weight_scale": torch.tensor(1.0), + "model.diffusion_model.other.weight": torch.randn(4, 4, dtype=torch.bfloat16), + } + out_sd, _ = comfy.utils.convert_old_quants( + state_dict, + model_prefix="model.diffusion_model.", + metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})}, + ) + self.assertIn("model.diffusion_model.proj_in.comfy_quant", out_sd) + self.assertNotIn("proj_in.comfy_quant", out_sd) + self.assertEqual(marker_json(out_sd, "model.diffusion_model.proj_in.comfy_quant")["format"], "float8_e4m3fn") + + # ---- scenario 2: sd already stripped of prefix, metadata key still prefixed (the bug) ---- + def test_scenario2_prefixed_metadata_matches_stripped_sd_after_fix(self): + layer_quant_config = {"model.diffusion_model.proj_in": {"format": "float8_e4m3fn"}} + state_dict = { + "proj_in.weight": torch.randn(4, 4, dtype=torch.float32).to(torch.float8_e4m3fn), + "proj_in.weight_scale": torch.tensor(1.0), + "other.weight": torch.randn(4, 4, dtype=torch.bfloat16), + } + out_sd, _ = comfy.utils.convert_old_quants( + state_dict, + model_prefix="model.diffusion_model.", + metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})}, + ) + # Before the fix this would have blindly written + # "model.diffusion_model.proj_in.comfy_quant", which never matches + # "proj_in.weight" -> detect_layer_quantization()/MixedPrecisionOps + # would find no marker and load the layer as a plain dtype tensor. + self.assertIn("proj_in.comfy_quant", out_sd) + self.assertNotIn("model.diffusion_model.proj_in.comfy_quant", out_sd) + self.assertEqual(marker_json(out_sd, "proj_in.comfy_quant")["format"], "float8_e4m3fn") + + # ---- scenario 3: metadata key already stripped, sd already stripped (today's working path) ---- + def test_scenario3_stripped_metadata_matches_stripped_sd_unchanged(self): + layer_quant_config = {"proj_in": {"format": "float8_e4m3fn"}} + state_dict = { + "proj_in.weight": torch.randn(4, 4, dtype=torch.float32).to(torch.float8_e4m3fn), + "proj_in.weight_scale": torch.tensor(1.0), + } + out_sd, _ = comfy.utils.convert_old_quants( + state_dict, + model_prefix="model.diffusion_model.", + metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})}, + ) + # Zero behavior change vs. today: direct match succeeds immediately, + # key is written exactly as it always was. + self.assertIn("proj_in.comfy_quant", out_sd) + self.assertEqual(marker_json(out_sd, "proj_in.comfy_quant")["format"], "float8_e4m3fn") + + # Same as scenario 3 but with model_prefix="" (the literal value passed + # by comfy/sd.py::load_diffusion_model_state_dict at both call sites). + def test_scenario3b_stripped_metadata_empty_model_prefix_unchanged(self): + layer_quant_config = {"proj_in": {"format": "nvfp4"}} + state_dict = { + "proj_in.weight": torch.randint(0, 255, (4, 2), dtype=torch.uint8), + "proj_in.weight_scale": torch.tensor(1.0), + "proj_in.weight_scale_2": torch.tensor(1.0), + } + out_sd, _ = comfy.utils.convert_old_quants( + state_dict, + model_prefix="", + metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})}, + ) + self.assertIn("proj_in.comfy_quant", out_sd) + + # ---- scenario 4: legacy scaled_fp8 branch must be completely unaffected ---- + def test_scenario4_legacy_scaled_fp8_branch_unaffected(self): + state_dict = { + "model.diffusion_model.scaled_fp8": torch.tensor([0.0], dtype=torch.float32), + "model.diffusion_model.proj_in.weight": torch.randn(4, 4, dtype=torch.float32).to(torch.float8_e4m3fn), + "model.diffusion_model.proj_in.scale_weight": torch.tensor(2.0), + "model.diffusion_model.other.weight": torch.randn(4, 4, dtype=torch.bfloat16), + } + out_sd, metadata = comfy.utils.convert_old_quants( + state_dict, + model_prefix="model.diffusion_model.", + metadata={}, + ) + # Old-format branch derives layer keys straight from state_dict's own + # (already correctly prefixed) keys, so resolution is a same-key + # direct match every time -- this path must be byte-for-byte identical + # to pre-fix behavior. + self.assertNotIn("model.diffusion_model.scaled_fp8", out_sd) + self.assertIn("model.diffusion_model.proj_in.weight_scale", out_sd) + self.assertIn("model.diffusion_model.proj_in.comfy_quant", out_sd) + self.assertEqual( + marker_json(out_sd, "model.diffusion_model.proj_in.comfy_quant")["format"], + "float8_e4m3fn", + ) + self.assertNotIn("proj_in.comfy_quant", out_sd) # not stripped/mismatched + + # ---- extra: idempotency across the exact two-call pattern comfy/sd.py uses ---- + def test_two_call_pattern_mirrors_load_diffusion_model_state_dict(self): + """Simulates comfy/sd.py::load_diffusion_model_state_dict() verbatim: + convert_old_quants(sd, "", metadata=metadata) is called once before + the diffusion_model_prefix strip and once after, both with an empty + model_prefix. This must work for BOTH metadata conventions without + any change to the call site (see PR #13328, closed for reordering + the calls and breaking the other convention instead).""" + for convention, layer_key in ( + ("prefixed", "model.diffusion_model.proj_in"), + ("stripped", "proj_in"), + ): + with self.subTest(convention=convention): + metadata = {"_quantization_metadata": json.dumps( + {"layers": {layer_key: {"format": "float8_e4m3fn"}}} + )} + sd = { + "model.diffusion_model.proj_in.weight": torch.randn(4, 4, dtype=torch.float32).to(torch.float8_e4m3fn), + "model.diffusion_model.proj_in.weight_scale": torch.tensor(1.0), + "unrelated.top_level.weight": torch.randn(2, 2, dtype=torch.bfloat16), + } + + # call 1: before stripping, model_prefix="" (as sd.py does) + sd, metadata = comfy.utils.convert_old_quants(sd, "", metadata=metadata) + + # simulate state_dict_prefix_replace(sd, {prefix: ""}, filter_keys=True) + prefix = "model.diffusion_model." + temp_sd = {k[len(prefix):]: v for k, v in sd.items() if k.startswith(prefix)} + self.assertGreater(len(temp_sd), 0) + sd = temp_sd + + # call 2: after stripping, model_prefix="" again (as sd.py does) + sd, metadata = comfy.utils.convert_old_quants(sd, "", metadata=metadata) + + self.assertIn("proj_in.comfy_quant", sd, + f"{convention} metadata convention did not resolve after the two-call dance") + self.assertEqual(marker_json(sd, "proj_in.comfy_quant")["format"], "float8_e4m3fn") + + # ---- extra: repeated calls with identical inputs don't duplicate/clobber ---- + def test_marker_write_is_idempotent(self): + layer_quant_config = {"proj_in": {"format": "float8_e4m3fn"}} + metadata = {"_quantization_metadata": json.dumps({"layers": layer_quant_config})} + state_dict = { + "proj_in.weight": torch.randn(4, 4, dtype=torch.float32).to(torch.float8_e4m3fn), + "proj_in.weight_scale": torch.tensor(1.0), + } + out_sd1, _ = comfy.utils.convert_old_quants(dict(state_dict), model_prefix="", metadata=dict(metadata)) + out_sd2, _ = comfy.utils.convert_old_quants(out_sd1, model_prefix="", metadata=dict(metadata)) + self.assertEqual(set(out_sd1.keys()), set(out_sd2.keys())) + self.assertTrue(torch.equal(out_sd1["proj_in.comfy_quant"], out_sd2["proj_in.comfy_quant"])) + + # ---- extra: functional end-to-end, proving the fixed layer actually loads as QuantizedTensor ---- + def test_functional_load_after_prefix_mismatch_fix(self): + layer_quant_config = {"model.diffusion_model.layer1": {"format": "float8_e4m3fn"}} + fp8_weight = torch.randn(20, 10, dtype=torch.float32).to(torch.float8_e4m3fn) + # sd already stripped of "model.diffusion_model." (as it is by the + # time comfy/sd.py's second convert_old_quants call runs), while + # metadata still carries the full prefix. + state_dict = { + "layer1.weight": fp8_weight, + "layer1.bias": torch.randn(20, dtype=torch.bfloat16), + "layer1.weight_scale": torch.tensor(2.0, dtype=torch.float32), + } + state_dict, _ = comfy.utils.convert_old_quants( + state_dict, + model_prefix="model.diffusion_model.", + metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})}, + ) + model = SimpleModel(operations=ops.mixed_precision_ops({})) + model.load_state_dict(state_dict, strict=False) + + self.assertIsInstance(model.layer1.weight, QuantizedTensor) + self.assertEqual(model.layer1.weight._layout_cls, "TensorCoreFP8E4M3Layout") + self.assertEqual(model.layer1.weight._params.scale.item(), 2.0) + + +class TestKnownResidualGap(unittest.TestCase): + """Documents a related but DISTINCT failure mode this PR does not close, + found while writing the regression tests above. Kept as an + expectedFailure so it stays visible instead of silently passing or + breaking CI. + + load_diffusion_model_state_dict() passes model_prefix="" (a string + literal, not the real diffusion-model prefix) to both of its + convert_old_quants() calls -- see comfy/sd.py, unchanged by this PR since + editing that call site is the path #13328 was closed for. When a + checkpoint's real weight keys carry NO wrapper prefix at all but its + _quantization_metadata layer keys DO carry a "model.diffusion_model." + style prefix, convert_old_quants() (both before and after this fix) has + no way to recognize the mismatch on the first call with model_prefix="", + so its fallback write -- required to be byte-for-byte identical to + today's behavior for the "neither convention matches" case -- pollutes + state_dict with spurious "model.diffusion_model.*.comfy_quant" keys. + comfy.model_detection.unet_prefix_from_state_dict() (a different + function, out of this PR's scope) then falsely detects + "model.diffusion_model." as the prefix from those spurious keys alone, + and state_dict_prefix_replace(..., filter_keys=True) strips using that + wrong prefix, discarding every real .weight/.weight_scale tensor. + + This is pre-existing: it reproduces identically with and without this + PR's fix (verified manually), because the root cause here is + unet_prefix_from_state_dict()/the sd.py call site, not the marker-key + resolution this PR changes. Closing it would require either passing the + real prefix into convert_old_quants() from sd.py (a call-site change) or + hardening unet_prefix_from_state_dict() to ignore .comfy_quant keys -- + both outside "only touch convert_old_quants()". + """ + + @unittest.expectedFailure + def test_empty_model_prefix_cannot_prevent_real_prefix_detection_poisoning(self): + prefix = "model.diffusion_model." + num_layers = 10 # unet_prefix_from_state_dict requires > 5 matches + sd = {} + layers_meta = {} + for i in range(num_layers): + local = f"block{i}" + sd[f"{local}.weight"] = torch.randn(4, 4, dtype=torch.float32).to(torch.float8_e4m3fn) + sd[f"{local}.weight_scale"] = torch.tensor(1.0) + layers_meta[f"{prefix}{local}"] = {"format": "float8_e4m3fn"} + metadata = {"_quantization_metadata": json.dumps({"layers": layers_meta})} + + # exact two-call dance from comfy/sd.py::load_diffusion_model_state_dict + sd, metadata = comfy.utils.convert_old_quants(sd, "", metadata=metadata) + diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd) + temp_sd = comfy.utils.state_dict_prefix_replace(sd, {diffusion_model_prefix: ""}, filter_keys=True) + if len(temp_sd) > 0: + sd = temp_sd + sd, metadata = comfy.utils.convert_old_quants(sd, "", metadata=metadata) + + remaining_weights = [k for k in sd if k.endswith(".weight")] + self.assertEqual(len(remaining_weights), num_layers) + + +if __name__ == "__main__": + unittest.main() From 50af1a356362d5401c1c557465bccd5d63d67bd3 Mon Sep 17 00:00:00 2001 From: Alan Li Date: Sun, 26 Jul 2026 21:35:50 -0400 Subject: [PATCH 2/3] Fix quant metadata marker collision handling --- comfy/utils.py | 16 ++-- .../test_convert_old_quants_prefix.py | 84 +++++-------------- 2 files changed, 28 insertions(+), 72 deletions(-) diff --git a/comfy/utils.py b/comfy/utils.py index 3df92aa59..e1a099fbf 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -1401,15 +1401,8 @@ def detect_layer_quantization(state_dict, prefix): return None def _resolve_quant_metadata_layer_key(state_dict, layer_key, model_prefix): - # _quantization_metadata layer keys come from two different checkpoint - # export conventions in the wild: some store them with the full - # diffusion-model prefix (e.g. "model.diffusion_model.proj_in"), others - # already strip it (e.g. "proj_in"). convert_old_quants() can also be - # called before or after model_prefix is stripped from state_dict - # (comfy/sd.py calls it up to twice around that strip, see #13328), so a - # single fixed convention silently mismatches one case or the other and - # the layer never gets wrapped in a QuantizedTensor (see #11864). Match - # against state_dict's real keys instead of assuming a convention. + # Metadata layer keys may use the full diffusion-model prefix or the + # already-stripped form, so resolve them against the state dict. if "{}.weight".format(layer_key) in state_dict: return layer_key if model_prefix: @@ -1479,8 +1472,9 @@ def convert_old_quants(state_dict, model_prefix="", metadata={}): for k, v in layers.items(): resolved_key = _resolve_quant_metadata_layer_key(state_dict, k, model_prefix) marker_key = "{}.comfy_quant".format(resolved_key) - if marker_key not in state_dict: # idempotent: convert_old_quants may run twice on the same checkpoint - state_dict[marker_key] = torch.tensor(list(json.dumps(v).encode('utf-8')), dtype=torch.uint8) + marker = torch.tensor(list(json.dumps(v).encode('utf-8')), dtype=torch.uint8) + if marker_key not in state_dict or not torch.equal(state_dict[marker_key], marker): + state_dict[marker_key] = marker return state_dict, metadata diff --git a/tests-unit/comfy_quant/test_convert_old_quants_prefix.py b/tests-unit/comfy_quant/test_convert_old_quants_prefix.py index 7645db68f..a2aa7da17 100644 --- a/tests-unit/comfy_quant/test_convert_old_quants_prefix.py +++ b/tests-unit/comfy_quant/test_convert_old_quants_prefix.py @@ -17,7 +17,6 @@ if not has_gpu(): from comfy import ops from comfy.quant_ops import QuantizedTensor import comfy.utils -import comfy.model_detection as model_detection def marker_json(state_dict, key): @@ -188,9 +187,29 @@ class TestConvertOldQuantsPrefixAware(unittest.TestCase): "proj_in.weight_scale": torch.tensor(1.0), } out_sd1, _ = comfy.utils.convert_old_quants(dict(state_dict), model_prefix="", metadata=dict(metadata)) - out_sd2, _ = comfy.utils.convert_old_quants(out_sd1, model_prefix="", metadata=dict(metadata)) - self.assertEqual(set(out_sd1.keys()), set(out_sd2.keys())) - self.assertTrue(torch.equal(out_sd1["proj_in.comfy_quant"], out_sd2["proj_in.comfy_quant"])) + keys_before = set(out_sd1.keys()) + original_marker = out_sd1["proj_in.comfy_quant"] + marker_before = original_marker.clone() + out_sd2, _ = comfy.utils.convert_old_quants(dict(out_sd1), model_prefix="", metadata=dict(metadata)) + self.assertEqual(keys_before, set(out_sd2.keys())) + self.assertTrue(torch.equal(marker_before, out_sd2["proj_in.comfy_quant"])) + self.assertIs(original_marker, out_sd2["proj_in.comfy_quant"]) + + def test_conflicting_marker_is_replaced_with_current_metadata(self): + layer_quant_config = {"proj_in": {"format": "float8_e4m3fn"}} + old_marker = torch.tensor(list(json.dumps({"format": "nvfp4"}).encode("utf-8")), dtype=torch.uint8) + state_dict = { + "proj_in.weight": torch.randn(4, 4, dtype=torch.float32).to(torch.float8_e4m3fn), + "proj_in.weight_scale": torch.tensor(1.0), + "proj_in.comfy_quant": old_marker, + } + out_sd, _ = comfy.utils.convert_old_quants( + state_dict, + model_prefix="", + metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})}, + ) + self.assertEqual(marker_json(out_sd, "proj_in.comfy_quant"), {"format": "float8_e4m3fn"}) + self.assertIsNot(old_marker, out_sd["proj_in.comfy_quant"]) # ---- extra: functional end-to-end, proving the fixed layer actually loads as QuantizedTensor ---- def test_functional_load_after_prefix_mismatch_fix(self): @@ -217,62 +236,5 @@ class TestConvertOldQuantsPrefixAware(unittest.TestCase): self.assertEqual(model.layer1.weight._params.scale.item(), 2.0) -class TestKnownResidualGap(unittest.TestCase): - """Documents a related but DISTINCT failure mode this PR does not close, - found while writing the regression tests above. Kept as an - expectedFailure so it stays visible instead of silently passing or - breaking CI. - - load_diffusion_model_state_dict() passes model_prefix="" (a string - literal, not the real diffusion-model prefix) to both of its - convert_old_quants() calls -- see comfy/sd.py, unchanged by this PR since - editing that call site is the path #13328 was closed for. When a - checkpoint's real weight keys carry NO wrapper prefix at all but its - _quantization_metadata layer keys DO carry a "model.diffusion_model." - style prefix, convert_old_quants() (both before and after this fix) has - no way to recognize the mismatch on the first call with model_prefix="", - so its fallback write -- required to be byte-for-byte identical to - today's behavior for the "neither convention matches" case -- pollutes - state_dict with spurious "model.diffusion_model.*.comfy_quant" keys. - comfy.model_detection.unet_prefix_from_state_dict() (a different - function, out of this PR's scope) then falsely detects - "model.diffusion_model." as the prefix from those spurious keys alone, - and state_dict_prefix_replace(..., filter_keys=True) strips using that - wrong prefix, discarding every real .weight/.weight_scale tensor. - - This is pre-existing: it reproduces identically with and without this - PR's fix (verified manually), because the root cause here is - unet_prefix_from_state_dict()/the sd.py call site, not the marker-key - resolution this PR changes. Closing it would require either passing the - real prefix into convert_old_quants() from sd.py (a call-site change) or - hardening unet_prefix_from_state_dict() to ignore .comfy_quant keys -- - both outside "only touch convert_old_quants()". - """ - - @unittest.expectedFailure - def test_empty_model_prefix_cannot_prevent_real_prefix_detection_poisoning(self): - prefix = "model.diffusion_model." - num_layers = 10 # unet_prefix_from_state_dict requires > 5 matches - sd = {} - layers_meta = {} - for i in range(num_layers): - local = f"block{i}" - sd[f"{local}.weight"] = torch.randn(4, 4, dtype=torch.float32).to(torch.float8_e4m3fn) - sd[f"{local}.weight_scale"] = torch.tensor(1.0) - layers_meta[f"{prefix}{local}"] = {"format": "float8_e4m3fn"} - metadata = {"_quantization_metadata": json.dumps({"layers": layers_meta})} - - # exact two-call dance from comfy/sd.py::load_diffusion_model_state_dict - sd, metadata = comfy.utils.convert_old_quants(sd, "", metadata=metadata) - diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd) - temp_sd = comfy.utils.state_dict_prefix_replace(sd, {diffusion_model_prefix: ""}, filter_keys=True) - if len(temp_sd) > 0: - sd = temp_sd - sd, metadata = comfy.utils.convert_old_quants(sd, "", metadata=metadata) - - remaining_weights = [k for k in sd if k.endswith(".weight")] - self.assertEqual(len(remaining_weights), num_layers) - - if __name__ == "__main__": unittest.main() From 48454d17b629705f415f1064c421dbbb912cc578 Mon Sep 17 00:00:00 2001 From: Alan Li Date: Sun, 26 Jul 2026 21:42:15 -0400 Subject: [PATCH 3/3] Avoid duplicate quant metadata markers --- comfy/utils.py | 5 +++++ tests-unit/comfy_quant/test_convert_old_quants_prefix.py | 1 + 2 files changed, 6 insertions(+) diff --git a/comfy/utils.py b/comfy/utils.py index e1a099fbf..799063273 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -1412,6 +1412,11 @@ def _resolve_quant_metadata_layer_key(state_dict, layer_key, model_prefix): candidate = "{}{}".format(model_prefix, layer_key) if "{}.weight".format(candidate) in state_dict: return candidate + parts = layer_key.split(".") + for i in range(1, len(parts)): + candidate = ".".join(parts[i:]) + if "{}.weight".format(candidate) in state_dict and "{}.comfy_quant".format(candidate) in state_dict: + return candidate return layer_key def convert_old_quants(state_dict, model_prefix="", metadata={}): diff --git a/tests-unit/comfy_quant/test_convert_old_quants_prefix.py b/tests-unit/comfy_quant/test_convert_old_quants_prefix.py index a2aa7da17..86366ab18 100644 --- a/tests-unit/comfy_quant/test_convert_old_quants_prefix.py +++ b/tests-unit/comfy_quant/test_convert_old_quants_prefix.py @@ -176,6 +176,7 @@ class TestConvertOldQuantsPrefixAware(unittest.TestCase): self.assertIn("proj_in.comfy_quant", sd, f"{convention} metadata convention did not resolve after the two-call dance") + self.assertNotIn("model.diffusion_model.proj_in.comfy_quant", sd) self.assertEqual(marker_json(sd, "proj_in.comfy_quant")["format"], "float8_e4m3fn") # ---- extra: repeated calls with identical inputs don't duplicate/clobber ----