Skip to content

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

Description

@Booyaka101

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

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions