Add Select Model/CLIP/VAE Device passthrough nodes

Replace the per-loader device widgets removed in the previous commit
with three small passthrough selector nodes registered under
advanced/multigpu:

- Select Model Device  (MODEL  in/out)  - options: default / cpu / gpu:N
- Select CLIP Device   (CLIP   in/out)  - options: default / cpu / gpu:N
- Select VAE Device    (VAE    in/out)  - options: default / gpu:N (no cpu)

Each node clones the inbound patcher (model.clone() / clip.clone() /
copy.copy(vae)+vae.patcher.clone()) and retargets load_device (and
offload_device for cpu / vae_offload_device for VAE).

Portability across machines with different GPU counts:
- VALIDATE_INPUTS returns True so an unknown gpu:N value (e.g. a
  workflow saved on a 2-GPU machine opened on a 1-GPU machine) does
  not error at validation time.
- At runtime, resolve_gpu_device_option(...) returns None for
  unknown options (with a warning), and each selector then logs a
  per-node info message and passes through unchanged, matching the
  no-op style used by MultiGPU CFG Split's
  "No extra torch devices need initialization..." log.

Also adds comfy.model_management.get_gpu_device_options_no_cpu() which
the VAE selector uses; on a single-GPU box this collapses to just
["default"], which is fine.

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:39:18 -07:00
parent 9a12a9328b
commit d7706091ae
2 changed files with 157 additions and 0 deletions

View File

@ -255,6 +255,14 @@ def get_gpu_device_options():
options.append(f"gpu:{i}") options.append(f"gpu:{i}")
return options return options
def get_gpu_device_options_no_cpu():
"""Variant of get_gpu_device_options that omits "cpu".
Intended for components like the VAE selector where running on CPU
is impractical and should not be offered as a choice.
"""
return [o for o in get_gpu_device_options() if o != "cpu"]
def resolve_gpu_device_option(option: str): def resolve_gpu_device_option(option: str):
"""Resolve a device option string to a torch.device. """Resolve a device option string to a torch.device.

View File

