From e50f4d2c8a5b95c3c21417ee55f65d559f721d7d Mon Sep 17 00:00:00 2001 From: Nathan Sapwell Date: Sat, 10 Jan 2026 08:40:46 +0800 Subject: [PATCH] fix: compatibility with existing multigpu packages --- __init__.py | 216 +++++++++++++++++++++++++++++++++++----------------- 1 file changed, 145 insertions(+), 71 deletions(-) diff --git a/__init__.py b/__init__.py index 23a7b85..1acc4b3 100644 --- a/__init__.py +++ b/__init__.py @@ -10,18 +10,150 @@ # Device utilities based on ComfyUI-MultiGPU by pollockjj (GPL-3.0) import logging -import torch -import comfy.model_management as mm +import sys + +logger = logging.getLogger("LTXV2MultiGPU") + +# ============================================================================ +# Detect if ComfyUI-MultiGPU is installed and use its infrastructure +# ============================================================================ + +_USING_MULTIGPU = False + +# Check for ComfyUI-MultiGPU module +for module_name in list(sys.modules.keys()): + if 'comfyui-multigpu' in module_name.lower() or module_name.startswith('custom_nodes.ComfyUI-MultiGPU'): + _USING_MULTIGPU = True + break + +# Also try direct import +if not _USING_MULTIGPU: + try: + # Try importing from custom_nodes path patterns + import importlib.util + import os + + # Check common custom_nodes locations + possible_paths = [] + for path in sys.path: + if 'custom_nodes' in path: + multigpu_path = os.path.join(os.path.dirname(path), 'ComfyUI-MultiGPU', '__init__.py') + if os.path.exists(multigpu_path): + _USING_MULTIGPU = True + break + # Also check sibling + multigpu_path = os.path.join(path, '..', 'ComfyUI-MultiGPU', '__init__.py') + if os.path.exists(os.path.normpath(multigpu_path)): + _USING_MULTIGPU = True + break + except Exception: + pass + +if _USING_MULTIGPU: + # ======================================================================== + # USE COMFYUI-MULTIGPU INFRASTRUCTURE + # ======================================================================== + logger.info("[LTXV2 MultiGPU] ComfyUI-MultiGPU detected, using its device management") + + try: + # Import from ComfyUI-MultiGPU + from custom_nodes import ComfyUI_MultiGPU as multigpu + + # Use their wrappers and device functions + override_class = multigpu.override_class + override_class_clip = multigpu.override_class_clip + get_device_list = multigpu.get_device_list + + logger.info("[LTXV2 MultiGPU] Successfully imported ComfyUI-MultiGPU functions") + except ImportError: + # Fallback: try alternate import path + try: + # Some installations use different module names + for mod_name, mod in sys.modules.items(): + if 'multigpu' in mod_name.lower() and hasattr(mod, 'override_class'): + override_class = mod.override_class + override_class_clip = mod.override_class_clip + get_device_list = mod.get_device_list + logger.info(f"[LTXV2 MultiGPU] Found ComfyUI-MultiGPU via {mod_name}") + break + else: + raise ImportError("Could not find ComfyUI-MultiGPU exports") + except Exception as e: + logger.warning(f"[LTXV2 MultiGPU] Failed to import from ComfyUI-MultiGPU: {e}") + logger.warning("[LTXV2 MultiGPU] Falling back to standalone mode") + _USING_MULTIGPU = False + +if not _USING_MULTIGPU: + # ======================================================================== + # STANDALONE MODE - Apply our own patches + # ======================================================================== + logger.info("[LTXV2 MultiGPU] Running in standalone mode") + + import torch + import comfy.model_management as mm + + from .device_utils import get_device_list + from .wrappers import ( + override_class, + override_class_clip, + get_current_device, + set_current_device, + get_current_text_encoder_device, + set_current_text_encoder_device, + ) + + # Initialize device state from ComfyUI defaults + _initial_device = mm.get_torch_device() + _initial_text_encoder_device = mm.text_encoder_device() + + set_current_device(_initial_device) + set_current_text_encoder_device(_initial_text_encoder_device) + + # Store original functions + _original_get_torch_device = mm.get_torch_device + _original_text_encoder_device = mm.text_encoder_device + + def get_torch_device_patched(): + """Return the user-selected device instead of ComfyUI default.""" + current = get_current_device() + if current is None: + return _original_get_torch_device() + + devs = get_device_list() + current_str = str(current) + + if mm.cpu_state == mm.CPUState.CPU or "cpu" in current_str.lower(): + return torch.device("cpu") + + if current_str in devs: + return torch.device(current_str) + + return _original_get_torch_device() + + def text_encoder_device_patched(): + """Return the user-selected text encoder device.""" + current = get_current_text_encoder_device() + if current is None: + return _original_text_encoder_device() + + devs = get_device_list() + current_str = str(current) + + if current_str in devs: + return torch.device(current_str) + + return _original_text_encoder_device() + + # Apply patches + mm.get_torch_device = get_torch_device_patched + mm.text_encoder_device = text_encoder_device_patched + + logger.info("[LTXV2 MultiGPU] Patched comfy.model_management for device selection") + +# ============================================================================ +# Import nodes (always needed) +# ============================================================================ -from .device_utils import get_device_list -from .wrappers import ( - override_class, - override_class_clip, - get_current_device, - set_current_device, - get_current_text_encoder_device, - set_current_text_encoder_device, -) from .nodes import ( LTXV2AudioVAELoader, LTXV2AVTextEncoderLoader, @@ -29,65 +161,6 @@ from .nodes import ( LTXV2CheckpointLoader, ) -logger = logging.getLogger("LTXV2MultiGPU") - -# ============================================================================ -# Initialize device state from ComfyUI defaults -# ============================================================================ - -_initial_device = mm.get_torch_device() -_initial_text_encoder_device = mm.text_encoder_device() - -set_current_device(_initial_device) -set_current_text_encoder_device(_initial_text_encoder_device) - -# ============================================================================ -# Patch ComfyUI model_management to use our device selection -# ============================================================================ - -_original_get_torch_device = mm.get_torch_device -_original_text_encoder_device = mm.text_encoder_device - - -def get_torch_device_patched(): - """Return the user-selected device instead of ComfyUI default.""" - current = get_current_device() - if current is None: - return _original_get_torch_device() - - devs = get_device_list() - current_str = str(current) - - if mm.cpu_state == mm.CPUState.CPU or "cpu" in current_str.lower(): - return torch.device("cpu") - - if current_str in devs: - return torch.device(current_str) - - return _original_get_torch_device() - - -def text_encoder_device_patched(): - """Return the user-selected text encoder device.""" - current = get_current_text_encoder_device() - if current is None: - return _original_text_encoder_device() - - devs = get_device_list() - current_str = str(current) - - if current_str in devs: - return torch.device(current_str) - - return _original_text_encoder_device() - - -# Apply patches -mm.get_torch_device = get_torch_device_patched -mm.text_encoder_device = text_encoder_device_patched - -logger.info("[LTXV2 MultiGPU] Patched comfy.model_management for device selection") - # ============================================================================ # Register nodes # ============================================================================ @@ -107,6 +180,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { } logger.info(f"[LTXV2 MultiGPU] Registered {len(NODE_CLASS_MAPPINGS)} nodes") -logger.info(f"[LTXV2 MultiGPU] Available devices: {get_device_list()}") +if not _USING_MULTIGPU: + logger.info(f"[LTXV2 MultiGPU] Available devices: {get_device_list()}") __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]