From 1173245de99adba691647eb22f1107cfc5d07a6a Mon Sep 17 00:00:00 2001 From: hum_ma Date: Fri, 6 Feb 2026 11:28:48 +0200 Subject: [PATCH] Add files via upload --- __init__.py | 124 ++++++++++++++++++++++++++++++++++++++ device_utils.py | 155 ++++++++++++++++++++++++++++++++++++++++++++++++ nodes.py | 127 +++++++++++++++++++++++++++++++++++++++ wrappers.py | 93 +++++++++++++++++++++++++++++ 4 files changed, 499 insertions(+) create mode 100644 __init__.py create mode 100644 device_utils.py create mode 100644 nodes.py create mode 100644 wrappers.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..ae4e7c4 --- /dev/null +++ b/__init__.py @@ -0,0 +1,124 @@ +import torch +import logging +import os +import folder_paths +import comfy.model_management as mm +from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS +from .device_utils import ( + get_device_list, + is_accelerator_available, +) + +MGPU_MM_LOG = False +DEBUG_LOG = False + +logger = logging.getLogger("GPUClip") +logger.propagate = False + +if not logger.handlers: + log_level = logging.DEBUG if DEBUG_LOG else logging.INFO + handler = logging.StreamHandler() + formatter = logging.Formatter('%(message)s') + handler.setFormatter(formatter) + logger.addHandler(handler) + logger.setLevel(log_level) + +def check_module_exists(module_path): + """Check if a custom node module exists in ComfyUI custom_nodes directory.""" + full_path = os.path.join(folder_paths.get_folder_paths("custom_nodes")[0], module_path) + logger.debug(f"[GPUClip] Checking for module at {full_path}") + if not os.path.exists(full_path): + logger.debug(f"[GPUClip] Module {module_path} not found - skipping") + return False + logger.debug(f"[GPUClip] Found {module_path}, creating compatible MultiGPU nodes") + return True + +current_device = mm.get_torch_device() +current_text_encoder_device = mm.text_encoder_device() + +def set_current_device(device): + """Set the current device context for MultiGPU operations.""" + global current_device + current_device = device + logger.debug(f"[GPUClip Initialization] current_device set to: {device}") + +def set_current_text_encoder_device(device): + """Set the current text encoder device context for CLIP models.""" + global current_text_encoder_device + current_text_encoder_device = device + logger.debug(f"[GPUClip Initialization] current_text_encoder_device set to: {device}") + +def get_torch_device_patched(): + """Return MultiGPU-aware device selection for patched mm.get_torch_device.""" + device = None + if (not is_accelerator_available() or mm.cpu_state == mm.CPUState.CPU or "cpu" in str(current_device).lower()): + device = torch.device("cpu") + else: + devs = set(get_device_list()) + device = torch.device(current_device) if str(current_device) in devs else torch.device("cpu") + logger.debug(f"[GPUClip Core Patching] get_torch_device_patched returning device: {device} (current_device={current_device})") + return device + +def text_encoder_device_patched(): + """Return MultiGPU-aware text encoder device for patched mm.text_encoder_device.""" + device = None + if (not is_accelerator_available() or mm.cpu_state == mm.CPUState.CPU or "cpu" in str(current_text_encoder_device).lower()): + device = torch.device("cpu") + else: + devs = set(get_device_list()) + device = torch.device(current_text_encoder_device) if str(current_text_encoder_device) in devs else torch.device("cpu") + logger.info(f"[GPUClip Core Patching] text_encoder_device_patched returning device: {device} (current_text_encoder_device={current_text_encoder_device})") + return device + +mm.get_torch_device = get_torch_device_patched +mm.text_encoder_device = text_encoder_device_patched + +from .nodes import ( + GPUClipDeviceSelector, + CLIPLoaderGGUF, + DualCLIPLoaderGGUF, + TripleCLIPLoaderGGUF, + QuadrupleCLIPLoaderGGUF, +) + +from .wrappers import ( + override_class, + override_class_clip, + override_class_clip_no_device, +) + +NODE_CLASS_MAPPINGS = { + "GPUClipDeviceSelector": GPUClipDeviceSelector, +} + +NODE_CLASS_MAPPINGS["GPUCLIPLoader"] = override_class_clip(GLOBAL_NODE_CLASS_MAPPINGS["CLIPLoader"]) +NODE_CLASS_MAPPINGS["GPUDualCLIPLoader"] = override_class_clip(GLOBAL_NODE_CLASS_MAPPINGS["DualCLIPLoader"]) +NODE_CLASS_MAPPINGS["GPUTripleCLIPLoader"] = override_class_clip_no_device(GLOBAL_NODE_CLASS_MAPPINGS["TripleCLIPLoader"]) +NODE_CLASS_MAPPINGS["GPUQuadrupleCLIPLoader"] = override_class_clip_no_device(GLOBAL_NODE_CLASS_MAPPINGS["QuadrupleCLIPLoader"]) + +def register_and_count(module_names, node_map): + """Register MultiGPU node wrappers for detected custom node modules.""" + found = False + for name in module_names: + if check_module_exists(name): + found = True + break + + count = 0 + if found: + initial_len = len(NODE_CLASS_MAPPINGS) + for key, value in node_map.items(): + NODE_CLASS_MAPPINGS[key] = value + count = len(NODE_CLASS_MAPPINGS) - initial_len + + return found + +gguf_nodes = { + "GPUCLIPLoaderGGUF": override_class_clip(CLIPLoaderGGUF), + "GPUDualCLIPLoaderGGUF": override_class_clip(DualCLIPLoaderGGUF), + "GPUTripleCLIPLoaderGGUF": override_class_clip_no_device(TripleCLIPLoaderGGUF), + "GPUQuadrupleCLIPLoaderGGUF": override_class_clip_no_device(QuadrupleCLIPLoaderGGUF) +} +register_and_count(["ComfyUI-GGUF", "comfyui-gguf"], gguf_nodes) + +logger.info(f"[GPUClip] mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}") diff --git a/device_utils.py b/device_utils.py new file mode 100644 index 0000000..6ed00fc --- /dev/null +++ b/device_utils.py @@ -0,0 +1,155 @@ +import torch +import logging +import hashlib +import psutil +import comfy.model_management as mm +import gc + +logger = logging.getLogger("GPUClip") + +_DEVICE_LIST_CACHE = None + +def get_device_list(): + """ + Enumerate ALL physically available devices that can store torch tensors. + This includes all device types supported by ComfyUI core. + Results are cached after first call since devices don't change during runtime. + + Returns a comprehensive list of all available devices across all types: + - CPU (always available) + - CUDA devices (NVIDIA GPUs) + - XPU devices (Intel GPUs) + - NPU devices (Ascend NPUs from Huawei) + - MLU devices (Cambricon MLUs) + - MPS device (Apple Metal) + - DirectML devices (Windows DirectML) + - CoreX/IXUCA devices + """ + global _DEVICE_LIST_CACHE + + if _DEVICE_LIST_CACHE is not None: + return _DEVICE_LIST_CACHE + + devs = [] + + devs.append("cpu") + + if hasattr(torch, "cuda") and hasattr(torch.cuda, "is_available") and torch.cuda.is_available(): + device_count = torch.cuda.device_count() + devs += [f"cuda:{i}" for i in range(device_count)] + logger.debug(f"[MultiGPU_Device_Utils] Found {device_count} CUDA device(s)") + + try: + import intel_extension_for_pytorch as ipex + except ImportError: + pass + + if hasattr(torch, "xpu") and hasattr(torch.xpu, "is_available") and torch.xpu.is_available(): + device_count = torch.xpu.device_count() + devs += [f"xpu:{i}" for i in range(device_count)] + logger.debug(f"[MultiGPU_Device_Utils] Found {device_count} XPU device(s)") + + try: + import torch_npu + if hasattr(torch, "npu") and hasattr(torch.npu, "is_available") and torch.npu.is_available(): + device_count = torch.npu.device_count() + devs += [f"npu:{i}" for i in range(device_count)] + logger.debug(f"[MultiGPU_Device_Utils] Found {device_count} NPU device(s)") + except ImportError: + pass + + try: + import torch_mlu + if hasattr(torch, "mlu") and hasattr(torch.mlu, "is_available") and torch.mlu.is_available(): + device_count = torch.mlu.device_count() + devs += [f"mlu:{i}" for i in range(device_count)] + logger.debug(f"[MultiGPU_Device_Utils] Found {device_count} MLU device(s)") + except ImportError: + pass + + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + devs.append("mps") + logger.debug("[MultiGPU_Device_Utils] Found MPS device") + + try: + import torch_directml + adapter_count = torch_directml.device_count() + if adapter_count > 0: + devs += [f"directml:{i}" for i in range(adapter_count)] + logger.debug(f"[MultiGPU_Device_Utils] Found {adapter_count} DirectML adapter(s)") + except ImportError: + pass + + try: + if hasattr(torch, "corex"): + if hasattr(torch.corex, "device_count"): + device_count = torch.corex.device_count() + devs += [f"corex:{i}" for i in range(device_count)] + logger.debug(f"[MultiGPU_Device_Utils] Found {device_count} CoreX device(s)") + else: + devs.append("corex:0") + logger.debug("[MultiGPU_Device_Utils] Found CoreX device") + except ImportError: + pass + + _DEVICE_LIST_CACHE = devs + + logger.debug(f"[MultiGPU_Device_Utils] Device list initialized: {devs}") + + return devs + +def is_accelerator_available(): + """Check if any GPU or accelerator device is available including CUDA, XPU, NPU, MLU, MPS, DirectML, or CoreX.""" + if hasattr(torch, "cuda") and torch.cuda.is_available(): + return True + + if hasattr(torch, "xpu") and hasattr(torch.xpu, "is_available") and torch.xpu.is_available(): + return True + + try: + import torch_npu + if hasattr(torch, "npu") and hasattr(torch.npu, "is_available") and torch.npu.is_available(): + return True + except ImportError: + pass + + try: + import torch_mlu + if hasattr(torch, "mlu") and hasattr(torch.mlu, "is_available") and torch.mlu.is_available(): + return True + except ImportError: + pass + + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return True + + try: + import torch_directml + if torch_directml.device_count() > 0: + return True + except ImportError: + pass + + if hasattr(torch, "corex"): + return True + + return False + +def is_device_compatible(device_string): + """Check if a device string represents a valid available device.""" + available_devices = get_device_list() + return device_string in available_devices + +def get_device_type(device_string): + """Extract device type from device string (e.g. 'cuda' from 'cuda:0').""" + if ":" in device_string: + return device_string.split(":")[0] + return device_string + +def parse_device_string(device_string): + """Parse device string into (device_type, device_index) tuple.""" + if ":" in device_string: + parts = device_string.split(":") + return parts[0], int(parts[1]) + return device_string, None + diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..cd0faf5 --- /dev/null +++ b/nodes.py @@ -0,0 +1,127 @@ +import torch +import folder_paths +from nodes import NODE_CLASS_MAPPINGS +from .device_utils import get_device_list + +class GPUClipDeviceSelector: + @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 = "GPUCLIP" + + def select_device(self, device): + """Select target device from available device list.""" + return (device,) + +class CLIPLoaderGGUF: + @classmethod + def INPUT_TYPES(s): + import nodes + base = nodes.CLIPLoader.INPUT_TYPES() + return { + "required": { + "clip_name": (s.get_filename_list(),), + "type": base["required"]["type"], + } + } + + RETURN_TYPES = ("CLIP",) + FUNCTION = "load_clip" + CATEGORY = "bootleg" + TITLE = "CLIPLoader (GGUF)" + + @classmethod + def get_filename_list(s): + """Get combined list of CLIP and CLIP_GGUF model files.""" + 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): + """Load CLIP model data from checkpoint paths.""" + original_loader = NODE_CLASS_MAPPINGS["CLIPLoaderGGUF"]() + return original_loader.load_data(ckpt_paths) + + def load_patcher(self, clip_paths, clip_type, clip_data): + """Create ModelPatcher for CLIP model.""" + 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", device=None): + """Load CLIP model from GGUF or standard format.""" + original_loader = NODE_CLASS_MAPPINGS["CLIPLoaderGGUF"]() + return original_loader.load_clip(clip_name, type) + +class DualCLIPLoaderGGUF(CLIPLoaderGGUF): + @classmethod + def INPUT_TYPES(s): + import nodes + base = nodes.DualCLIPLoader.INPUT_TYPES() + file_options = (s.get_filename_list(), ) + return { + "required": { + "clip_name1": file_options, + "clip_name2": file_options, + "type": base["required"]["type"], + } + } + + TITLE = "DualCLIPLoader (GGUF)" + + def load_clip(self, clip_name1, clip_name2, type, device=None): + """Load dual CLIP model configuration.""" + 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 + + +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"): + """Load triple CLIP model configuration for SD3.""" + original_loader = NODE_CLASS_MAPPINGS["TripleCLIPLoaderGGUF"]() + return original_loader.load_clip(clip_name1, clip_name2, clip_name3, type) + +class QuadrupleCLIPLoaderGGUF(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, + "clip_name4": file_options, + } + } + + TITLE = "QuadrupleCLIPLoader (GGUF)" + + def load_clip(self, clip_name1, clip_name2, clip_name3, clip_name4, type="stable_diffusion"): + """Load quadruple CLIP model configuration.""" + original_loader = NODE_CLASS_MAPPINGS["QuadrupleCLIPLoaderGGUF"]() + return original_loader.load_clip(clip_name1, clip_name2, clip_name3, clip_name4, type) + diff --git a/wrappers.py b/wrappers.py new file mode 100644 index 0000000..80bdf80 --- /dev/null +++ b/wrappers.py @@ -0,0 +1,93 @@ +""" +ComfyUI-GPUClip Wrapper Functions +All node override/wrapper generation functions consolidated in one location +""" + +import copy +from .device_utils import get_device_list + + +# ============================================================================ +# STANDARD MULTIGPU WRAPPERS (Device selection without DisTorch) +# ============================================================================ + +def override_class(cls): + """Standard MultiGPU device override for UNet/VAE models""" + from . import set_current_device + + 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 = "GPUCLIP" + FUNCTION = "override" + + def override(self, *args, device=None, **kwargs): + if device is not None: + set_current_device(device) + fn = getattr(super(), cls.FUNCTION) + out = fn(*args, **kwargs) + return out + + return NodeOverride + + +def override_class_clip(cls): + """Standard MultiGPU device override for CLIP models (with device kwarg workaround)""" + from . import set_current_text_encoder_device + + 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 = "GPUCLIP" + FUNCTION = "override" + + def override(self, *args, device=None, **kwargs): + if device is not None: + set_current_text_encoder_device(device) + kwargs['device'] = 'default' + fn = getattr(super(), cls.FUNCTION) + out = fn(*args, **kwargs) + return out + + return NodeOverride + + +def override_class_clip_no_device(cls): + """Standard MultiGPU device override for Triple/Quad CLIP models (no device kwarg workaround)""" + from . import set_current_text_encoder_device + + 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 = "GPUCLIP" + FUNCTION = "override" + + def override(self, *args, device=None, **kwargs): + if device is not None: + set_current_text_encoder_device(device) + fn = getattr(super(), cls.FUNCTION) + out = fn(*args, **kwargs) + return out + + return NodeOverride