From 4e650055d0b7ea2eef4a0ac5091fadd5bcf553e1 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Fri, 22 May 2026 21:46:07 -0700 Subject: [PATCH] SelectXDevice nodes: register new load_device with ModelPatcherDynamic 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 --- comfy/model_patcher.py | 19 +++++++++++++++---- comfy_extras/nodes_multigpu.py | 6 ++++++ 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 087b0fbfa..2bb363fab 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -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 diff --git a/comfy_extras/nodes_multigpu.py b/comfy_extras/nodes_multigpu.py index 9e03c56f0..df701af56 100644 --- a/comfy_extras/nodes_multigpu.py +++ b/comfy_extras/nodes_multigpu.py @@ -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)