import time import copy import torch import sys import comfy.model_management import os from pathlib import Path import importlib.util import logging import folder_paths from collections import defaultdict import builtins current_device = comfy.model_management.get_torch_device() current_offload_device = comfy.model_management.get_torch_device() def get_torch_device_patched(): device = None if (not torch.cuda.is_available() or comfy.model_management.cpu_state == comfy.model_management.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 comfy.model_management.cpu_state == comfy.model_management.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 comfy.model_management.cpu_state == comfy.model_management.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 comfy.model_management.cpu_state == comfy.model_management.CPUState.CPU or "cpu" in str(current_offload_device).lower()): device = torch.device("cpu") else: device = torch.device(current_offload_device) return device comfy.model_management.get_torch_device = get_torch_device_patched comfy.model_management.unet_offload_device = unet_offload_device_patched comfy.model_management.text_encoder_device = text_encoder_device_patched comfy.model_management.text_encoder_offload_device = text_encoder_offload_device_patched def register_patched_ggufmodelpatcher(node_instance): 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 logging.info("MultiGPU: GGUFDisTorch - GGUF ModelPatcher not yet patched, applying patch") def new_load(self, *args, force_patch_weights=False, **kwargs): super(module.GGUFModelPatcher, self).load(*args, force_patch_weights=True, **kwargs) 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) logging.info(f"MultiGPU: GGUFDisTorch - Weight Module {n} on device {device}, offload_device is {self.offload_device}") if device is not None: linked.append((n, m)) continue if hasattr(m, "bias"): device = getattr(m.bias, "device", None) #logging.info(f"MultiGPU: GGUFDisTorch - Bias Module {n} on device {device}, offload_device is {self.offload_device}") if device is not None: linked.append((n, m)) continue logging.info(f"MultiGPU: GGUFDisTorch - Found {len(linked)} linked modules out of {module_count} total modules") if linked: logging.info(f"MultiGPU: GGUFDisTorch - Found {len(linked)} linked modules, computing reallocation") device_assignments = analyze_ggml_loading(self.model, node_instance.distorch_allocations)['device_assignments'] for device, layers in device_assignments.items(): logging.info(f"MultiGPU: GGUFDisTorch - Moving {len(layers)} layers to {device}") target_device = torch.device(device) #logging.info(f"MultiGPU: GGUFDisTorch - Moving {len(layers)} layers to {device}") for n, m, _ in layers: m.to(self.load_device).to(target_device) self.mmap_released = True logging.info("MultiGPU: GGUFDisTorch - self.mmap_released = True") module.GGUFModelPatcher.load = new_load module.GGUFModelPatcher._patched = True logging.info("MultiGPU: GGUFDisTorch - Successfully patched GGUF ModelPatcher") else: logging.info("MultiGPU: GGUFDisTorch - GGUF ModelPatcher already patched") def analyze_ggml_loading(model, distorch_allocations): DEVICE_RATIOS_DISTORCH = {} device_table = {} primary_dev_name = distorch_allocations.get("compute_device") primary_total_mem_bytes = comfy.model_management.get_total_memory(torch.device(primary_dev_name)) primary_fraction = distorch_allocations.get("compute_device_alloc", 0.0) primary_alloc_gb = (primary_total_mem_bytes * primary_fraction) / (1024**3) DEVICE_RATIOS_DISTORCH[primary_dev_name] = primary_alloc_gb device_table[primary_dev_name] = {"fraction": primary_fraction,"total_gb": primary_total_mem_bytes / (1024**3),"alloc_gb": primary_alloc_gb} i = 1 while f"distorch{i}_device" in distorch_allocations: dev_key = f"distorch{i}_device" alloc_key = f"distorch{i}_alloc" dev_name = distorch_allocations[dev_key] dev_total_mem_bytes = comfy.model_management.get_total_memory(torch.device(dev_name)) dev_fraction = distorch_allocations.get(alloc_key, 0.0) dev_alloc_gb = (dev_total_mem_bytes * dev_fraction) / (1024**3) DEVICE_RATIOS_DISTORCH[dev_name] = dev_alloc_gb device_table[dev_name] = {"fraction": dev_fraction,"total_gb": dev_total_mem_bytes / (1024**3),"alloc_gb": dev_alloc_gb} i += 1 cpu_dev_name = distorch_allocations.get("distorch_cpu", "cpu") cpu_total_mem_bytes = comfy.model_management.get_total_memory(torch.device(cpu_dev_name)) cpu_fraction = distorch_allocations.get("distorch_cpu_alloc", 0.0) cpu_alloc_gb = (cpu_total_mem_bytes * cpu_fraction) / (1024**3) DEVICE_RATIOS_DISTORCH[cpu_dev_name] = cpu_alloc_gb device_table[cpu_dev_name] = {"fraction": cpu_fraction,"total_gb": cpu_total_mem_bytes / (1024**3),"alloc_gb": cpu_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 UNet 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,) 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) return fn(*args, **kwargs) 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): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.distorch_allocations = {} self.distorch_compute_device = None @classmethod def INPUT_TYPES(s): inputs = copy.deepcopy(cls.INPUT_TYPES()) devices = [d for d in get_device_list() if d != "cpu"] inputs["optional"] = inputs.get("optional", {}) inputs["required"]["compute_device"] = (devices, {"default": devices[0], "tooltip": "Device model will use for computation"}) inputs["required"]["compute_device_alloc"] = ("FLOAT", {"default": 0.15, "step": 0.01, "tooltip": "Fraction of memory NOT allocated to active latent space computation, recommended <= 15%"}) for i in range(len(devices) - 1): inputs["optional"][f"distorch{i+1}_device"] = (devices, {"default": devices[i+1], "tooltip": f"Device for distorch{i+1} model layer VRAM allocation"}) inputs["optional"][f"distorch{i+1}_alloc"] = ("FLOAT", {"default": 0.9, "step": 0.01, "tooltip": f"Fraction of memory allocated to distorch{i+1} model layer, recommended >= 90%"}) inputs["optional"]["distorch_cpu"] = (["cpu"], {"default": "cpu", "tooltip": "Device for distorch CPU memory allocation"}) inputs["optional"]["distorch_cpu_alloc"] = ("FLOAT", {"default": 0.0, "step": 0.01, "tooltip": "Fraction of memory allocated to distorch CPU memory (potentially slower than cuda)"}) return inputs CATEGORY = "multigpu" FUNCTION = "override" def override(self, *args, **kwargs): self.distorch_allocations = {} self.distorch_compute_device = kwargs.get("compute_device", None) if self.distorch_compute_device is not None: global current_device current_device = self.distorch_compute_device register_patched_ggufmodelpatcher(self) for key, value in list(kwargs.items()): if key not in {"unet_name", "clip_name1", "clip_name2", "clip_name2", "type"}: logging.info(f"MultiGPU: Removing {key} from kwargs") logging.info(f"MultiGPU: Value: {value}") self.distorch_allocations[key] = kwargs.pop(key) fn = getattr(super(), cls.FUNCTION) return fn(*args, **kwargs) return NodeOverrideDisTorch NODE_CLASS_MAPPINGS = {"DeviceSelectorMultiGPU": DeviceSelectorMultiGPU} def check_module_exists(module_path): full_path = os.path.join("custom_nodes", 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 def register_module(target_nodes): from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS for node in target_nodes: NODE_CLASS_MAPPINGS[f"{node}MultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS[node]) return def register_UnetLoaderGGUFMultiGPU(): global NODE_CLASS_MAPPINGS # First define the base UnetLoaderGGUF class class UnetLoaderGGUF: @classmethod def INPUT_TYPES(s): unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")] return { "required": { "unet_name": (unet_names,), } } RETURN_TYPES = ("MODEL",) FUNCTION = "load_unet" CATEGORY = "bootleg" TITLE = "Unet Loader (GGUF)" def load_unet(self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["UnetLoaderGGUF"]() return original_loader.load_unet(unet_name, dequant_dtype, patch_dtype, patch_on_device) # Create the MultiGPU version of the base class UnetLoaderGGUFMultiGPU = override_class(UnetLoaderGGUF) NODE_CLASS_MAPPINGS["UnetLoaderGGUFMultiGPU"] = UnetLoaderGGUFMultiGPU logging.info(f"MultiGPU: Registered UnetLoaderGGUFMultiGPU") # Now create the advanced version that inherits from the MultiGPU base class class UnetLoaderGGUFAdvanced(UnetLoaderGGUF): @classmethod def INPUT_TYPES(s): unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")] return { "required": { "unet_name": (unet_names,), "dequant_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}), "patch_dtype": (["default", "target", "float32", "float16", "bfloat16"], {"default": "default"}), "patch_on_device": ("BOOLEAN", {"default": False}), } } TITLE = "Unet Loader (GGUF/Advanced)" # Create the MultiGPU version of the advanced class UnetLoaderGGUFAdvancedMultiGPU = override_class(UnetLoaderGGUFAdvanced) NODE_CLASS_MAPPINGS["UnetLoaderGGUFAdvancedMultiGPU"] = UnetLoaderGGUFAdvancedMultiGPU logging.info(f"MultiGPU: Registered UnetLoaderGGUFAdvancedMultiGPU") def register_UnetLoaderGGUFDisTorchMultiGPU(): global NODE_CLASS_MAPPINGS global distorch_compute_device # First define the base UnetLoaderGGUFDisTorch class class UnetLoaderGGUFDisTorch: @classmethod def INPUT_TYPES(s): unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")] return { "required": { "unet_name": (unet_names,), } } RETURN_TYPES = ("MODEL",) FUNCTION = "load_unet" CATEGORY = "bootleg" TITLE = "Unet Loader (GGUFDisTorch)" def load_unet(self, unet_name, dequant_dtype=None, patch_dtype=None, patch_on_device=None): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["UnetLoaderGGUF"]() return original_loader.load_unet(unet_name, dequant_dtype, patch_dtype, patch_on_device) # Create the MultiGPU version of the base class UnetLoaderGGUFDisTorchMultiGPU = override_class_with_distorch(UnetLoaderGGUFDisTorch) NODE_CLASS_MAPPINGS["UnetLoaderGGUFDisTorchMultiGPU"] = UnetLoaderGGUFDisTorchMultiGPU logging.info(f"MultiGPU: Registered UnetLoaderGGUFDisTorchMultiGPU") def register_CLIPLoaderGGUFMultiGPU(): global NODE_CLASS_MAPPINGS class CLIPLoaderGGUF: @classmethod def INPUT_TYPES(s): return { "required": { "clip_name": (s.get_filename_list(),), "type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv"],), } } RETURN_TYPES = ("CLIP",) FUNCTION = "load_clip" CATEGORY = "bootleg" TITLE = "CLIPLoader (GGUF)" @classmethod def get_filename_list(s): files = [] files += folder_paths.get_filename_list("clip") files += folder_paths.get_filename_list("clip_gguf") return sorted(files) def load_data(self, ckpt_paths): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["CLIPLoaderGGUF"]() return original_loader.load_data(ckpt_paths) def load_patcher(self, clip_paths, clip_type, clip_data): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["CLIPLoaderGGUF"]() return original_loader.load_patcher(clip_paths, clip_type, clip_data) def load_clip(self, clip_name, type="stable_diffusion"): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["CLIPLoaderGGUF"]() return original_loader.load_clip(clip_name, type) # Create the MultiGPU version of the base class CLIPLoaderGGUFMultiGPU = override_class(CLIPLoaderGGUF) NODE_CLASS_MAPPINGS["CLIPLoaderGGUFMultiGPU"] = CLIPLoaderGGUFMultiGPU logging.info(f"MultiGPU: Registered CLIPLoaderGGUFMultiGPU") CLIPLoaderGGUFDisTorchMultiGPU = override_class_with_distorch(CLIPLoaderGGUF) NODE_CLASS_MAPPINGS["CLIPLoaderGGUFDisTorchMultiGPU"] = CLIPLoaderGGUFDisTorchMultiGPU logging.info(f"MultiGPU: Registered CLIPLoaderGGUFDisTorchMultiGPU") # Now create the advanced version that inherits from the MultiGPU base class class DualCLIPLoaderGGUF(CLIPLoaderGGUF): @classmethod def INPUT_TYPES(s): file_options = (s.get_filename_list(), ) return { "required": { "clip_name1": file_options, "clip_name2": file_options, "type": (("sdxl", "sd3", "flux", "hunyuan_video"),), } } TITLE = "DualCLIPLoader (GGUF)" def load_clip(self, clip_name1, clip_name2, type): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["DualCLIPLoaderGGUF"]() clip = original_loader.load_clip(clip_name1, clip_name2, type) clip[0].patcher.load(force_patch_weights=True) return clip # Create the MultiGPU version of the advanced class DualCLIPLoaderGGUFMultiGPU = override_class(DualCLIPLoaderGGUF) NODE_CLASS_MAPPINGS["DualCLIPLoaderGGUFMultiGPU"] = DualCLIPLoaderGGUFMultiGPU logging.info(f"MultiGPU: Registered DualCLIPLoaderGGUFMultiGPU") DualCLIPLoaderGGUFDisTorchMultiGPU = override_class_with_distorch(DualCLIPLoaderGGUF) NODE_CLASS_MAPPINGS["DualCLIPLoaderGGUFDisTorchMultiGPU"] = DualCLIPLoaderGGUFDisTorchMultiGPU logging.info(f"MultiGPU: Registered DualCLIPLoaderGGUFDisTorchMultiGPU") class TripleCLIPLoaderGGUF(CLIPLoaderGGUF): @classmethod def INPUT_TYPES(s): file_options = (s.get_filename_list(), ) return { "required": { "clip_name1": file_options, "clip_name2": file_options, "clip_name3": file_options, } } TITLE = "TripleCLIPLoader (GGUF)" def load_clip(self, clip_name1, clip_name2, clip_name3, type="sd3"): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["TripleCLIPLoaderGGUF"]() return original_loader.load_clip(clip_name1, clip_name2, clip_name3, type) # Create the MultiGPU version of the advanced class TripleCLIPLoaderGGUFMultiGPU = override_class(TripleCLIPLoaderGGUF) NODE_CLASS_MAPPINGS["TripleCLIPLoaderGGUFMultiGPU"] = TripleCLIPLoaderGGUFMultiGPU logging.info(f"MultiGPU: Registered TripleCLIPLoaderGGUFMultiGPU") TripleCLIPLoaderGGUFDisTorchMultiGPU = override_class_with_distorch(TripleCLIPLoaderGGUF) NODE_CLASS_MAPPINGS["TripleCLIPLoaderGGUFDisTorchMultiGPU"] = TripleCLIPLoaderGGUFDisTorchMultiGPU logging.info(f"MultiGPU: Registered TripleCLIPLoaderGGUFDisTorchMultiGPU") def register_LTXVLoaderMultiGPU(): global NODE_CLASS_MAPPINGS class LTXVLoader: @classmethod def INPUT_TYPES(s): return { "required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), {"tooltip": "The name of the checkpoint (model) to load."}), "dtype": (["bfloat16", "float32"], {"default": "bfloat16"}) } } RETURN_TYPES = ("MODEL", "VAE") RETURN_NAMES = ("model", "vae") FUNCTION = "load" CATEGORY = "lightricks/LTXV" TITLE = "LTXV Loader" OUTPUT_NODE = False def load(self, ckpt_name, dtype): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["LTXVLoader"]() return original_loader.load(ckpt_name, dtype) def _load_unet(self, load_device, offload_device, weights, num_latent_channels, dtype, config=None ): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["LTXVLoader"]() return original_loader._load_unet(load_device, offload_device, weights, num_latent_channels, dtype, config=None ) def _load_vae(self, weights, config=None): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["LTXVLoader"]() return original_loader._load_vae(weights, config=None) NODE_CLASS_MAPPINGS["LTXVLoaderMultiGPU"] = override_class(LTXVLoader) logging.info(f"MultiGPU: Registered LTXVLoaderMultiGPU") def register_Florence2ModelLoaderMultiGPU(): global NODE_CLASS_MAPPINGS class Florence2ModelLoader: @classmethod def INPUT_TYPES(s): return {"required": { "model": ([item.name for item in Path(folder_paths.models_dir, "LLM").iterdir() if item.is_dir()], {"tooltip": "models are expected to be in Comfyui/models/LLM folder"}), "precision": (['fp16','bf16','fp32'],), "attention": (['flash_attention_2', 'sdpa', 'eager'], {"default": 'sdpa'}), }, "optional": { "lora": ("PEFTLORA",), }} RETURN_TYPES = ("FL2MODEL",) RETURN_NAMES = ("florence2_model",) FUNCTION = "loadmodel" CATEGORY = "Florence2" def loadmodel(self, model, precision, attention, lora=None): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["Florence2ModelLoader"]() return original_loader.loadmodel(model, precision, attention, lora) NODE_CLASS_MAPPINGS["Florence2ModelLoaderMultiGPU"] = override_class(Florence2ModelLoader) logging.info(f"MultiGPU: Registered Florence2ModelLoaderMultiGPU") def register_DownloadAndLoadFlorence2ModelMultiGPU(): global NODE_CLASS_MAPPINGS class DownloadAndLoadFlorence2Model: @classmethod def INPUT_TYPES(s): return {"required": { "model": ([ 'microsoft/Florence-2-base', 'microsoft/Florence-2-base-ft', 'microsoft/Florence-2-large', 'microsoft/Florence-2-large-ft', 'HuggingFaceM4/Florence-2-DocVQA', 'thwri/CogFlorence-2.1-Large', 'thwri/CogFlorence-2.2-Large', 'gokaygokay/Florence-2-SD3-Captioner', 'gokaygokay/Florence-2-Flux-Large', 'MiaoshouAI/Florence-2-base-PromptGen-v1.5', 'MiaoshouAI/Florence-2-large-PromptGen-v1.5', 'MiaoshouAI/Florence-2-base-PromptGen-v2.0', 'MiaoshouAI/Florence-2-large-PromptGen-v2.0' ], {"default": 'microsoft/Florence-2-base'}), "precision": (['fp16','bf16','fp32'], {"default": 'fp16'}), "attention": (['flash_attention_2', 'sdpa', 'eager'], {"default": 'sdpa'}), }, "optional": { "lora": ("PEFTLORA",), }} RETURN_TYPES = ("FL2MODEL",) RETURN_NAMES = ("florence2_model",) FUNCTION = "loadmodel" CATEGORY = "Florence2" def loadmodel(self, model, precision, attention, lora=None): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["DownloadAndLoadFlorence2Model"]() return original_loader.loadmodel(model, precision, attention, lora) NODE_CLASS_MAPPINGS["DownloadAndLoadFlorence2ModelMultiGPU"] = override_class(DownloadAndLoadFlorence2Model) logging.info(f"MultiGPU: Registered DownloadAndLoadFlorence2ModelMultiGPU") def register_CheckpointLoaderNF4(): global NODE_CLASS_MAPPINGS class CheckpointLoaderNF4: @classmethod def INPUT_TYPES(s): return {"required": { "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), }} RETURN_TYPES = ("MODEL", "CLIP", "VAE") FUNCTION = "load_checkpoint" CATEGORY = "loaders" def load_checkpoint(self, ckpt_name): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["CheckpointLoaderNF4"]() return original_loader.load_checkpoint(ckpt_name) NODE_CLASS_MAPPINGS["CheckpointLoaderNF4MultiGPU"] = override_class(CheckpointLoaderNF4) logging.info(f"MultiGPU: Registered CheckpointLoaderNF4MultiGPU") def register_LoadFluxControlNetMultiGPU(): global NODE_CLASS_MAPPINGS class LoadFluxControlNet: @classmethod def INPUT_TYPES(s): return {"required": {"model_name": (["flux-dev", "flux-dev-fp8", "flux-schnell"],), "controlnet_path": (folder_paths.get_filename_list("xlabs_controlnets"), ), }} RETURN_TYPES = ("FluxControlNet",) RETURN_NAMES = ("ControlNet",) FUNCTION = "loadmodel" CATEGORY = "XLabsNodes" def loadmodel(self, model_name, controlnet_path): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["LoadFluxControlNet"]() return original_loader.loadmodel(model_name, controlnet_path) NODE_CLASS_MAPPINGS["LoadFluxControlNetMultiGPU"] = override_class(LoadFluxControlNet) logging.info(f"MultiGPU: Registered LoadFluxControlNetMultiGPU") def register_MMAudioModelLoaderMultiGPU(): global NODE_CLASS_MAPPINGS class MMAudioModelLoader: @classmethod def INPUT_TYPES(s): return { "required": { "mmaudio_model": (folder_paths.get_filename_list("mmaudio"), {"tooltip": "These models are loaded from the 'ComfyUI/models/mmaudio' -folder",}), "base_precision": (["fp16", "fp32", "bf16"], {"default": "fp16"}), }, } RETURN_TYPES = ("MMAUDIO_MODEL",) RETURN_NAMES = ("mmaudio_model", ) FUNCTION = "loadmodel" CATEGORY = "MMAudio" def loadmodel(self, mmaudio_model, base_precision): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["MMAudioModelLoader"]() return original_loader.loadmodel(mmaudio_model, base_precision) NODE_CLASS_MAPPINGS["MMAudioModelLoaderMultiGPU"] = override_class(MMAudioModelLoader) logging.info(f"MultiGPU: Registered MMAudioModelLoaderMultiGPU") def register_MMAudioFeatureUtilsLoaderMultiGPU(): global NODE_CLASS_MAPPINGS class MMAudioFeatureUtilsLoader: @classmethod def INPUT_TYPES(s): return { "required": { "vae_model": (folder_paths.get_filename_list("mmaudio"), {"tooltip": "These models are loaded from 'ComfyUI/models/mmaudio'"}), "synchformer_model": (folder_paths.get_filename_list("mmaudio"), {"tooltip": "These models are loaded from 'ComfyUI/models/mmaudio'"}), "clip_model": (folder_paths.get_filename_list("mmaudio"), {"tooltip": "These models are loaded from 'ComfyUI/models/mmaudio'"}), }, "optional": { "bigvgan_vocoder_model": ("VOCODER_MODEL", {"tooltip": "These models are loaded from 'ComfyUI/models/mmaudio'"}), "mode": (["16k", "44k"], {"default": "44k"}), "precision": (["fp16", "fp32", "bf16"], {"default": "fp16"} ), } } RETURN_TYPES = ("MMAUDIO_FEATUREUTILS",) RETURN_NAMES = ("mmaudio_featureutils", ) FUNCTION = "loadmodel" CATEGORY = "MMAudio" def loadmodel(self, vae_model, precision, synchformer_model, clip_model, mode, bigvgan_vocoder_model=None): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["MMAudioFeatureUtilsLoader"]() return original_loader.loadmodel(vae_model, precision, synchformer_model, clip_model, mode, bigvgan_vocoder_model) NODE_CLASS_MAPPINGS["MMAudioFeatureUtilsLoaderMultiGPU"] = override_class(MMAudioFeatureUtilsLoader) logging.info(f"MultiGPU: Registered MMAudioFeatureUtilsLoaderMultiGPU") def register_MMAudioSamplerMultiGPU(): global NODE_CLASS_MAPPINGS class MMAudioSampler: @classmethod def INPUT_TYPES(s): return { "required": { "mmaudio_model": ("MMAUDIO_MODEL",), "feature_utils": ("MMAUDIO_FEATUREUTILS",), "duration": ("FLOAT", {"default": 8, "step": 0.01, "tooltip": "Duration of the audio in seconds"}), "steps": ("INT", {"default": 25, "step": 1, "tooltip": "Number of steps to interpolate"}), "cfg": ("FLOAT", {"default": 4.5, "step": 0.1, "tooltip": "Strength of the conditioning"}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "prompt": ("STRING", {"default": "", "multiline": True} ), "negative_prompt": ("STRING", {"default": "", "multiline": True} ), "mask_away_clip": ("BOOLEAN", {"default": False, "tooltip": "If true, the clip video will be masked away"}), "force_offload": ("BOOLEAN", {"default": True, "tooltip": "If true, the model will be offloaded to the offload device"}), }, "optional": { "images": ("IMAGE",), }, } RETURN_TYPES = ("AUDIO",) RETURN_NAMES = ("audio", ) FUNCTION = "sample" CATEGORY = "MMAudio" def sample(self, mmaudio_model, seed, feature_utils, duration, steps, cfg, prompt, negative_prompt, mask_away_clip, force_offload, images=None): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["MMAudioSampler"]() return original_loader.sample(mmaudio_model, seed, feature_utils, duration, steps, cfg, prompt, negative_prompt, mask_away_clip, force_offload, images) NODE_CLASS_MAPPINGS["MMAudioSamplerMultiGPU"] = override_class(MMAudioSampler) logging.info(f"MultiGPU: Registered MMAudioSamplerMultiGPU") def register_PulidModelLoader(): global NODE_CLASS_MAPPINGS class PulidModelLoader: @classmethod def INPUT_TYPES(s): return {"required": { "pulid_file": (folder_paths.get_filename_list("pulid"), )}} RETURN_TYPES = ("PULID",) FUNCTION = "load_model" CATEGORY = "pulid" def load_model(self, pulid_file): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["PulidModelLoader"]() return original_loader.load_model(pulid_file) NODE_CLASS_MAPPINGS["PulidModelLoaderMultiGPU"] = override_class(PulidModelLoader) logging.info(f"MultiGPU: Registered PulidModelLoaderMultiGPU") def register_PulidInsightFaceLoader(): global NODE_CLASS_MAPPINGS class PulidInsightFaceLoader: @classmethod def INPUT_TYPES(s): return { "required": { "provider": (["CPU", "CUDA", "ROCM", "CoreML"], ), }, } RETURN_TYPES = ("FACEANALYSIS",) FUNCTION = "load_insightface" CATEGORY = "pulid" def load_insightface(self, provider): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["PulidInsightFaceLoader"]() return original_loader.load_insightface(provider) NODE_CLASS_MAPPINGS["PulidInsightFaceLoaderMultiGPU"] = override_class(PulidInsightFaceLoader) logging.info(f"MultiGPU: Registered PulidInsightFaceLoaderMultiGPU") def register_PulidEvaClipLoader(): global NODE_CLASS_MAPPINGS class PulidEvaClipLoader: @classmethod def INPUT_TYPES(s): return { "required": {}, } RETURN_TYPES = ("EVA_CLIP",) FUNCTION = "load_eva_clip" CATEGORY = "pulid" def load_eva_clip(self): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["PulidEvaClipLoader"]() return original_loader.load_eva_clip() NODE_CLASS_MAPPINGS["PulidEvaClipLoaderMultiGPU"] = override_class(PulidEvaClipLoader) logging.info(f"MultiGPU: Registered PulidEvaClipLoaderMultiGPU") def register_HyVideoModelLoader(): global NODE_CLASS_MAPPINGS # Keep original MultiGPU wrapper unchanged class HyVideoModelLoader: @classmethod def INPUT_TYPES(s): return { "required": { "model": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder",}), "base_precision": (["fp32", "bf16"], {"default": "bf16"}), "quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_scaled', 'torchao_fp8dq', "torchao_fp8dqrow", "torchao_int8dq", "torchao_fp6", "torchao_int4", "torchao_int8"], {"default": 'disabled', "tooltip": "optional quantization method"}), "load_device": (["main_device"], {"default": "main_device"}), }, "optional": { "attention_mode": ([ "sdpa", "flash_attn_varlen", "sageattn_varlen", "comfy", ], {"default": "flash_attn"}), "compile_args": ("COMPILEARGS", ), "block_swap_args": ("BLOCKSWAPARGS", ), "lora": ("HYVIDLORA", {"default": None}), "auto_cpu_offload": ("BOOLEAN", {"default": False, "tooltip": "Enable auto offloading for reduced VRAM usage, implementation from DiffSynth-Studio, slightly different from block swapping and uses even less VRAM, but can be slower as you can't define how much VRAM to use"}), } } RETURN_TYPES = ("HYVIDEOMODEL",) RETURN_NAMES = ("model", ) FUNCTION = "loadmodel" CATEGORY = "HunyuanVideoWrapper" def loadmodel(self, model, base_precision, load_device, quantization, compile_args=None, attention_mode="sdpa", block_swap_args=None, lora=None, auto_cpu_offload=False): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["HyVideoModelLoader"]() return original_loader.loadmodel(model, base_precision, load_device, quantization, compile_args, attention_mode, block_swap_args, lora, auto_cpu_offload) # Add new DiffSynth-style node class HyVideoModelLoaderDiffSynth: @classmethod def INPUT_TYPES(s): return { "required": { "model": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder",}), "base_precision": (["fp32", "bf16"], {"default": "bf16"}), "quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_scaled', 'torchao_fp8dq', "torchao_fp8dqrow", "torchao_int8dq", "torchao_fp6", "torchao_int4", "torchao_int8"], {"default": 'disabled', "tooltip": "optional quantization method"}), }, "optional": { "attention_mode": ([ "sdpa", "flash_attn_varlen", "sageattn_varlen", "comfy", ], {"default": "flash_attn"}), "compile_args": ("COMPILEARGS", ), "block_swap_args": ("BLOCKSWAPARGS", ), "lora": ("HYVIDLORA", {"default": None}), } } RETURN_TYPES = ("HYVIDEOMODEL",) RETURN_NAMES = ("model", ) FUNCTION = "loadmodel" CATEGORY = "HunyuanVideoWrapper" def loadmodel(self, model, base_precision, quantization, compile_args=None, attention_mode="sdpa", block_swap_args=None, lora=None): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["HyVideoModelLoader"]() # Use DiffSynth's auto offloading approach return original_loader.loadmodel(model, base_precision, "main_device", quantization, compile_args, attention_mode, block_swap_args, lora, auto_cpu_offload=True) # Register both with MultiGPU wrapper NODE_CLASS_MAPPINGS["HyVideoModelLoaderMultiGPU"] = override_class(HyVideoModelLoader) NODE_CLASS_MAPPINGS["HyVideoModelLoaderDiffSynthMultiGPU"] = override_class_with_offload(HyVideoModelLoaderDiffSynth) logging.info(f"MultiGPU: Registered HyVideoModelLoader nodes") def register_HyVideoVAELoader(): global NODE_CLASS_MAPPINGS class HyVideoVAELoader: @classmethod def INPUT_TYPES(s): return { "required": { "model_name": (folder_paths.get_filename_list("vae"), {"tooltip": "These models are loaded from 'ComfyUI/models/vae'"}), }, "optional": { "precision": (["fp16", "fp32", "bf16"], {"default": "bf16"} ), "compile_args":("COMPILEARGS", ), } } RETURN_TYPES = ("VAE",) RETURN_NAMES = ("vae", ) FUNCTION = "loadmodel" CATEGORY = "HunyuanVideoWrapper" DESCRIPTION = "Loads Hunyuan VAE model from 'ComfyUI/models/vae'" def loadmodel(self, model_name, precision, compile_args=None): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["HyVideoVAELoader"]() return original_loader.loadmodel(model_name, precision, compile_args) NODE_CLASS_MAPPINGS["HyVideoVAELoaderMultiGPU"] = override_class(HyVideoVAELoader) logging.info(f"MultiGPU: Registered HyVideoVAELoaderMultiGPU") def register_DownloadAndLoadHyVideoTextEncoder(): global NODE_CLASS_MAPPINGS class DownloadAndLoadHyVideoTextEncoder: @classmethod def INPUT_TYPES(s): return { "required": { "llm_model": (["Kijai/llava-llama-3-8b-text-encoder-tokenizer","xtuner/llava-llama-3-8b-v1_1-transformers"],), "clip_model": (["disabled","openai/clip-vit-large-patch14",],), "precision": (["fp16", "fp32", "bf16"], {"default": "bf16"} ), }, "optional": { "apply_final_norm": ("BOOLEAN", {"default": False}), "hidden_state_skip_layer": ("INT", {"default": 2}), "quantization": (['disabled', 'bnb_nf4', "fp8_e4m3fn"], {"default": 'disabled'}), } } RETURN_TYPES = ("HYVIDTEXTENCODER",) RETURN_NAMES = ("hyvid_text_encoder", ) FUNCTION = "loadmodel" CATEGORY = "HunyuanVideoWrapper" DESCRIPTION = "Loads Hunyuan text_encoder model from 'ComfyUI/models/LLM'" def loadmodel(self, llm_model, clip_model, precision, apply_final_norm=False, hidden_state_skip_layer=2, quantization="disabled"): from nodes import NODE_CLASS_MAPPINGS original_loader = NODE_CLASS_MAPPINGS["DownloadAndLoadHyVideoTextEncoder"]() return original_loader.loadmodel(llm_model, clip_model, precision, apply_final_norm, hidden_state_skip_layer, quantization) NODE_CLASS_MAPPINGS["DownloadAndLoadHyVideoTextEncoderMultiGPU"] = override_class(DownloadAndLoadHyVideoTextEncoder) logging.info(f"MultiGPU: Registered DownloadAndLoadHyVideoTextEncoderMultiGPU") # Register desired nodes register_module(["UNETLoader", "VAELoader", "CLIPLoader", "DualCLIPLoader", "TripleCLIPLoader", "CheckpointLoaderSimple", "ControlNetLoader"]) if check_module_exists("ComfyUI-LTXVideo"): register_LTXVLoaderMultiGPU() if check_module_exists("ComfyUI-Florence2"): register_Florence2ModelLoaderMultiGPU() register_DownloadAndLoadFlorence2ModelMultiGPU() if check_module_exists("ComfyUI_bitsandbytes_NF4"): register_CheckpointLoaderNF4() if check_module_exists("x-flux-comfyui"): register_LoadFluxControlNetMultiGPU() if check_module_exists("ComfyUI-MMAudio"): register_MMAudioModelLoaderMultiGPU() register_MMAudioFeatureUtilsLoaderMultiGPU() register_MMAudioSamplerMultiGPU() if check_module_exists("ComfyUI-GGUF"): register_UnetLoaderGGUFMultiGPU() register_UnetLoaderGGUFDisTorchMultiGPU() register_CLIPLoaderGGUFMultiGPU() if check_module_exists("PuLID_ComfyUI"): register_PulidModelLoader() register_PulidInsightFaceLoader() register_PulidEvaClipLoader() if check_module_exists("ComfyUI-HunyuanVideoWrapper"): register_HyVideoModelLoader() register_HyVideoVAELoader() register_DownloadAndLoadHyVideoTextEncoder() logging.info(f"MultiGPU: Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}")