wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-05-24 07:57:29 +08:00
Code Issues Actions 16 Packages Projects Releases Wiki Activity

SelectXDevice nodes: register new load_device with ModelPatcherDynamic

Browse Source 
When --enable-dynamic-vram is on, every ModelPatcher is a
ModelPatcherDynamic whose underlying model has a per-device dynamic_pins
dict, initialized in __init__ for self.load_device only. If a cloned
patcher's load_device is later reassigned (as the Select{Model,CLIP,VAE}
Device nodes do), the new device key is missing and partially_unload_ram
raises KeyError: device(type='cuda', index=N).

Fix:
- Extract the per-device dynamic_pins init in ModelPatcherDynamic.__init__
  into a new helper method register_load_device(device) which is now also
  called from __init__.
- Each Select*Device node calls clone.patcher.register_load_device(resolved)
  after retargeting load_device, guarded by hasattr so non-dynamic
  patchers (plain ModelPatcher in non-dynamic-vram installs) skip it.

Caught by happy-path test where SelectCLIPDevice retargeted CLIP from
cuda:0 to cuda:1 and CLIPTextEncode then crashed in
partially_unload_ram -> dynamic_pins[cuda:1].

Amp-Thread-ID: https://ampcode.com/threads/T-019e52b4-31ee-72cd-996b-64ecd9420e13
Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
Jedrzej Kosinski 2026-05-22 21:46:07 -07:00
parent d7706091ae
commit 4e650055d0
2 changed files with 21 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
comfy/model_patcher.py
View File

@ -1692,16 +1692,27 @@ class ModelPatcherDynamic(ModelPatcher):
self.model.dynamic_vbars = {}
if not hasattr(self.model, "dynamic_pins"):
self.model.dynamic_pins = {}
if self.load_device not in self.model.dynamic_pins:
self.model.dynamic_pins[self.load_device] = {
self.register_load_device(self.load_device)
self.non_dynamic_delegate_model = None
assert load_device is not None
def register_load_device(self, device):
"""Ensure dynamic_pins has an entry for *device*.
Called from __init__ and also from any code that retargets an
already-constructed patcher to a new load_device (e.g. the
Select{Model,CLIP,VAE}Device selector nodes); without this entry
partially_unload_ram() raises KeyError when it tries to read the
per-device pin state.
"""
if device not in self.model.dynamic_pins:
self.model.dynamic_pins[device] = {
"weights": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0]),
"patches": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0]),
"hostbufs_initialized": False,
"failed": False,
"active": False,
}
self.non_dynamic_delegate_model = None
assert load_device is not None
def is_dynamic(self):
return True

6
comfy_extras/nodes_multigpu.py
View File

@ -90,6 +90,8 @@ class SelectModelDeviceNode(io.ComfyNode):
model.load_device = resolved
if resolved.type == "cpu":
model.offload_device = resolved
if hasattr(model, "register_load_device"):
model.register_load_device(resolved)
return io.NodeOutput(model)
@ -135,6 +137,8 @@ class SelectCLIPDeviceNode(io.ComfyNode):
clip.patcher.load_device = resolved
if resolved.type == "cpu":
clip.patcher.offload_device = resolved
if hasattr(clip.patcher, "register_load_device"):
clip.patcher.register_load_device(resolved)
return io.NodeOutput(clip)
@ -185,6 +189,8 @@ class SelectVAEDeviceNode(io.ComfyNode):
vae.device = resolved
vae.patcher.load_device = resolved
vae.patcher.offload_device = comfy.model_management.vae_offload_device()
if hasattr(vae.patcher, "register_load_device"):
vae.patcher.register_load_device(resolved)
return io.NodeOutput(vae)