@ -1,5 +1,7 @@
from __future__ import annotations from __future__ import annotations
import copy
import logging
from inspect import cleandoc from inspect import cleandoc
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from typing_extensions import override from typing_extensions import override
@ -8,6 +10,8 @@ from comfy_api.latest import ComfyExtension, io
if TYPE_CHECKING: if TYPE_CHECKING:
from comfy.model_patcher import ModelPatcher from comfy.model_patcher import ModelPatcher
from comfy.sd import CLIP, VAE
import comfy.model_management
import comfy.multigpu import comfy.multigpu
@ -42,6 +46,148 @@ class MultiGPUCFGSplitNode(io.ComfyNode):
return io.NodeOutput(model) return io.NodeOutput(model)
class SelectModelDeviceNode(io.ComfyNode):
"""
Place the diffusion model on a specific device (default / cpu / gpu:N).
When the selected device does not exist on the current machine
(e.g. a workflow built on a 2-GPU box opened on a 1-GPU box),
the node passes the model through unchanged and logs a message
instead of failing. This keeps workflows portable across machines
with different GPU counts.
"""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="SelectModelDevice",
display_name="Select Model Device",
category="advanced/multigpu",
description=cleandoc(cls.__doc__),
inputs=[
io.Model.Input("model"),
io.Combo.Input("device", options=comfy.model_management.get_gpu_device_options()),
],
outputs=[
io.Model.Output(),
],
)
@classmethod
def VALIDATE_INPUTS(cls, device="default"):
# Allow unknown gpu:N values so portable workflows do not error
# at validation time; runtime fallback will handle them.
return True
@classmethod
def execute(cls, model: ModelPatcher, device: str = "default") -> io.NodeOutput:
model = model.clone()
resolved = comfy.model_management.resolve_gpu_device_option(device)
if resolved is None:
if device not in (None, "default"):
logging.info(f"Select Model Device: requested device '{device}' not available, passing through unchanged.")
return io.NodeOutput(model)
model.load_device = resolved
if resolved.type == "cpu":
model.offload_device = resolved
return io.NodeOutput(model)
class SelectCLIPDeviceNode(io.ComfyNode):
"""
Place the CLIP text encoder on a specific device (default / cpu / gpu:N).
When the selected device does not exist on the current machine
(e.g. a workflow built on a 2-GPU box opened on a 1-GPU box),
the node passes the CLIP through unchanged and logs a message
instead of failing. This keeps workflows portable across machines
with different GPU counts.
"""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="SelectCLIPDevice",
display_name="Select CLIP Device",
category="advanced/multigpu",
description=cleandoc(cls.__doc__),
inputs=[
io.Clip.Input("clip"),
io.Combo.Input("device", options=comfy.model_management.get_gpu_device_options()),
],
outputs=[
io.Clip.Output(),
],
)
@classmethod
def VALIDATE_INPUTS(cls, device="default"):
return True
@classmethod
def execute(cls, clip: CLIP, device: str = "default") -> io.NodeOutput:
clip = clip.clone()
resolved = comfy.model_management.resolve_gpu_device_option(device)
if resolved is None:
if device not in (None, "default"):
logging.info(f"Select CLIP Device: requested device '{device}' not available, passing through unchanged.")
return io.NodeOutput(clip)
clip.patcher.load_device = resolved
if resolved.type == "cpu":
clip.patcher.offload_device = resolved
return io.NodeOutput(clip)
class SelectVAEDeviceNode(io.ComfyNode):
"""
Place the VAE on a specific device (default / gpu:N).
CPU is intentionally not offered as a choice; VAE on CPU is impractical.
When the selected device does not exist on the current machine
(e.g. a workflow built on a 2-GPU box opened on a 1-GPU box),
the node passes the VAE through unchanged and logs a message
instead of failing. This keeps workflows portable across machines
with different GPU counts.
"""
@classmethod
def define_schema(cls):
return io.Schema(
node_id="SelectVAEDevice",
display_name="Select VAE Device",
category="advanced/multigpu",
description=cleandoc(cls.__doc__),
inputs=[
io.Vae.Input("vae"),
io.Combo.Input("device", options=comfy.model_management.get_gpu_device_options_no_cpu()),
],
outputs=[
io.Vae.Output(),
],
)
@classmethod
def VALIDATE_INPUTS(cls, device="default"):
return True
@classmethod
def execute(cls, vae: VAE, device: str = "default") -> io.NodeOutput:
# VAE has no .clone(); shallow-copy the wrapper and clone the patcher
# so we can retarget load/offload device without affecting the input VAE.
vae = copy.copy(vae)
vae.patcher = vae.patcher.clone()
resolved = comfy.model_management.resolve_gpu_device_option(device)
if resolved is None:
if device not in (None, "default"):
logging.info(f"Select VAE Device: requested device '{device}' not available, passing through unchanged.")
return io.NodeOutput(vae)
vae.device = resolved
vae.patcher.load_device = resolved
vae.patcher.offload_device = comfy.model_management.vae_offload_device()
return io.NodeOutput(vae)
class MultiGPUOptionsNode(io.ComfyNode): class MultiGPUOptionsNode(io.ComfyNode):
""" """
Select the relative speed of GPUs in the special case they have significantly different performance from one another. Select the relative speed of GPUs in the special case they have significantly different performance from one another.
@ -92,6 +238,9 @@ class MultiGPUExtension(ComfyExtension):
async def get_node_list(self) -> list[type[io.ComfyNode]]: async def get_node_list(self) -> list[type[io.ComfyNode]]:
return [ return [
MultiGPUCFGSplitNode, MultiGPUCFGSplitNode,
SelectModelDeviceNode,
SelectCLIPDeviceNode,
SelectVAEDeviceNode,
# MultiGPUOptionsNode, # MultiGPUOptionsNode,
] ]