get_key_weight() crashes with AttributeError when hooks are used on a mixed-precision quantized model
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
## Summary
Applying a hook-based model patch (e.g. via `comfy.hooks.create_hook_lora` + masked/regional conditioning) against a model that uses mixed-precision quantization — where some layers are quantized (carry a `weight_scale` tensor) and others are plain `nn.Linear` — crashes with an unguarded `AttributeError` inside `ModelPatcher.get_key_patches()`.
## Environment
- ComfyUI version: 0.34.0
- Model: Krea 2 (FLUX-architecture) loaded with mixed-precision quantization (`Using MixedPrecisionOps for text encoder` / `Detected mixed precision quantization` / `Native ops: nvfp4, int8_tensorwise, convrot_w4a4, asym_w4a8_int8, mxfp8, float8_e5m2, float8_e4m3fn`)
- GPU: NVIDIA GeForce RTX 5090, `cudaMallocAsync`
- PyTorch: 2.10.0+cu130, Python 3.13.12
## Steps to reproduce
1. Load a checkpoint that uses mixed-precision quantization, such that its `model.state_dict()` includes `*.weight_scale` keys for *some* layers but not others (i.e. not every layer is uniformly quantized).
2. Build a `HookGroup` via `comfy.hooks.create_hook_lora(...)` and attach it to a conditioning (e.g. via `comfy.hooks.set_conds_props(..., hooks=hook_group)`), as the regional/masked-LoRA and hook-scheduling nodes do.
3. Sample. On the first step, the sampler calls `model.current_patcher.apply_hooks(hooks=hooks)` → `patch_hooks()`.
## Actual behavior
```
AttributeError: 'Linear' object has no attribute 'weight_scale'
File "comfy/model_patcher.py", line 1617, in patch_hooks
original_weights = self.get_key_patches()
File "comfy/model_patcher.py", line 875, in get_key_patches
weight, set_func, convert_func = get_key_weight(self.model, k)
File "comfy/model_patcher.py", line 234, in get_key_weight
weight = getattr(op, op_keys[1])
File "torch/nn/modules/module.py", line 1965, in __getattr__
raise AttributeError(...)
```
## Root cause
`patch_hooks()` calls `self.get_key_patches()` (model_patcher.py:1617) to snapshot **every** weight in the model, not just the keys the active hook actually targets. `get_key_patches()` (model_patcher.py:869-887) iterates every key returned by `model.state_dict()` and calls `get_key_weight(self.model, k)` for each.
`get_key_weight()` (model_patcher.py:216-238) splits the key into a module path and an attribute name, then does:
```python
weight = getattr(op, op_keys[1])
```
with no guard — unlike the two lookups immediately above it (`set_` / `convert_`), which are both wrapped in `try/except AttributeError`. For a mixed-precision model, `model.state_dict()` can report a `*.weight_scale` key for a module whose *live* instance is a plain `nn.Linear` with no such attribute (i.e. `state_dict()` and the actual live attribute set disagree for that layer), and the unguarded `getattr` throws.
This snapshot is only used as optional side-context in `comfy.lora.calculate_weight(..., original_weights=...)` for the hook's *actual* patch keys, which are separately validated (`patch_hooks()` line 1618-1620 checks `if key not in model_sd_keys: continue` before ever using `original_weights[key]`). So a missing/`None` entry for an unrelated phantom key is inert — it doesn't need to hard-fail the whole snapshot.
## Suggested fix
Make the lookup fail soft, matching the existing pattern for `set_func`/`convert_func` two lines above:
```diff
- weight = getattr(op, op_keys[1])
+ try:
+ weight = getattr(op, op_keys[1])
+ except AttributeError:
+ weight = None
if convert_func is not None:
weight = comfy.utils.get_attr(model, key)
```
Verified locally: this resolves the crash and sampling proceeds normally. Checked the other callers of `get_key_patches()` (the model-merging nodes in `comfy_extras/nodes_model_merging.py`) — none of them are affected differently by this change; a model without this quirk never hits the new branch at all.
Contributor guide
Research direction
Start in comfy/model_patcher.py at get_key_weight() and trace its callers through get_key_patches() and patch_hooks(). Reproduce the failure with a mixed-precision model using hooks, then verify sampling completes without the AttributeError and that the model-merging callers in comfy_extras/nodes_model_merging.py remain unaffected.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 88/100