298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-04-17 09:33:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added a 'max_gpus' input

Browse Source
This commit is contained in:
Jedrzej Kosinski 2025-01-14 13:45:14 -06:00
parent 198953cd08
commit 25818dc848
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
comfy_extras/nodes_multigpu.py
View File

@ -11,6 +11,9 @@ class MultiGPUInitialize:
return { return {
"required": { "required": {
"model": ("MODEL",), "model": ("MODEL",),
},
"optional": {
"max_gpus" : ("INT", {"default": 8, "min": 1, "step": 1}),
} }
} }
@ -18,9 +21,10 @@ class MultiGPUInitialize:
FUNCTION = "init_multigpu" FUNCTION = "init_multigpu"
CATEGORY = "DevTools" CATEGORY = "DevTools"
def init_multigpu(self, model: ModelPatcher): def init_multigpu(self, model: ModelPatcher, max_gpus: int):
model = model.clone() model = model.clone()
extra_devices = comfy.model_management.get_all_torch_devices(exclude_current=True) extra_devices = comfy.model_management.get_all_torch_devices(exclude_current=True)
extra_devices = extra_devices[:max_gpus-1]
if len(extra_devices) > 0: if len(extra_devices) > 0:
comfy.model_management.unload_all_models() comfy.model_management.unload_all_models()
for device in extra_devices: for device in extra_devices: