diff --git a/comfy/ops.py b/comfy/ops.py index 14599997b..314ba005b 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -1151,7 +1151,10 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat layer_conf = state_dict.pop(f"{prefix}comfy_quant", None) if layer_conf is not None: - layer_conf = json.loads(layer_conf.numpy().tobytes()) + raw_conf = layer_conf.numpy().tobytes() + # Some quantizers mark unquantized layers with an all-NUL comfy_quant + # placeholder instead of omitting it; treat that the same as absent. + layer_conf = json.loads(raw_conf) if raw_conf.strip(b"\x00") else None if layer_conf is None: module.weight = torch.nn.Parameter(weight.to(device=device, dtype=compute_dtype), requires_grad=False) @@ -1587,7 +1590,10 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec weight_key = f"{prefix}weight" layer_conf = state_dict.pop(f"{prefix}comfy_quant", None) if layer_conf is not None: - layer_conf = json.loads(layer_conf.numpy().tobytes()) + raw_conf = layer_conf.numpy().tobytes() + # Some quantizers mark unquantized layers with an all-NUL comfy_quant + # placeholder instead of omitting it; treat that the same as absent. + layer_conf = json.loads(raw_conf) if raw_conf.strip(b"\x00") else None # Only fp8 and int8_tensorwise support per-row dequant via index select. # Block-scaled formats (NVFP4, MXFP8) can't do per-row lookup efficiently. diff --git a/tests-unit/comfy_quant/test_mixed_precision.py b/tests-unit/comfy_quant/test_mixed_precision.py index 7bbc96616..f697eaf8c 100644 --- a/tests-unit/comfy_quant/test_mixed_precision.py +++ b/tests-unit/comfy_quant/test_mixed_precision.py @@ -228,6 +228,33 @@ class TestMixedPrecisionOps(unittest.TestCase): with self.assertRaises(KeyError): model.load_state_dict(state_dict, strict=False) + def test_all_nul_comfy_quant_marker_loads_as_unquantized(self): + """Some quantizers mark unquantized layers with an all-NUL comfy_quant + placeholder instead of omitting the key; it must load as plain weight, + not crash decoding it as JSON or raise for a missing format.""" + state_dict = { + "layer1.weight": torch.randn(20, 10, dtype=torch.bfloat16), + "layer1.bias": torch.randn(20, dtype=torch.bfloat16), + "layer1.comfy_quant": torch.zeros(29, dtype=torch.uint8), + "layer2.weight": torch.randn(30, 20, dtype=torch.bfloat16), + "layer2.bias": torch.randn(30, dtype=torch.bfloat16), + "layer3.weight": torch.randn(40, 30, dtype=torch.bfloat16), + "layer3.bias": torch.randn(40, dtype=torch.bfloat16), + } + + model = SimpleModel(operations=ops.mixed_precision_ops({})) + model.load_state_dict(state_dict, strict=False) + + self.assertNotIsInstance(model.layer1.weight, QuantizedTensor) + + for layer in [model.layer1, model.layer2, model.layer3]: + layer.weight_function = [] + layer.bias_function = [] + + input_tensor = torch.randn(5, 10, dtype=torch.bfloat16) + output = model(input_tensor) + self.assertEqual(output.shape, (5, 40)) + def test_int8_convrot_metadata_loads_into_params(self): """ConvRot metadata must reach TensorWiseINT8Layout params.""" torch.manual_seed(123)