Add files via upload

2026-02-06 11:28:48 +02:00
parent b712149195
commit 1173245de9
4 changed files with 499 additions and 0 deletions
@@ -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())}")
@@ -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
+
@@ -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)
+
@@ -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