Comfy-Org / Comfy-Org/ComfyUI

Way to skip self.model.to(device_to) inside ModelPatcher.unpatch_model

Open
#14,142 3 comments 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 7h
Merged PRs (30d)
158

Description

### Feature Idea

`ModelPatcher.unpatch_model(unpatch_weights=True)` ends its weight-unpatch block with `self.model.to(device_to)` (master `comfy/model_patcher.py` around the `current_weight_patches_uuid = None` / `backup.clear()` lines). That walk uses `nn.Module.to` which iterates every parameter, ignoring `comfy_cast_weights = True` on custom layers.

For custom ops that intentionally keep weights on `offload_device` and only cast on the fly through `comfy_cast_weights`, the walk hits parameters that are not safe to `.to(...)` (mmap-backed, lazily quantized, etc.) and faults. ComfyUI-GGUF hits this with Windows `EXCEPTION_ACCESS_VIOLATION` on mmap-backed quantized tensors (city96/ComfyUI-GGUF#444).

The only existing escape hatch is `unpatch_weights=False`, but that block also runs `unpatch_hooks`, `unpin_all_weights`, lowvram cleanup, non-quantized backup restoration, `model_loaded_weight_memory = 0`, `model_offload_buffer_memory = 0`, and `comfy_patched_weights` deletion. Custom patchers that want to keep all of that but skip the device walk end up mirroring the whole block by hand. city96/ComfyUI-GGUF#445 is the latest iteration of that workaround there.

Two shapes that would work:

1. Add `unpatch_device: bool = True` to `unpatch_model`, gating just the `self.model.to(device_to)` line. Defaults preserve behavior.
2. Replace `self.model.to(device_to)` with a per-module walk that skips modules with `comfy_cast_weights = True` (their weights are managed by the op layer, not by `ModelPatcher`).

Either would let custom patchers drop their `unpatch_model` overrides entirely.

### Existing Solutions

- ComfyUI-GGUF currently maintains a multi-iteration downstream workaround in `nodes.py:GGUFModelPatcher.unpatch_model` (see history at `6dbb4ba` -> `717a0e1` -> #355 -> #357 -> #445). The maintainer flagged this as needing an upstream fix in `717a0e1`'s commit message in September 2024.
- Other quantization custom nodes (e.g. bitsandbytes-based loaders) avoid the issue because their weights survive `.to`, but they would still benefit from the cleaner contract.

### Other

The `nn.Module.to` walk is also redundant on the `partially_load` call path (`comfy/model_patcher.py` `partially_load` calls `unpatch_model(self.offload_device, ...)` and then immediately re-loads), so option (2) would also avoid wasted work on that hot path. The `detach` path is the only caller that genuinely needs the weights on `offload_device` after unpatch.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.