Fix crash loading quantized checkpoints with all-NUL comfy_quant markers

[chelsealong]

Summary

Some quantizers publish mixed-precision checkpoints where unquantized
layers carry an all-NUL comfy_quant placeholder tensor instead of
omitting the key entirely, and where genuinely-quantized layers carry
a valid comfy_quant payload with a weight_scale but no explicit
format key (e.g. the official MiniMax H3 nvfp4-AWQ text encoder,
Comfy-Org/MiniMax-H3). Loading such a checkpoint currently fails in
three stages:

1. json.loads(layer_conf.numpy().tobytes()) misdetects the leading
   NUL bytes of an all-NUL marker as UTF-32 and raises
   UnicodeDecodeError.
2. After working around that, the resulting {} config still fails
   with Unknown quantization format for layer ... because the loader
   treats any present comfy_quant key as requiring an explicit
   format, even when the payload is empty.
3. Separately, layers whose comfy_quant payload is valid, non-empty
   JSON with a weight_scale but no format key (e.g. a
   q_proj-style layer in the same checkpoint) hit the same
   Unknown quantization format error, since the loader never falls
   back to inferring a format for those.

This fixes all three in comfy/ops.py's two load paths
(_load_quantized_module, used by Linear/MoE Linear, and the
mixed-precision Embedding's own inline handling) by:

• treating an all-NUL marker the same as an absent one, so the layer
  loads as a plain unquantized weight in compute dtype, and
• when format is missing but a weight_scale key is present,
  inferring the format from the on-disk weight dtype
  (torch.int8 -> int8_tensorwise, torch.float8_e4m3fn ->
  float8_e4m3fn, torch.uint8 -> nvfp4), matching the workaround
  documented by the issue reporter.

Fixes #15400

Test plan

• test_all_nul_comfy_quant_marker_loads_as_unquantized (existing):
  builds a state dict with a torch.zeros(29, dtype=torch.uint8)
comfy_quant marker (matching the shape/dtype reported in the
  issue) and asserts the layer loads as a plain torch.nn.Parameter
  and the model runs a forward pass.
• test_formatless_scaled_comfy_quant_infers_format_from_dtype (new):
  builds a state dict with a Linear layer whose comfy_quant is
  valid empty JSON ({}), an int8 weight, and a weight_scale
  key — mirroring a real q_proj-style layer — and asserts the format
  is inferred as int8_tensorwise and the model runs a forward pass.
• test_formatless_scaled_comfy_quant_embedding_infers_format_from_dtype
  (new): same scenario against the Embedding load path, which has
  its own inline comfy_quant handling separate from
  _load_quantized_module.

Confirmed both new tests fail on the pre-fix code
(git stash the comfy/ops.py change) with the exact
Unknown quantization format for layer layer1 error from the issue:

$ python -m pytest tests-unit/comfy_quant -q -k formatless
...
ValueError: Unknown quantization format for layer layer1
2 failed, 8 deselected in 2.48s

With the fix in place:

$ python -m pytest tests-unit/comfy_quant -q
..........
10 passed in 2.40s

ruff check comfy/ops.py tests-unit/comfy_quant/test_mixed_precision.py
passes.

---

This change was written with AI assistance (Claude).

[coderabbitai]

📝 Walkthrough

Walkthrough

The change adds dtype-based quantization format inference for int8, FP8, and NVFP4 weights when format metadata is missing. All-NUL comfy_quant metadata is treated as absent for quantized modules and embeddings. Embedding loading rejects inferred NVFP4 formats. Regression tests cover unquantized loading, format inference, stale-state reset, error handling, and forward execution.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the primary fix for loading quantized checkpoints with all-NUL markers.
Description check	✅ Passed	The description directly explains the loading failures, implemented fixes, affected paths, and test results.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@comfy/ops.py`:
- Around line 1169-1174: In the layer_conf is None branch of the weight-loading
logic, clear the module’s existing quantization state before installing the
plain Parameter: reset both layout_type and quant_format to None. Leave the
quantized configuration handling in the else branch unchanged.
- Around line 1613-1616: In the quantization inference block of the
embedding-loading path, after
`_QUANT_FORMAT_BY_WEIGHT_DTYPE.get(_stored_weight.dtype)` assigns
`quant_format`, explicitly reject the inferred `nvfp4` format with a clear error
before the default loader can process the uint8 weight. Preserve existing
behavior for supported formats and cases where no format is inferred.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 94311e10-b304-4ca3-89ed-acf3e441e09a

📥 Commits

Reviewing files that changed from the base of the PR and between 43cb4ff and 6fc556d.

📒 Files selected for processing (2)

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

📜 Review details

🧰 Additional context used

📓 Path-based instructions (6)

**/*

📄 CodeRabbit inference engine (AGENTS.md)

**/*: Keep changes small, direct, and limited to the narrowest necessary code path and smallest number of files.
Prefer practical fixes, minimal dependencies, and existing repository patterns; remove obsolete, dead, unreachable, or unused code.
Preserve existing APIs, node names, model-loading behavior, file layout, and workflow compatibility unless replacement is explicitly intended.
Core ComfyUI must not add outbound internet requests, telemetry, tracking, reporting, remote configuration, or background network activity. User-authorized model downloads are limited to the requested artifact and must exclude telemetry and unrelated metadata.

Files:

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: Keep state and capability flags on the object that owns the behavior. Prefer explicit parent-owned attributes over probing child objects with getattr; use child checks only when the child owns the delegated behavior.
Preserve shared method signatures, argument order, return shapes, side effects, and error behavior unless every affected caller and interface is intentionally updated.
Do not add unused compatibility parameters, flags, attributes, constructor options, or model-specific options to shared helpers; keep one-off behavior at the integration boundary.
Normalize third-party return conventions at integration boundaries so core code receives the expected type and shape; avoid undocumented caller-side unwrapping.
Do not add torch.no_grad, torch.inference_mode, or inference-mode wrappers. Do not add model freeze/unfreeze toggles; only disable globally enabled inference mode when a training path requires gradients.
Remove inference-only training behavior such as dropout while preserving checkpoint and state-dict compatibility; use nn.Identity when deleting a module would alter keys or ordering.
Keep imports at module scope except established optional-backend probes or imports required to avoid cycles; avoid unnecessary try/except blocks and use specific exceptions with useful fallbacks.
Do not add workarounds for unsupported library versions, especially PyTorch exception-and-float-cast retries, unless a comment names the exact versions still requiring them.
Let unsupported model formats, invalid quantization metadata, and bad states fail with clear errors instead of silently degrading output.
Match local style, keep comments sparse and useful, and remove comments that merely restate obvious code.
Treat dtype, device placement, VRAM use, and offloading as correctness concerns across CPU, CUDA, ROCm, MPS, DirectML, XPU, NPU, and low-VRAM environments.
Prefer existing ComfyUI and Comfy Kitchen operations, quantization helpers, cast/offload helpe...

Files:

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

**/*.{py,json}

📄 CodeRabbit inference engine (AGENTS.md)

Treat legacy combo, io.Combo, and io.DynamicCombo values affecting filesystem access as untrusted; revalidate them at load/save boundaries with folder_paths, containment checks, or fixed allowlists.

Files:

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

**/*.{py,md,txt,json}

📄 CodeRabbit inference engine (AGENTS.md)

Keep warning and info messages short and actionable, remove noisy or misleading logging, and make documentation edits concise, factual, and tied to changed behavior.

Files:

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

**

⚙️ CodeRabbit configuration file

**: IMPORTANT: Only comment on issues directly introduced by this PR's code changes.
Treat AGENTS.md as mandatory repository policy, not optional style guidance.
Flag PR changes that violate AGENTS.md even when the code is otherwise functional.
In particular, enforce architecture boundaries, dtype/device/memory rules,
interface contracts, import style, no unnecessary try/except blocks, no inline
imports, no outbound internet paths in core ComfyUI, and narrow scoped fixes.
Prefer direct findings over suggestions when a rule is violated. Only ignore
AGENTS.md when it clearly conflicts with a newer explicit maintainer instruction
in the PR.
Do NOT flag pre-existing issues in code that was merely moved, re-indented,
de-indented, or reformatted without logic changes. If code appears in the diff
only due to whitespace or structural reformatting (e.g., removing a with: block),
treat it as unchanged. Contributors should not feel obligated to address
pre-existing issues outside the scope of their contribution.

Files:

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

comfy/**

⚙️ CodeRabbit configuration file

comfy/**: Core ML/diffusion engine. Focus on:

• Backward compatibility (breaking changes affect all custom nodes)
• Memory management and GPU resource handling
• Performance implications in hot paths
• Thread safety for concurrent execution

Files:

• comfy/ops.py

🧠 Learnings (3)

📚 Learning: 2026-02-21T14:01:41.482Z

Learnt from: pythongosssss
Repo: Comfy-Org/ComfyUI PR: 12555
File: comfy_extras/nodes_glsl.py:719-724
Timestamp: 2026-02-21T14:01:41.482Z
Learning: In PyOpenGL, bare Python scalars can be accepted for 1-element array parameters by NumberHandler. This means you can pass an int/float directly to OpenGL texture deletion (e.g., glDeleteTextures(tex)) without wrapping in a list. Verify function-specific expectations and ensure types match what the OpenGL call expects; use explicit lists only when the API requires an array.

Applied to files:

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

📚 Learning: 2026-05-13T12:31:45.069Z

Learnt from: rattus128
Repo: Comfy-Org/ComfyUI PR: 13802
File: comfy/pinned_memory.py:19-30
Timestamp: 2026-05-13T12:31:45.069Z
Learning: When reviewing code that uses comfy/pinned_memory.py’s `HostBuffer.extend(size=..., reallocate=...)`: by default (`reallocate` is not True / False), `extend(size=...)` is a *relative increment* that grows the buffer by `size` bytes—so slicing like `[offset:offset+size]` after `hostbuf.extend(size=size)` is correct and the argument should not be rewritten to `offset + size`. Only in the single-segment reallocation mode (`reallocate=True`, e.g., as used by `resize_pin_buffer()` in `comfy/model_management.py`) should `size` be treated as an *absolute target* and the call/arguments should be checked accordingly.

Applied to files:

• comfy/ops.py

📚 Learning: 2026-08-06T22:18:59.719Z

Learnt from: kijai
Repo: Comfy-Org/ComfyUI PR: 15362
File: comfy/ldm/wan/model_animate2.py:186-223
Timestamp: 2026-08-06T22:18:59.719Z
Learning: When reviewing ComfyUI quantization code, treat `comfy.quant_ops.TensorWiseINT8Layout` and `comfy.quant_ops.TensorCoreConvRotW4A4Layout` as re-exports from `comfy_kitchen`. Validate their behavior against the re-exported `comfy_kitchen` implementations rather than assuming they are local fallback classes.

Applied to files:

• comfy/ops.py

🪛 ast-grep (0.45.0)

tests-unit/comfy_quant/test_mixed_precision.py

[info] 264-264: use jsonify instead of json.dumps for JSON output
Context: json.dumps({})
Note: [CWE-116] Improper Encoding or Escaping of Output.

(use-jsonify)

---

[info] 300-300: use jsonify instead of json.dumps for JSON output
Context: json.dumps({})
Note: [CWE-116] Improper Encoding or Escaping of Output.

(use-jsonify)

🔇 Additional comments (1)

tests-unit/comfy_quant/test_mixed_precision.py (1)

231-310: LGTM!

[chelsealong]

Addressed both inline review comments: _load_quantized_module now clears quant_format/layout_type when a module reloads an unquantized weight (avoids Linear.forward taking the stale quantized path against a plain Parameter), and the Embedding load path now raises ValueError when an inferred format is nvfp4 instead of silently loading raw quantized bytes as an unquantized weight. Added test_reload_unquantized_resets_stale_quant_state and test_formatless_scaled_comfy_quant_embedding_rejects_nvfp4; both fail on pre-fix code. Full suite: python -m pytest tests-unit/comfy_quant -q → 12 passed. ruff check comfy/ops.py tests-unit/comfy_quant/test_mixed_precision.py passes.

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

comfy/ops.py (1)

1605-1620: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Reset stale embedding quantization state.

At Line 1610, an all-NUL marker now selects plain-weight loading. If this Embedding previously loaded quantized data, self.quant_format and self.layout_type still describe the old quantized layout. Reset both attributes when layer_conf is None. Add a reload regression test for this path.

Proposed fix

                 if layer_conf is not None:
                     raw_conf = layer_conf.numpy().tobytes()
                     layer_conf = json.loads(raw_conf) if raw_conf.strip(b"\x00") else None
 
+                if layer_conf is None:
+                    self.quant_format = None
+                    self.layout_type = None
+
                 quant_format = layer_conf.get("format") if layer_conf is not None else None

As per path instructions, “Clear stale quantization state when transitioning to unquantized weights.”

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@comfy/ops.py` around lines 1605 - 1620, When the embedding quantization
metadata is absent, including the all-NUL marker path in the surrounding loading
logic, reset self.quant_format and self.layout_type to their unquantized
defaults before continuing plain-weight loading. Add a regression test that
loads quantized data into the same Embedding, reloads it with an all-NUL or
missing comfy_quant marker, and verifies both attributes no longer retain the
stale quantized state.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@comfy/ops.py`:
- Around line 1605-1620: When the embedding quantization metadata is absent,
including the all-NUL marker path in the surrounding loading logic, reset
self.quant_format and self.layout_type to their unquantized defaults before
continuing plain-weight loading. Add a regression test that loads quantized data
into the same Embedding, reloads it with an all-NUL or missing comfy_quant
marker, and verifies both attributes no longer retain the stale quantized state.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: 2f431195-2019-42e9-819f-63e398312407

📥 Commits

Reviewing files that changed from the base of the PR and between 6fc556d and 9601f47.

📒 Files selected for processing (2)

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

📜 Review details

🧰 Additional context used

📓 Path-based instructions (6)

**/*

📄 CodeRabbit inference engine (AGENTS.md)

**/*: Keep changes small, direct, and limited to the narrowest necessary code path and smallest number of files.
Prefer practical fixes, minimal dependencies, and existing repository patterns; remove obsolete, dead, unreachable, or unused code.
Preserve existing APIs, node names, model-loading behavior, file layout, and workflow compatibility unless replacement is explicitly intended.
Core ComfyUI must not add outbound internet requests, telemetry, tracking, reporting, remote configuration, or background network activity. User-authorized model downloads are limited to the requested artifact and must exclude telemetry and unrelated metadata.

Files:

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: Keep state and capability flags on the object that owns the behavior. Prefer explicit parent-owned attributes over probing child objects with getattr; use child checks only when the child owns the delegated behavior.
Preserve shared method signatures, argument order, return shapes, side effects, and error behavior unless every affected caller and interface is intentionally updated.
Do not add unused compatibility parameters, flags, attributes, constructor options, or model-specific options to shared helpers; keep one-off behavior at the integration boundary.
Normalize third-party return conventions at integration boundaries so core code receives the expected type and shape; avoid undocumented caller-side unwrapping.
Do not add torch.no_grad, torch.inference_mode, or inference-mode wrappers. Do not add model freeze/unfreeze toggles; only disable globally enabled inference mode when a training path requires gradients.
Remove inference-only training behavior such as dropout while preserving checkpoint and state-dict compatibility; use nn.Identity when deleting a module would alter keys or ordering.
Keep imports at module scope except established optional-backend probes or imports required to avoid cycles; avoid unnecessary try/except blocks and use specific exceptions with useful fallbacks.
Do not add workarounds for unsupported library versions, especially PyTorch exception-and-float-cast retries, unless a comment names the exact versions still requiring them.
Let unsupported model formats, invalid quantization metadata, and bad states fail with clear errors instead of silently degrading output.
Match local style, keep comments sparse and useful, and remove comments that merely restate obvious code.
Treat dtype, device placement, VRAM use, and offloading as correctness concerns across CPU, CUDA, ROCm, MPS, DirectML, XPU, NPU, and low-VRAM environments.
Prefer existing ComfyUI and Comfy Kitchen operations, quantization helpers, cast/offload helpe...

Files:

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

**/*.{py,json}

📄 CodeRabbit inference engine (AGENTS.md)

Treat legacy combo, io.Combo, and io.DynamicCombo values affecting filesystem access as untrusted; revalidate them at load/save boundaries with folder_paths, containment checks, or fixed allowlists.

Files:

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

**/*.{py,md,txt,json}

📄 CodeRabbit inference engine (AGENTS.md)

Keep warning and info messages short and actionable, remove noisy or misleading logging, and make documentation edits concise, factual, and tied to changed behavior.

Files:

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

**

⚙️ CodeRabbit configuration file

**: IMPORTANT: Only comment on issues directly introduced by this PR's code changes.
Treat AGENTS.md as mandatory repository policy, not optional style guidance.
Flag PR changes that violate AGENTS.md even when the code is otherwise functional.
In particular, enforce architecture boundaries, dtype/device/memory rules,
interface contracts, import style, no unnecessary try/except blocks, no inline
imports, no outbound internet paths in core ComfyUI, and narrow scoped fixes.
Prefer direct findings over suggestions when a rule is violated. Only ignore
AGENTS.md when it clearly conflicts with a newer explicit maintainer instruction
in the PR.
Do NOT flag pre-existing issues in code that was merely moved, re-indented,
de-indented, or reformatted without logic changes. If code appears in the diff
only due to whitespace or structural reformatting (e.g., removing a with: block),
treat it as unchanged. Contributors should not feel obligated to address
pre-existing issues outside the scope of their contribution.

Files:

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

comfy/**

⚙️ CodeRabbit configuration file

comfy/**: Core ML/diffusion engine. Focus on:

• Backward compatibility (breaking changes affect all custom nodes)
• Memory management and GPU resource handling
• Performance implications in hot paths
• Thread safety for concurrent execution

Files:

• comfy/ops.py

🧠 Learnings (3)

📚 Learning: 2026-02-21T14:01:41.482Z

Learnt from: pythongosssss
Repo: Comfy-Org/ComfyUI PR: 12555
File: comfy_extras/nodes_glsl.py:719-724
Timestamp: 2026-02-21T14:01:41.482Z
Learning: In PyOpenGL, bare Python scalars can be accepted for 1-element array parameters by NumberHandler. This means you can pass an int/float directly to OpenGL texture deletion (e.g., glDeleteTextures(tex)) without wrapping in a list. Verify function-specific expectations and ensure types match what the OpenGL call expects; use explicit lists only when the API requires an array.

Applied to files:

• comfy/ops.py
• tests-unit/comfy_quant/test_mixed_precision.py

📚 Learning: 2026-05-13T12:31:45.069Z

Learnt from: rattus128
Repo: Comfy-Org/ComfyUI PR: 13802
File: comfy/pinned_memory.py:19-30
Timestamp: 2026-05-13T12:31:45.069Z
Learning: When reviewing code that uses comfy/pinned_memory.py’s `HostBuffer.extend(size=..., reallocate=...)`: by default (`reallocate` is not True / False), `extend(size=...)` is a *relative increment* that grows the buffer by `size` bytes—so slicing like `[offset:offset+size]` after `hostbuf.extend(size=size)` is correct and the argument should not be rewritten to `offset + size`. Only in the single-segment reallocation mode (`reallocate=True`, e.g., as used by `resize_pin_buffer()` in `comfy/model_management.py`) should `size` be treated as an *absolute target* and the call/arguments should be checked accordingly.

Applied to files:

• comfy/ops.py

📚 Learning: 2026-08-06T22:18:59.719Z

Learnt from: kijai
Repo: Comfy-Org/ComfyUI PR: 15362
File: comfy/ldm/wan/model_animate2.py:186-223
Timestamp: 2026-08-06T22:18:59.719Z
Learning: When reviewing ComfyUI quantization code, treat `comfy.quant_ops.TensorWiseINT8Layout` and `comfy.quant_ops.TensorCoreConvRotW4A4Layout` as re-exports from `comfy_kitchen`. Validate their behavior against the re-exported `comfy_kitchen` implementations rather than assuming they are local fallback classes.

Applied to files:

• comfy/ops.py

🪛 ast-grep (0.45.0)

tests-unit/comfy_quant/test_mixed_precision.py

[info] 331-331: use jsonify instead of json.dumps for JSON output
Context: json.dumps({"layers": layer_quant_config})
Note: [CWE-116] Improper Encoding or Escaping of Output.

(use-jsonify)

---

[info] 374-374: use jsonify instead of json.dumps for JSON output
Context: json.dumps({})
Note: [CWE-116] Improper Encoding or Escaping of Output.

(use-jsonify)

🔇 Additional comments (1)

tests-unit/comfy_quant/test_mixed_precision.py (1)

311-359: LGTM!

Also applies to: 361-381