From 25818dc848f8db6f79b4410e46b06133165d35a2 Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 14 Jan 2025 13:45:14 -0600 Subject: [PATCH] Added a 'max_gpus' input --- comfy_extras/nodes_multigpu.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/comfy_extras/nodes_multigpu.py b/comfy_extras/nodes_multigpu.py index 929151b5..3ba55862 100644 --- a/comfy_extras/nodes_multigpu.py +++ b/comfy_extras/nodes_multigpu.py @@ -11,6 +11,9 @@ class MultiGPUInitialize: return { "required": { "model": ("MODEL",), + }, + "optional": { + "max_gpus" : ("INT", {"default": 8, "min": 1, "step": 1}), } } @@ -18,9 +21,10 @@ class MultiGPUInitialize: FUNCTION = "init_multigpu" CATEGORY = "DevTools" - def init_multigpu(self, model: ModelPatcher): + def init_multigpu(self, model: ModelPatcher, max_gpus: int): model = model.clone() extra_devices = comfy.model_management.get_all_torch_devices(exclude_current=True) + extra_devices = extra_devices[:max_gpus-1] if len(extra_devices) > 0: comfy.model_management.unload_all_models() for device in extra_devices: