From 50af1a356362d5401c1c557465bccd5d63d67bd3 Mon Sep 17 00:00:00 2001 From: Alan Li Date: Sun, 26 Jul 2026 21:35:50 -0400 Subject: [PATCH] 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()