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 <amp@ampcode.com>

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.register_load_device(self.load_device)
-            self.model.dynamic_pins[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

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)