ryan/ComfyUI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Extracted multigpu core code into multigpu.py, added load_balance_devices to get subdivision of work based on available devices and splittable work item count, added MultiGPU Options nodes to set relative_speed of specific devices; does not change behavior yet

Browse Source
This commit is contained in:
Jedrzej Kosinski
2025-01-27 06:25:48 -06:00
parent e3298b84de
commit eda866bf51
2 changed files with 113 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

58
comfy_extras/nodes_multigpu.py
View File

@@ -1,10 +1,10 @@
from __future__ import annotations
import torch
from comfy.model_patcher import ModelPatcher
import comfy.utils
import comfy.patcher_extension
import comfy.model_management
import comfy.multigpu
class MultiGPUInitialize:
@@ -26,7 +26,7 @@ class MultiGPUInitialize:
FUNCTION = "init_multigpu"
CATEGORY = "advanced/multigpu"
def init_multigpu(self, model: ModelPatcher, max_gpus: int, gpu_options: GPUOptionsGroup=None):
def init_multigpu(self, model: ModelPatcher, max_gpus: int, gpu_options: comfy.multigpu.GPUOptionsGroup=None):
extra_devices = comfy.model_management.get_all_torch_devices(exclude_current=True)
extra_devices = extra_devices[:max_gpus-1]
if len(extra_devices) > 0:
@@ -39,7 +39,7 @@ class MultiGPUInitialize:
multigpu_models.append(device_patcher)
model.set_additional_models("multigpu", multigpu_models)
if gpu_options is None:
gpu_options = GPUOptionsGroup()
gpu_options = comfy.multigpu.GPUOptionsGroup()
gpu_options.register(model)
return (model,)
@@ -62,63 +62,17 @@ class MultiGPUOptionsNode:
FUNCTION = "create_gpu_options"
CATEGORY = "advanced/multigpu"
def create_gpu_options(self, device_index: int, relative_speed: float, gpu_options: GPUOptionsGroup=None):
def create_gpu_options(self, device_index: int, relative_speed: float, gpu_options: comfy.multigpu.GPUOptionsGroup=None):
if not gpu_options:
gpu_options = GPUOptionsGroup()
gpu_options = comfy.multigpu.GPUOptionsGroup()
gpu_options.clone()
opt = GPUOptions(device_index=device_index, relative_speed=relative_speed)
opt = comfy.multigpu.GPUOptions(device_index=device_index, relative_speed=relative_speed)
gpu_options.add(opt)
return (gpu_options,)
class GPUOptions:
def __init__(self, device_index: int, relative_speed: float):
self.device_index = device_index
self.relative_speed = relative_speed
def clone(self):
return GPUOptions(self.device_index, self.relative_speed)
def create_dict(self):
return {
"relative_speed": self.relative_speed
}
class GPUOptionsGroup:
def __init__(self):
self.options: dict[int, GPUOptions] = {}
def add(self, info: GPUOptions):
self.options[info.device_index] = info
def clone(self):
c = GPUOptionsGroup()
for opt in self.options.values():
c.add(opt)
return c
def register(self, model: ModelPatcher):
opts_dict = {}
# get devices that are valid for this model
devices: list[torch.device] = [model.load_device]
for extra_model in model.get_additional_models_with_key("multigpu"):
extra_model: ModelPatcher
devices.append(extra_model.load_device)
# create dictionary with actual device mapped to its GPUOptions
device_opts_list: list[GPUOptions] = []
for device in devices:
device_opts = self.options.get(device.index, GPUOptions(device_index=device.index, relative_speed=1.0))
opts_dict[device] = device_opts.create_dict()
device_opts_list.append(device_opts)
# make relative_speed relative to 1.0
max_speed = max([x.relative_speed for x in device_opts_list])
for value in opts_dict.values():
value["relative_speed"] /= max_speed
model.model_options["multigpu_options"] = opts_dict
node_list = [
MultiGPUInitialize,
MultiGPUOptionsNode