437 lines
18 KiB
Python
437 lines
18 KiB
Python
import copy
|
|
import torch
|
|
import sys
|
|
import comfy.model_management as mm
|
|
import os
|
|
from pathlib import Path
|
|
import logging
|
|
import folder_paths
|
|
from collections import defaultdict
|
|
import hashlib
|
|
|
|
current_device = mm.get_torch_device()
|
|
current_offload_device = mm.get_torch_device()
|
|
model_allocation_store = {}
|
|
|
|
def get_torch_device_patched():
|
|
device = None
|
|
if (not torch.cuda.is_available() or mm.cpu_state == mm.CPUState.CPU or "cpu" in str(current_device).lower()):
|
|
device = torch.device("cpu")
|
|
else:
|
|
device = torch.device(current_device)
|
|
return device
|
|
|
|
def text_encoder_device_patched():
|
|
device = None
|
|
if (not torch.cuda.is_available() or mm.cpu_state == mm.CPUState.CPU or "cpu" in str(current_device).lower()):
|
|
device = torch.device("cpu")
|
|
else:
|
|
device = torch.device(current_device)
|
|
return device
|
|
|
|
def unet_offload_device_patched():
|
|
device = None
|
|
if (not torch.cuda.is_available() or mm.cpu_state == mm.CPUState.CPU or "cpu" in str(current_offload_device).lower()):
|
|
device = torch.device("cpu")
|
|
else:
|
|
device = torch.device(current_offload_device)
|
|
return device
|
|
|
|
def text_encoder_offload_device_patched():
|
|
device = None
|
|
if (not torch.cuda.is_available() or mm.cpu_state == mm.CPUState.CPU or "cpu" in str(current_offload_device).lower()):
|
|
device = torch.device("cpu")
|
|
else:
|
|
device = torch.device(current_offload_device)
|
|
return device
|
|
|
|
mm.get_torch_device = get_torch_device_patched
|
|
mm.unet_offload_device = unet_offload_device_patched
|
|
mm.text_encoder_device = text_encoder_device_patched
|
|
mm.text_encoder_offload_device = text_encoder_offload_device_patched
|
|
|
|
|
|
def create_model_hash(model, caller):
|
|
|
|
model_type = type(model.model).__name__
|
|
model_size = model.model_size()
|
|
first_layers = str(list(model.model_state_dict().keys())[:3])
|
|
identifier = f"{model_type}_{model_size}_{first_layers}"
|
|
final_hash = hashlib.sha256(identifier.encode()).hexdigest()
|
|
|
|
return final_hash
|
|
|
|
def register_patched_ggufmodelpatcher():
|
|
from nodes import NODE_CLASS_MAPPINGS
|
|
original_loader = NODE_CLASS_MAPPINGS["UnetLoaderGGUF"]
|
|
module = sys.modules[original_loader.__module__]
|
|
|
|
if not hasattr(module.GGUFModelPatcher, '_patched'):
|
|
original_load = module.GGUFModelPatcher.load
|
|
|
|
def new_load(self, *args, force_patch_weights=False, **kwargs):
|
|
global model_allocation_store
|
|
|
|
super(module.GGUFModelPatcher, self).load(*args, force_patch_weights=True, **kwargs)
|
|
debug_hash = create_model_hash(self, "patcher")
|
|
linked = []
|
|
module_count = 0
|
|
for n, m in self.model.named_modules():
|
|
module_count += 1
|
|
if hasattr(m, "weight"):
|
|
device = getattr(m.weight, "device", None)
|
|
if device is not None:
|
|
linked.append((n, m))
|
|
continue
|
|
if hasattr(m, "bias"):
|
|
device = getattr(m.bias, "device", None)
|
|
if device is not None:
|
|
linked.append((n, m))
|
|
continue
|
|
if linked:
|
|
if hasattr(self, 'model'):
|
|
debug_hash = create_model_hash(self, "patcher")
|
|
debug_allocations = model_allocation_store.get(debug_hash)
|
|
if debug_allocations:
|
|
device_assignments = analyze_ggml_loading(self.model, debug_allocations)['device_assignments']
|
|
for device, layers in device_assignments.items():
|
|
target_device = torch.device(device)
|
|
for n, m, _ in layers:
|
|
m.to(self.load_device).to(target_device)
|
|
|
|
self.mmap_released = True
|
|
|
|
module.GGUFModelPatcher.load = new_load
|
|
module.GGUFModelPatcher._patched = True
|
|
|
|
def analyze_ggml_loading(model, allocations_str):
|
|
DEVICE_RATIOS_DISTORCH = {}
|
|
device_table = {}
|
|
|
|
for allocation in allocations_str.split(';'):
|
|
dev_name, fraction = allocation.split(',')
|
|
fraction = float(fraction)
|
|
total_mem_bytes = mm.get_total_memory(torch.device(dev_name))
|
|
alloc_gb = (total_mem_bytes * fraction) / (1024**3)
|
|
DEVICE_RATIOS_DISTORCH[dev_name] = alloc_gb
|
|
device_table[dev_name] = {
|
|
"fraction": fraction,
|
|
"total_gb": total_mem_bytes / (1024**3),
|
|
"alloc_gb": alloc_gb
|
|
}
|
|
|
|
eq_line = "=" * 47
|
|
dash_line = "-" * 47
|
|
fmt_alloc = "{:<12}{:>10}{:>14}{:>10}"
|
|
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
|
logging.info(eq_line)
|
|
logging.info(" DisTorch Analysis")
|
|
logging.info(eq_line)
|
|
logging.info(dash_line)
|
|
logging.info(" DisTorch Device Allocations")
|
|
logging.info(dash_line)
|
|
logging.info(fmt_alloc.format("Device", "Alloc %", "Total (GB)", " Alloc (GB)"))
|
|
logging.info(dash_line)
|
|
|
|
sorted_devices = sorted(device_table.keys(), key=lambda d: (d == "cpu", d))
|
|
|
|
for dev in sorted_devices:
|
|
frac = device_table[dev]["fraction"]
|
|
tot_gb = device_table[dev]["total_gb"]
|
|
alloc_gb = device_table[dev]["alloc_gb"]
|
|
logging.info(fmt_alloc.format(dev,f"{int(frac * 100)}%",f"{tot_gb:.2f}",f"{alloc_gb:.2f}"))
|
|
|
|
logging.info(dash_line)
|
|
|
|
layer_summary = {}
|
|
layer_list = []
|
|
memory_by_type = defaultdict(int)
|
|
total_memory = 0
|
|
|
|
for name, module in model.named_modules():
|
|
if hasattr(module, "weight"):
|
|
layer_type = type(module).__name__
|
|
layer_summary[layer_type] = layer_summary.get(layer_type, 0) + 1
|
|
layer_list.append((name, module, layer_type))
|
|
layer_memory = 0
|
|
if module.weight is not None:
|
|
layer_memory += module.weight.numel() * module.weight.element_size()
|
|
if hasattr(module, "bias") and module.bias is not None:
|
|
layer_memory += module.bias.numel() * module.bias.element_size()
|
|
memory_by_type[layer_type] += layer_memory
|
|
total_memory += layer_memory
|
|
|
|
logging.info(" DisTorch GGML Layer Distribution")
|
|
logging.info(dash_line)
|
|
fmt_layer = "{:<12}{:>10}{:>14}{:>10}"
|
|
logging.info(fmt_layer.format("Layer Type", "Layers", "Memory (MB)", "% Total"))
|
|
logging.info(dash_line)
|
|
for layer_type, count in layer_summary.items():
|
|
mem_mb = memory_by_type[layer_type] / (1024 * 1024)
|
|
mem_percent = (memory_by_type[layer_type] / total_memory) * 100 if total_memory > 0 else 0
|
|
logging.info(fmt_layer.format(layer_type,str(count),f"{mem_mb:.2f}",f"{mem_percent:.1f}%"))
|
|
logging.info(dash_line)
|
|
|
|
nonzero_devices = [d for d, r in DEVICE_RATIOS_DISTORCH.items() if r > 0]
|
|
nonzero_total_ratio = sum(DEVICE_RATIOS_DISTORCH[d] for d in nonzero_devices)
|
|
device_assignments = {device: [] for device in DEVICE_RATIOS_DISTORCH.keys()}
|
|
total_layers = len(layer_list)
|
|
current_layer = 0
|
|
|
|
for idx, device in enumerate(nonzero_devices):
|
|
ratio = DEVICE_RATIOS_DISTORCH[device]
|
|
if idx == len(nonzero_devices) - 1:
|
|
device_layer_count = total_layers - current_layer
|
|
else:
|
|
device_layer_count = int((ratio / nonzero_total_ratio) * total_layers)
|
|
start_idx = current_layer
|
|
end_idx = current_layer + device_layer_count
|
|
device_assignments[device] = layer_list[start_idx:end_idx]
|
|
current_layer += device_layer_count
|
|
|
|
logging.info(" DisTorch Final Device/Layer Assignments")
|
|
logging.info(dash_line)
|
|
fmt_assign = "{:<12}{:>10}{:>14}{:>10}"
|
|
logging.info(fmt_assign.format("Device", "Layers", "Memory (MB)", "% Total"))
|
|
logging.info(dash_line)
|
|
total_assigned_memory = 0
|
|
device_memories = {}
|
|
for device, layers in device_assignments.items():
|
|
device_memory = 0
|
|
for layer_type in layer_summary:
|
|
type_layers = sum(1 for _, _, lt in layers if lt == layer_type)
|
|
if layer_summary[layer_type] > 0:
|
|
mem_per_layer = memory_by_type[layer_type] / layer_summary[layer_type]
|
|
device_memory += mem_per_layer * type_layers
|
|
device_memories[device] = device_memory
|
|
total_assigned_memory += device_memory
|
|
|
|
sorted_assignments = sorted(device_assignments.keys(), key=lambda d: (d == "cpu", d))
|
|
|
|
for dev in sorted_assignments:
|
|
layers = device_assignments[dev]
|
|
mem_mb = device_memories[dev] / (1024 * 1024)
|
|
mem_percent = (device_memories[dev] / total_memory) * 100 if total_memory > 0 else 0
|
|
logging.info(fmt_assign.format(dev,str(len(layers)),f"{mem_mb:.2f}",f"{mem_percent:.1f}%"))
|
|
logging.info(dash_line)
|
|
|
|
return {"device_assignments": device_assignments}
|
|
|
|
|
|
def get_device_list():
|
|
import torch
|
|
return ["cpu"] + [f"cuda:{i}" for i in range(torch.cuda.device_count())]
|
|
|
|
class DeviceSelectorMultiGPU:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
devices = get_device_list()
|
|
return {
|
|
"required": {
|
|
"device": (devices, {"default": devices[1] if len(devices) > 1 else devices[0]})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = (get_device_list(),)
|
|
RETURN_NAMES = ("device",)
|
|
FUNCTION = "select_device"
|
|
CATEGORY = "multigpu"
|
|
|
|
def select_device(self, device):
|
|
return (device,)
|
|
|
|
class HunyuanVideoEmbeddingsAdapter:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"hyvid_embeds": ("HYVIDEMBEDS",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING",)
|
|
FUNCTION = "adapt_embeddings"
|
|
CATEGORY = "multigpu"
|
|
|
|
def adapt_embeddings(self, hyvid_embeds):
|
|
cond = hyvid_embeds["prompt_embeds"]
|
|
|
|
pooled_dict = {
|
|
"pooled_output": hyvid_embeds["prompt_embeds_2"],
|
|
"cross_attn": hyvid_embeds["prompt_embeds"],
|
|
"attention_mask": hyvid_embeds["attention_mask"],
|
|
}
|
|
|
|
if hyvid_embeds["attention_mask_2"] is not None:
|
|
pooled_dict["attention_mask_controlnet"] = hyvid_embeds["attention_mask_2"]
|
|
|
|
if hyvid_embeds["cfg"] is not None:
|
|
pooled_dict["guidance"] = float(hyvid_embeds["cfg"])
|
|
pooled_dict["start_percent"] = float(hyvid_embeds["start_percent"]) if hyvid_embeds["start_percent"] is not None else 0.0
|
|
pooled_dict["end_percent"] = float(hyvid_embeds["end_percent"]) if hyvid_embeds["end_percent"] is not None else 1.0
|
|
|
|
return ([[cond, pooled_dict]],)
|
|
|
|
def override_class(cls):
|
|
class NodeOverride(cls):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
inputs = copy.deepcopy(cls.INPUT_TYPES())
|
|
devices = get_device_list()
|
|
default_device = devices[1] if len(devices) > 1 else devices[0]
|
|
inputs["optional"] = inputs.get("optional", {})
|
|
inputs["optional"]["device"] = (devices, {"default": default_device})
|
|
return inputs
|
|
|
|
CATEGORY = "multigpu"
|
|
FUNCTION = "override"
|
|
|
|
def override(self, *args, device=None, **kwargs):
|
|
global current_device
|
|
if device is not None:
|
|
current_device = device
|
|
fn = getattr(super(), cls.FUNCTION)
|
|
out = fn(*args, **kwargs)
|
|
return out
|
|
|
|
return NodeOverride
|
|
|
|
def override_class_with_offload(cls):
|
|
class NodeOverrideDiffSynth(cls):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
inputs = copy.deepcopy(cls.INPUT_TYPES())
|
|
devices = get_device_list()
|
|
default_device = devices[1] if len(devices) > 1 else devices[0]
|
|
inputs["optional"] = inputs.get("optional", {})
|
|
inputs["optional"]["device"] = (devices, {"default": default_device})
|
|
inputs["optional"]["offload_device"] = (devices, {"default": "cpu"})
|
|
return inputs
|
|
|
|
CATEGORY = "multigpu"
|
|
FUNCTION = "override"
|
|
|
|
def override(self, *args, device=None, offload_device=None, **kwargs):
|
|
global current_device
|
|
global current_offload_device
|
|
if device is not None:
|
|
current_device = device
|
|
if offload_device is not None:
|
|
current_offload_device = offload_device
|
|
fn = getattr(super(), cls.FUNCTION)
|
|
return fn(*args, **kwargs)
|
|
|
|
return NodeOverrideDiffSynth
|
|
|
|
|
|
def override_class_with_distorch(cls):
|
|
class NodeOverrideDisTorch(cls):
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
inputs = copy.deepcopy(cls.INPUT_TYPES())
|
|
devices = get_device_list()
|
|
default_device = devices[1] if len(devices) > 1 else devices[0]
|
|
inputs["optional"] = inputs.get("optional", {})
|
|
inputs["optional"]["device"] = (devices, {"default": default_device})
|
|
inputs["optional"]["allocations"] = ("STRING", {"multiline": False, "default": "cuda:0,0.15;cpu,0.5"})
|
|
return inputs
|
|
|
|
CATEGORY = "multigpu"
|
|
FUNCTION = "override"
|
|
|
|
def override(self, *args, device=None, allocations=None, **kwargs):
|
|
global current_device
|
|
if device is not None:
|
|
current_device = device
|
|
|
|
register_patched_ggufmodelpatcher()
|
|
fn = getattr(super(), cls.FUNCTION)
|
|
out = fn(*args, **kwargs)
|
|
|
|
if hasattr(out[0], 'model'):
|
|
model_hash = create_model_hash(out[0], "override")
|
|
model_allocation_store[model_hash] = allocations
|
|
elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'):
|
|
model_hash = create_model_hash(out[0].patcher, "override")
|
|
model_allocation_store[model_hash] = allocations
|
|
|
|
return out
|
|
|
|
return NodeOverrideDisTorch
|
|
|
|
def check_module_exists(module_path):
|
|
full_path = os.path.join(folder_paths.get_folder_paths("custom_nodes")[0], module_path)
|
|
logging.info(f"MultiGPU: Checking for module at {full_path}")
|
|
if not os.path.exists(full_path):
|
|
logging.info(f"MultiGPU: Module {module_path} not found - skipping")
|
|
return False
|
|
logging.info(f"MultiGPU: Found {module_path}, creating compatible MultiGPU nodes")
|
|
return True
|
|
|
|
from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS
|
|
from .nodes import (
|
|
UnetLoaderGGUF, UnetLoaderGGUFAdvanced,
|
|
CLIPLoaderGGUF, DualCLIPLoaderGGUF, TripleCLIPLoaderGGUF,
|
|
LTXVLoader,
|
|
Florence2ModelLoader, DownloadAndLoadFlorence2Model,
|
|
CheckpointLoaderNF4,
|
|
LoadFluxControlNet,
|
|
MMAudioModelLoader, MMAudioFeatureUtilsLoader, MMAudioSampler,
|
|
PulidModelLoader, PulidInsightFaceLoader, PulidEvaClipLoader,
|
|
HyVideoModelLoader, HyVideoModelLoaderDiffSynth, HyVideoVAELoader, DownloadAndLoadHyVideoTextEncoder
|
|
)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DeviceSelectorMultiGPU": DeviceSelectorMultiGPU,
|
|
"HunyuanVideoEmbeddingsAdapter": HunyuanVideoEmbeddingsAdapter
|
|
}
|
|
NODE_CLASS_MAPPINGS["UNETLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["UNETLoader"])
|
|
NODE_CLASS_MAPPINGS["VAELoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["VAELoader"])
|
|
NODE_CLASS_MAPPINGS["CLIPLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["CLIPLoader"])
|
|
NODE_CLASS_MAPPINGS["DualCLIPLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["DualCLIPLoader"])
|
|
NODE_CLASS_MAPPINGS["TripleCLIPLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["TripleCLIPLoader"])
|
|
NODE_CLASS_MAPPINGS["CheckpointLoaderSimpleMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"])
|
|
NODE_CLASS_MAPPINGS["ControlNetLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["ControlNetLoader"])
|
|
|
|
if check_module_exists("ComfyUI-LTXVideo") or check_module_exists("comfyui-ltxvideo"):
|
|
NODE_CLASS_MAPPINGS["LTXVLoaderMultiGPU"] = override_class(LTXVLoader)
|
|
|
|
if check_module_exists("ComfyUI-Florence2") or check_module_exists("comfyui-florence2"):
|
|
NODE_CLASS_MAPPINGS["Florence2ModelLoaderMultiGPU"] = override_class(Florence2ModelLoader)
|
|
NODE_CLASS_MAPPINGS["DownloadAndLoadFlorence2ModelMultiGPU"] = override_class(DownloadAndLoadFlorence2Model)
|
|
|
|
if check_module_exists("ComfyUI_bitsandbytes_NF4") or check_module_exists("comfyui_bitsandbytes_nf4"):
|
|
NODE_CLASS_MAPPINGS["CheckpointLoaderNF4MultiGPU"] = override_class(CheckpointLoaderNF4)
|
|
|
|
if check_module_exists("x-flux-comfyui") or check_module_exists("x-flux-comfyui"):
|
|
NODE_CLASS_MAPPINGS["LoadFluxControlNetMultiGPU"] = override_class(LoadFluxControlNet)
|
|
|
|
if check_module_exists("ComfyUI-MMAudio") or check_module_exists("comfyui-mmaudio"):
|
|
NODE_CLASS_MAPPINGS["MMAudioModelLoaderMultiGPU"] = override_class(MMAudioModelLoader)
|
|
NODE_CLASS_MAPPINGS["MMAudioFeatureUtilsLoaderMultiGPU"] = override_class(MMAudioFeatureUtilsLoader)
|
|
NODE_CLASS_MAPPINGS["MMAudioSamplerMultiGPU"] = override_class(MMAudioSampler)
|
|
|
|
if check_module_exists("ComfyUI-GGUF") or check_module_exists("comfyui-gguf"):
|
|
NODE_CLASS_MAPPINGS["UnetLoaderGGUFMultiGPU"] = override_class(UnetLoaderGGUF)
|
|
NODE_CLASS_MAPPINGS["UnetLoaderGGUFDisTorchMultiGPU"] = override_class_with_distorch(UnetLoaderGGUF)
|
|
NODE_CLASS_MAPPINGS["UnetLoaderGGUFAdvancedMultiGPU"] = override_class(UnetLoaderGGUFAdvanced)
|
|
NODE_CLASS_MAPPINGS["UnetLoaderGGUFAdvancedDisTorchMultiGPU"] = override_class_with_distorch(UnetLoaderGGUFAdvanced)
|
|
NODE_CLASS_MAPPINGS["CLIPLoaderGGUFMultiGPU"] = override_class(CLIPLoaderGGUF)
|
|
NODE_CLASS_MAPPINGS["CLIPLoaderGGUFDisTorchMultiGPU"] = override_class_with_distorch(CLIPLoaderGGUF)
|
|
NODE_CLASS_MAPPINGS["DualCLIPLoaderGGUFMultiGPU"] = override_class(DualCLIPLoaderGGUF)
|
|
NODE_CLASS_MAPPINGS["DualCLIPLoaderGGUFDisTorchMultiGPU"] = override_class_with_distorch(DualCLIPLoaderGGUF)
|
|
NODE_CLASS_MAPPINGS["TripleCLIPLoaderGGUFMultiGPU"] = override_class(TripleCLIPLoaderGGUF)
|
|
NODE_CLASS_MAPPINGS["TripleCLIPLoaderGGUFDisTorchMultiGPU"] = override_class_with_distorch(TripleCLIPLoaderGGUF)
|
|
|
|
if check_module_exists("PuLID_ComfyUI") or check_module_exists("pulid_comfyui"):
|
|
NODE_CLASS_MAPPINGS["PulidModelLoaderMultiGPU"] = override_class(PulidModelLoader)
|
|
NODE_CLASS_MAPPINGS["PulidInsightFaceLoaderMultiGPU"] = override_class(PulidInsightFaceLoader)
|
|
NODE_CLASS_MAPPINGS["PulidEvaClipLoaderMultiGPU"] = override_class(PulidEvaClipLoader)
|
|
|
|
if check_module_exists("ComfyUI-HunyuanVideoWrapper") or check_module_exists("comfyui-hunyuanvideowrapper"):
|
|
NODE_CLASS_MAPPINGS["HyVideoModelLoaderMultiGPU"] = override_class(HyVideoModelLoader)
|
|
NODE_CLASS_MAPPINGS["HyVideoVAELoaderMultiGPU"] = override_class(HyVideoVAELoader)
|
|
NODE_CLASS_MAPPINGS["DownloadAndLoadHyVideoTextEncoderMultiGPU"] = override_class(DownloadAndLoadHyVideoTextEncoder)
|
|
|
|
logging.info(f"MultiGPU: Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}")
|