nodes/pollockjj-ComfyUI-MultiGPU
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
bd672479fa700e29b62da809c2b4bdda1f05bb42
pollockjj-ComfyUI-MultiGPU/__init__.py
T
John Pollock bd672479fa refactor: eliminate circular import by separating model management functions 
- Create model_management_mgpu.py for centralized model lifecycle tracking
- Move memory management functions from device_utils.py to new module:
  * multigpu_memory_log, track_modelpatcher, trigger_executor_cache_reset
  * check_cpu_memory_threshold, prune_distorch_stores, try_malloc_trim
  * force_full_system_cleanup
- Update imports across codebase (distorch_2.py, distorch.py, __init__.py,
  nodes.py, checkpoint_multigpu.py)
- Resolves device_utils.py ↔ distorch_2.py circular dependency
- Follows established clean coding patterns with fail-fast error handling

Addresses critical CPU memory leak investigation infrastructure by ensuring
proper module separation for comprehensive memory management utilities.
2025-09-24 17:38:15 -05:00

836 lines
39 KiB
Python
Raw Blame History

import torch
import logging
import weakref
import os
import copy
from pathlib import Path
import folder_paths
import comfy.model_management as mm
import comfy.model_patcher
from nodes import NODE_CLASS_MAPPINGS as GLOBAL_NODE_CLASS_MAPPINGS
from .device_utils import (
get_device_list,
is_accelerator_available,
soft_empty_cache_multigpu,
)
from .model_management_mgpu import (
trigger_executor_cache_reset,
check_cpu_memory_threshold,
multigpu_memory_log,
prune_distorch_stores,
try_malloc_trim,
track_modelpatcher,
force_full_system_cleanup,
)
# --- DisTorch V2 Logging Configuration ---
# Set to "E" for Engineering (DEBUG) or "P" for Production (INFO)
LOG_LEVEL = "P"
# Configure logger
logger = logging.getLogger("MultiGPU")
logger.propagate = False
if not logger.handlers:
log_level = logging.DEBUG if LOG_LEVEL == "E" else logging.INFO
handler = logging.StreamHandler()
formatter = logging.Formatter('%(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(log_level)
# --- MultiGPU Cleanup Policy Configuration ---
# Policy: off | threshold | every_load | every_load+threshold (alias threshold+every_load)
MGPU_CLEANUP_POLICY = os.getenv("MULTIGPU_CLEANUP_POLICY", "off").lower()
try:
MGPU_CPU_RESET_THRESHOLD = float(os.getenv("MULTIGPU_CPU_RESET_THRESHOLD", "0.85"))
except Exception:
MGPU_CPU_RESET_THRESHOLD = 0.85
# Malloc trim (not part of Comfy Core): on | off
MGPU_MALLOC_TRIM = os.getenv("MULTIGPU_MALLOC_TRIM", "on").lower()
logger.info(f"[MultiGPU Config] cleanup_policy={MGPU_CLEANUP_POLICY}, cpu_reset_threshold={MGPU_CPU_RESET_THRESHOLD:.2f}, malloc_trim={MGPU_MALLOC_TRIM}")
MGPU_MM_LOG = True
def mgpu_mm_log_method(self, msg):
if MGPU_MM_LOG:
self.info(f"[MultiGPU Model Management] {msg}")
logger.mgpu_mm_log = mgpu_mm_log_method.__get__(logger, type(logger))
# Global device state management
current_device = mm.get_torch_device()
current_text_encoder_device = mm.text_encoder_device()
def set_current_device(device):
global current_device
current_device = device
logger.debug(f"[MultiGPU Initialization] current_device set to: {device}")
def set_current_text_encoder_device(device):
global current_text_encoder_device
current_text_encoder_device = device
logger.debug(f"[MultiGPU Initialization] current_text_encoder_device set to: {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):
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):
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):
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):
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):
if device is not None:
set_current_text_encoder_device(device)
fn = getattr(super(), cls.FUNCTION)
out = fn(*args, **kwargs)
return out
return NodeOverride
def get_torch_device_patched():
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"[MultiGPU Core Patching] get_torch_device_patched returning device: {device} (current_device={current_device})")
return device
def text_encoder_device_patched():
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.debug(f"[MultiGPU Core Patching] text_encoder_device_patched returning device: {device} (current_text_encoder_device={current_text_encoder_device})")
return device
logger.info(f"[MultiGPU Core Patching] Patching mm.get_torch_device and mm.text_encoder_device")
logger.debug(f"[MultiGPU DEBUG] Initial current_device: {current_device}")
logger.debug(f"[MultiGPU DEBUG] Initial current_text_encoder_device: {current_text_encoder_device}")
mm.get_torch_device = get_torch_device_patched
mm.text_encoder_device = text_encoder_device_patched
# ==========================================================================================
# Core Patching: ModelPatcher Lifecycle Tracking (__init__)
# ==========================================================================================
logger.info("[MultiGPU Core Patching] Applying ModelPatcher lifecycle tracking patch (__init__).")
if not hasattr(comfy.model_patcher.ModelPatcher, '_mgpu_lifecycle_patched'):
try:
_mgpu_original_modelpatcher_init = comfy.model_patcher.ModelPatcher.__init__
def _mgpu_patched_modelpatcher_init(self, *args, **kwargs):
_mgpu_original_modelpatcher_init(self, *args, **kwargs)
# Track all ModelPatcher instances at construction time
try:
track_modelpatcher(self)
except Exception:
pass
comfy.model_patcher.ModelPatcher.__init__ = _mgpu_patched_modelpatcher_init
comfy.model_patcher.ModelPatcher._mgpu_lifecycle_patched = True
logger.info("[MultiGPU Core Patching] ModelPatcher.__init__ patched for lifecycle tracking.")
except Exception as e:
logger.error(f"[MultiGPU Core Patching] FAILED to patch ModelPatcher.__init__: {e}")
# ==========================================================================================
# Core Patching: Fix Potential Reference Cycles in LoadedModel
# ==========================================================================================
if hasattr(mm, 'LoadedModel') and hasattr(mm.LoadedModel, '_set_model'):
logger.info("[MultiGPU Core Patching] Patching mm.LoadedModel._set_model and _switch_parent to reduce reference cycles.")
_mgpu_original_set_model = mm.LoadedModel._set_model
def _mgpu_patched_set_model(self, model):
patcher_id = id(model)
# Ensure attributes exist
if not hasattr(self, '_model'):
self._model = None
if not hasattr(self, '_parent_model'):
self._parent_model = None
if not hasattr(self, '_patcher_finalizer'):
self._patcher_finalizer = None
# Reset refs
self._model = weakref.ref(model)
self._parent_model = None
# Detach any previous finalizer
if self._patcher_finalizer is not None:
try:
self._patcher_finalizer.detach()
except Exception:
pass
self._patcher_finalizer = None
# If clone, set parent and attach a weakref-based finalizer
parent = getattr(model, 'parent', None)
if parent is not None:
self._parent_model = weakref.ref(parent)
self_weak = weakref.ref(self)
def _mgpu_finalize_clone():
s = self_weak()
if s is not None and hasattr(s, '_switch_parent'):
logger.mgpu_mm_log(f"[MultiGPU_LoadedModel_Patch] Clone Patcher {patcher_id} GC'd. Switching LoadedModel to parent.")
s._switch_parent()
else:
logger.mgpu_mm_log(f"[MultiGPU_LoadedModel_Patch] Clone Patcher {patcher_id} GC'd. LoadedModel already gone or missing _switch_parent.")
try:
self._patcher_finalizer = weakref.finalize(model, _mgpu_finalize_clone)
except Exception:
self._patcher_finalizer = None
else:
logger.mgpu_mm_log(f"[MultiGPU_LoadedModel_Patch] Set base model Patcher {patcher_id}.")
mm.LoadedModel._set_model = _mgpu_patched_set_model
# Patch _switch_parent to clear references explicitly
if hasattr(mm.LoadedModel, '_switch_parent'):
_mgpu_original_switch_parent = mm.LoadedModel._switch_parent
def _mgpu_patched_switch_parent(self):
_mgpu_original_switch_parent(self)
# Clear parent and detach finalizer to avoid cycles
if hasattr(self, '_parent_model'):
self._parent_model = None
if hasattr(self, '_patcher_finalizer') and self._patcher_finalizer is not None:
try:
self._patcher_finalizer.detach()
except Exception:
pass
self._patcher_finalizer = None
mm.LoadedModel._switch_parent = _mgpu_patched_switch_parent
else:
# Fallback if core ever changes
def _mgpu_fallback_switch_parent(self):
if hasattr(self, '_parent_model') and self._parent_model is not None:
parent_model = self._parent_model()
if parent_model is not None:
self._set_model(parent_model)
self._parent_model = None
mm.LoadedModel._switch_parent = _mgpu_fallback_switch_parent
else:
logger.warning("[MultiGPU Core Patching] mm.LoadedModel not found or missing _set_model; skip cycle patch.")
def check_module_exists(module_path):
full_path = os.path.join(folder_paths.get_folder_paths("custom_nodes")[0], module_path)
logger.debug(f"[MultiGPU] Checking for module at {full_path}")
if not os.path.exists(full_path):
logger.debug(f"[MultiGPU] Module {module_path} not found - skipping")
return False
logger.debug(f"[MultiGPU] Found {module_path}, creating compatible MultiGPU nodes")
return True
# Import from nodes.py
from .nodes import (
DeviceSelectorMultiGPU,
HunyuanVideoEmbeddingsAdapter,
UnetLoaderGGUF,
UnetLoaderGGUFAdvanced,
CLIPLoaderGGUF,
DualCLIPLoaderGGUF,
TripleCLIPLoaderGGUF,
QuadrupleCLIPLoaderGGUF,
LTXVLoader,
Florence2ModelLoader,
DownloadAndLoadFlorence2Model,
CheckpointLoaderNF4,
LoadFluxControlNet,
MMAudioModelLoader,
MMAudioFeatureUtilsLoader,
MMAudioSampler,
PulidModelLoader,
PulidInsightFaceLoader,
PulidEvaClipLoader,
HyVideoModelLoader,
HyVideoVAELoader,
DownloadAndLoadHyVideoTextEncoder,
FullCleanupMultiGPU,
)
# Import from wanvideo.py
from .wanvideo import (
WanVideoModelLoader,
WanVideoModelLoader_2,
WanVideoVAELoader,
LoadWanVideoT5TextEncoder,
LoadWanVideoClipTextEncoder,
WanVideoTextEncode,
WanVideoBlockSwap,
WanVideoSampler
)
# Import from distorch.py
from .distorch import (
model_allocation_store,
create_model_hash,
register_patched_ggufmodelpatcher,
analyze_ggml_loading,
calculate_vvram_allocation_string,
override_class_with_distorch_gguf,
override_class_with_distorch_gguf_v2,
override_class_with_distorch_clip,
override_class_with_distorch_clip_no_device,
override_class_with_distorch
)
# Import from distorch_2.py for DisTorch v2 SafeTensor support
from .distorch_2 import (
safetensor_allocation_store,
create_safetensor_model_hash,
register_patched_safetensor_modelpatcher,
analyze_safetensor_loading,
calculate_safetensor_vvram_allocation,
override_class_with_distorch_safetensor_v2,
override_class_with_distorch_safetensor_v2_clip,
override_class_with_distorch_safetensor_v2_clip_no_device
)
logger.info("[MultiGPU Core Patching] Patching mm.soft_empty_cache for Comprehensive Memory Management (VRAM + CPU + Store Pruning)")
original_soft_empty_cache = mm.soft_empty_cache
def soft_empty_cache_distorch2_patched(force=False):
"""
Patched mm.soft_empty_cache.
- Prunes DisTorch store bookkeeping to avoid stale references
- Manages VRAM: if DisTorch2 models are active, clear allocator caches on all devices;
otherwise delegate to original mm.soft_empty_cache.
- Manages CPU RAM: adaptive threshold-based PromptExecutor cache reset;
and force-triggered reset when explicitly requested (mirrors ComfyUI 'Free memory' button).
"""
multigpu_memory_log("patched_soft_empty", f"start:force={force}")
# Prune DisTorch stores before any clearing to drop stale references
try:
prune_distorch_stores()
except Exception:
pass
is_distorch_active = False
# Detect DisTorch2-managed models
for lm in mm.current_loaded_models:
mp = lm.model # weakref call to ModelPatcher
if mp is not None:
model_hash = create_safetensor_model_hash(mp, "cache_patch_check")
if model_hash in safetensor_allocation_store and safetensor_allocation_store.get(model_hash):
is_distorch_active = True
break
# Phase 2: adaptive CPU memory management
check_cpu_memory_threshold()
# VRAM allocator management
if is_distorch_active:
logger.mgpu_mm_log("DisTorch2 active: clearing allocator caches on all devices (VRAM)")
soft_empty_cache_multigpu()
else:
logger.mgpu_mm_log("DisTorch2 not active: delegating allocator cache clear (VRAM) to original mm.soft_empty_cache")
original_soft_empty_cache(force)
# Optional: return CPU heap to OS (not part of Comfy Core)
if MGPU_MALLOC_TRIM != "off":
try:
try_malloc_trim()
except Exception:
pass
# Phase 1/3: forced executor reset mirrors ComfyUI 'Free memory' semantics
if force:
logger.mgpu_mm_log("Force flag active: triggering executor cache reset (CPU)")
trigger_executor_cache_reset(reason="forced_soft_empty", force=True)
multigpu_memory_log("patched_soft_empty", "end")
mm.soft_empty_cache = soft_empty_cache_distorch2_patched
LARGE_MODEL_THRESHOLD = 2 * (1024**3) # 2 GB threshold for "large" models
# Patch only once (handles reloads)
if hasattr(mm, 'load_models_gpu') and not hasattr(mm.load_models_gpu, "_distorch2_proactive_patched"):
logger.info("[MultiGPU Core Patching] Patching mm.load_models_gpu for DisTorch2 proactive unloading")
original_load_models_gpu = mm.load_models_gpu
def patched_load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None, force_full_load=False):
"""
Proactively unload large models that are not needed when loading a large DisTorch2 model.
This frees both compute and donor device memory ahead of ComfyUI's compute-only check.
"""
multigpu_memory_log("patched_load_models_gpu", "start")
# Validate models argument loudly
if not isinstance(models, (list, tuple, set)):
logger.error("[MultiGPU Core Patching] CRITICAL: mm.load_models_gpu 'models' is not a list/tuple/set. Bypassing proactive patch.")
return original_load_models_gpu(models, memory_required, force_patch_weights, minimum_memory_required, force_full_load)
# Detect incoming DisTorch2 request
incoming_is_distorch = False
incoming_distorch_nonzero = False
incoming_patchers = set()
incoming_loaded_names = []
incoming_allowed_devices = None
incoming_compute_device = None
incoming_required_bytes = 0
incoming_compute_planned_bytes = 0
for lm in models:
# Identify ModelPatcher (prefer direct; fall back to .patcher)
if hasattr(lm, "load_device"):
patcher = lm
elif hasattr(lm, "patcher"):
patcher = lm.patcher
else:
patcher = None
model_for_hash = patcher if patcher is not None else getattr(lm, "model", lm)
if patcher is not None:
incoming_patchers.add(patcher)
# Determine required memory directly from ModelPatcher (no wrapper; no side effects)
device_str = str(patcher.load_device)
if patcher.current_loaded_device() == patcher.load_device:
required_bytes = patcher.model_size() - patcher.loaded_size()
else:
required_bytes = patcher.model_size()
else:
device_str = "n/a"
required_bytes = 0
# Check DisTorch2 management via allocation store (unchanged trigger)
model_hash = create_safetensor_model_hash(model_for_hash, "load_patch_check")
if model_hash in safetensor_allocation_store and safetensor_allocation_store.get(model_hash):
incoming_is_distorch = True
if required_bytes > 0:
incoming_distorch_nonzero = True
if incoming_allowed_devices is None:
# Derive compute/donor devices from allocation string
alloc_str = safetensor_allocation_store.get(model_hash, "")
allowed = set()
if alloc_str:
parts = alloc_str.split("#", 1)
if len(parts) == 2 and parts[1]:
vram = parts[1]
segs = vram.split(";")
# compute device
if len(segs) >= 1 and segs[0]:
allowed.add(segs[0].strip())
# donors list (comma-separated)
if len(segs) >= 3 and segs[2]:
for d in segs[2].split(","):
d = d.strip()
if d:
allowed.add(d)
else:
# Expert fraction string: "dev,fraction;dev2,fraction2;..."
for token in alloc_str.split(";"):
if "," in token:
dev, frac = token.split(",", 1)
fs = frac.strip()
numlike = fs.replace(".", "", 1).isdigit()
if numlike and float(fs) > 0.0:
allowed.add(dev.strip())
if not allowed:
allowed = {str(patcher.load_device), "cpu"}
incoming_allowed_devices = allowed
# Determine compute device and planned bytes from allocation string
alloc = safetensor_allocation_store.get(model_hash, "")
if "#" in alloc:
vram = alloc.split("#", 1)[1]
segs = vram.split(";")
if len(segs) >= 2 and segs[0]:
incoming_compute_device = segs[0].strip()
try:
vvram_gb = float(segs[1])
incoming_compute_planned_bytes = int(vvram_gb * (1024**3))
except Exception:
incoming_compute_planned_bytes = 0
else:
# Expert fractions: "dev,fraction;dev2,fraction2;..."
tokens = [t for t in alloc.split(";") if "," in t]
frac_map = {}
for t in tokens:
dev, frac = t.split(",", 1)
try:
frac_val = float(frac.strip())
except Exception:
continue
frac_map[dev.strip()] = frac_val
if frac_map:
ld = str(patcher.load_device)
# Prefer the explicit load_device if present and > 0
target_dev = ld if (ld in frac_map and frac_map[ld] > 0.0) else None
if target_dev is None:
# Otherwise pick highest positive fraction
target_dev = max((d for d,v in frac_map.items() if v > 0.0), key=lambda d: frac_map[d], default=None)
if target_dev is not None:
incoming_compute_device = target_dev
total = mm.get_total_memory(torch.device(target_dev))
incoming_compute_planned_bytes = int(frac_map[target_dev] * (total or 0))
if incoming_compute_device is None:
incoming_compute_device = str(patcher.load_device)
if incoming_compute_planned_bytes <= 0:
incoming_compute_planned_bytes = required_bytes
# Log informational context with required bytes and device
try:
model_name = type(getattr(model_for_hash, "model", model_for_hash)).__name__
except Exception:
model_name = "UnknownModel"
incoming_loaded_names.append(f"{model_name}:{required_bytes/(1024**3):.2f}GB req on {device_str}")
if incoming_loaded_names:
logger.mgpu_mm_log(f"Incoming models summary: {', '.join(incoming_loaded_names)}")
if incoming_distorch_nonzero:
logger.mgpu_mm_log("Non-Zero incoming DisTorch2 model detected. Initiating proactive unload.")
# Proactively clear PromptExecutor caches ahead of major DisTorch2 load (Phase 1)
trigger_executor_cache_reset(reason="proactive_distorch_load", force=False)
if not hasattr(mm, 'current_loaded_models'):
raise AttributeError("comfy.model_management is missing 'current_loaded_models'. Proactive unload check failed.")
needed_patchers = incoming_patchers
# Need-based free on compute device only (scale-aware; core-aligned)
dev_str = incoming_compute_device or (next(iter(incoming_allowed_devices)) if incoming_allowed_devices else None)
freed_bytes = 0
to_unload_indices = []
unload_summaries = []
if dev_str is not None:
dev_obj = torch.device(dev_str)
free_now = mm.get_free_memory(dev_obj)
try:
free_now_val = free_now[0] if isinstance(free_now, tuple) else free_now
except Exception:
free_now_val = free_now
# Use core-aligned immediate needs: planned vs. memory_required vs. minimum_memory_required
effective_needed = max(incoming_compute_planned_bytes or 0, memory_required or 0, minimum_memory_required or 0)
need_bytes = max(0, effective_needed - (free_now_val or 0))
logger.mgpu_mm_log(f"Need calc on {dev_str}: effective_needed={effective_needed/(1024**3):.2f}GB, free_now={((free_now_val or 0)/(1024**3)):.2f}GB, need_bytes={need_bytes/(1024**3):.2f}GB")
if need_bytes > 0:
logger.mgpu_mm_log(f"Need-based unload on {dev_str}: need ~{need_bytes/(1024**3):.2f}GB")
# Build candidates on this device only, excluding needed patchers
candidates = []
for idx, lm_cur in enumerate(mm.current_loaded_models):
mp_cur = getattr(lm_cur, 'model', None)
if mp_cur is None or mp_cur in needed_patchers:
continue
if str(getattr(lm_cur, "device", "")) != dev_str:
continue
size_cur = 0
if hasattr(lm_cur, 'model_memory'):
try:
size_cur = lm_cur.model_memory()
except Exception:
size_cur = 0
if size_cur <= 0 and hasattr(mp_cur, 'model_size'):
size_cur = mp_cur.model_size()
candidates.append((size_cur, idx, lm_cur, mp_cur))
# Sort by size descending
candidates.sort(key=lambda x: x[0], reverse=True)
for size_cur, idx, lm_cur, mp_cur in candidates:
model_name = type(getattr(mp_cur, 'model', mp_cur)).__name__
logger.mgpu_mm_log(f"Unloading model on {dev_str}: {model_name} (~{size_cur/(1024**3):.2f}GB)")
success = False
if hasattr(lm_cur, 'model_unload'):
success = lm_cur.model_unload(memory_to_free=None, unpatch_weights=True)
if success:
to_unload_indices.append(idx)
unload_summaries.append(f"{model_name}:{size_cur/(1024**3):.2f}GB")
freed_bytes += size_cur
if freed_bytes >= need_bytes:
break
# Remove from management list and clear caches
unloaded_count = 0
for idx in sorted(to_unload_indices, reverse=True):
mm.current_loaded_models.pop(idx)
unloaded_count += 1
if unloaded_count > 0:
logger.mgpu_mm_log(f"Proactively unloaded {unloaded_count} large model(s): {', '.join(unload_summaries)}")
logger.mgpu_mm_log("Performing multi-device cache clear after proactive unload")
# Force multi-device cache clear via patched soft_empty_cache (which detects DisTorch2)
mm.soft_empty_cache(force=True)
else:
# Lineage-aligned cache clear when no unloads happened: apply core 25% rule, per DisTorch devices
if incoming_allowed_devices is not None and mm.vram_state != mm.VRAMState.HIGH_VRAM:
triggered = []
for dev_str in incoming_allowed_devices:
try:
dev_obj = torch.device(dev_str)
except Exception:
continue
free_total, free_torch = mm.get_free_memory(dev_obj, torch_free_too=True)
# free_total: system free; free_torch: torch reserved-but-free
if free_torch > free_total * 0.25:
triggered.append(dev_str)
if triggered:
logger.mgpu_mm_log(f"No unloads; 25% torch-cache rule triggered on: {', '.join(triggered)}. Calling soft_empty_cache()")
mm.soft_empty_cache(force=True)
else:
logger.mgpu_mm_log("No unloads; 25% torch-cache rule not met on DisTorch devices; skipping cache clear")
else:
logger.mgpu_mm_log("No unload candidates matched criteria and either HIGH_VRAM or no DisTorch devices; skipping cache clear")
elif incoming_is_distorch:
logger.mgpu_mm_log("Incoming DisTorch2 model requires 0.00GB; skipping proactive unload")
# Memory Logging
multigpu_memory_log("patched_load_models_gpu", "pre-original-call")
result = original_load_models_gpu(models, memory_required, force_patch_weights, minimum_memory_required, force_full_load)
multigpu_memory_log("patched_load_models_gpu", "post-original-call")
# Cleanup policy triggers (flags-only, Manager semantics)
if MGPU_CLEANUP_POLICY in ("threshold", "every_load+threshold", "threshold+every_load"):
try:
check_cpu_memory_threshold(threshold_percent=MGPU_CPU_RESET_THRESHOLD * 100.0)
except Exception:
pass
if MGPU_CLEANUP_POLICY in ("every_load", "every_load+threshold", "threshold+every_load"):
try:
# flags-only; prompt worker performs unload/reset/gc
force_full_system_cleanup(reason="policy_every_load", force=False)
except Exception:
pass
return result
# Mark and apply the patch
patched_load_models_gpu._distorch2_proactive_patched = True
mm.load_models_gpu = patched_load_models_gpu
else:
if not hasattr(mm, 'load_models_gpu'):
raise AttributeError("comfy.model_management is missing 'load_models_gpu'. Core patching failed.")
else:
logger.debug("[MultiGPU Core Patching] mm.load_models_gpu already patched; skipping")
# Import advanced checkpoint loaders
from .checkpoint_multigpu import (
CheckpointLoaderAdvancedMultiGPU,
CheckpointLoaderAdvancedDisTorch2MultiGPU
)
# Initialize NODE_CLASS_MAPPINGS
NODE_CLASS_MAPPINGS = {
"DeviceSelectorMultiGPU": DeviceSelectorMultiGPU,
"HunyuanVideoEmbeddingsAdapter": HunyuanVideoEmbeddingsAdapter,
"CheckpointLoaderAdvancedMultiGPU": CheckpointLoaderAdvancedMultiGPU,
"CheckpointLoaderAdvancedDisTorch2MultiGPU": CheckpointLoaderAdvancedDisTorch2MultiGPU,
}
# Standard MultiGPU nodes
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_clip(GLOBAL_NODE_CLASS_MAPPINGS["CLIPLoader"])
NODE_CLASS_MAPPINGS["DualCLIPLoaderMultiGPU"] = override_class_clip(GLOBAL_NODE_CLASS_MAPPINGS["DualCLIPLoader"])
if "TripleCLIPLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
NODE_CLASS_MAPPINGS["TripleCLIPLoaderMultiGPU"] = override_class_clip_no_device(GLOBAL_NODE_CLASS_MAPPINGS["TripleCLIPLoader"])
if "QuadrupleCLIPLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
NODE_CLASS_MAPPINGS["QuadrupleCLIPLoaderMultiGPU"] = override_class_clip_no_device(GLOBAL_NODE_CLASS_MAPPINGS["QuadrupleCLIPLoader"])
NODE_CLASS_MAPPINGS["CLIPVisionLoaderMultiGPU"] = override_class_clip_no_device(GLOBAL_NODE_CLASS_MAPPINGS["CLIPVisionLoader"])
NODE_CLASS_MAPPINGS["CheckpointLoaderSimpleMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"])
NODE_CLASS_MAPPINGS["ControlNetLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["ControlNetLoader"])
if "DiffusersLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
NODE_CLASS_MAPPINGS["DiffusersLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["DiffusersLoader"])
if "DiffControlNetLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
NODE_CLASS_MAPPINGS["DiffControlNetLoaderMultiGPU"] = override_class(GLOBAL_NODE_CLASS_MAPPINGS["DiffControlNetLoader"])
# DisTorch 2 SafeTensor nodes for FLUX and other safetensor models
NODE_CLASS_MAPPINGS["UNETLoaderDisTorch2MultiGPU"] = override_class_with_distorch_safetensor_v2(GLOBAL_NODE_CLASS_MAPPINGS["UNETLoader"])
NODE_CLASS_MAPPINGS["VAELoaderDisTorch2MultiGPU"] = override_class_with_distorch_safetensor_v2(GLOBAL_NODE_CLASS_MAPPINGS["VAELoader"])
NODE_CLASS_MAPPINGS["CLIPLoaderDisTorch2MultiGPU"] = override_class_with_distorch_safetensor_v2_clip(GLOBAL_NODE_CLASS_MAPPINGS["CLIPLoader"])
NODE_CLASS_MAPPINGS["DualCLIPLoaderDisTorch2MultiGPU"] = override_class_with_distorch_safetensor_v2_clip(GLOBAL_NODE_CLASS_MAPPINGS["DualCLIPLoader"])
if "TripleCLIPLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
NODE_CLASS_MAPPINGS["TripleCLIPLoaderDisTorch2MultiGPU"] = override_class_with_distorch_safetensor_v2_clip_no_device(GLOBAL_NODE_CLASS_MAPPINGS["TripleCLIPLoader"])
if "QuadrupleCLIPLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
NODE_CLASS_MAPPINGS["QuadrupleCLIPLoaderDisTorch2MultiGPU"] = override_class_with_distorch_safetensor_v2_clip_no_device(GLOBAL_NODE_CLASS_MAPPINGS["QuadrupleCLIPLoader"])
NODE_CLASS_MAPPINGS["CLIPVisionLoaderDisTorch2MultiGPU"] = override_class_with_distorch_safetensor_v2_clip_no_device(GLOBAL_NODE_CLASS_MAPPINGS["CLIPVisionLoader"])
NODE_CLASS_MAPPINGS["CheckpointLoaderSimpleDisTorch2MultiGPU"] = override_class_with_distorch_safetensor_v2(GLOBAL_NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"])
NODE_CLASS_MAPPINGS["ControlNetLoaderDisTorch2MultiGPU"] = override_class_with_distorch_safetensor_v2(GLOBAL_NODE_CLASS_MAPPINGS["ControlNetLoader"])
if "DiffusersLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
NODE_CLASS_MAPPINGS["DiffusersLoaderDisTorch2MultiGPU"] = override_class_with_distorch_safetensor_v2(GLOBAL_NODE_CLASS_MAPPINGS["DiffusersLoader"])
if "DiffControlNetLoader" in GLOBAL_NODE_CLASS_MAPPINGS:
NODE_CLASS_MAPPINGS["DiffControlNetLoaderDisTorch2MultiGPU"] = override_class_with_distorch_safetensor_v2(GLOBAL_NODE_CLASS_MAPPINGS["DiffControlNetLoader"])
# --- Registration Table ---
logger.info("[MultiGPU] Initiating custom_node Registration. . .")
dash_line = "-" * 47
fmt_reg = "{:<30}{:>5}{:>10}"
logger.info(dash_line)
logger.info(fmt_reg.format("custom_node", "Found", "Nodes"))
logger.info(dash_line)
registration_data = []
def register_and_count(module_names, node_map):
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
registration_data.append({"name": module_names[0], "found": "Y" if found else "N", "count": count})
return found
# ComfyUI-LTXVideo
ltx_nodes = {"LTXVLoaderMultiGPU": override_class(LTXVLoader)}
register_and_count(["ComfyUI-LTXVideo", "comfyui-ltxvideo"], ltx_nodes)
# ComfyUI-Florence2
florence_nodes = {
"Florence2ModelLoaderMultiGPU": override_class(Florence2ModelLoader),
"DownloadAndLoadFlorence2ModelMultiGPU": override_class(DownloadAndLoadFlorence2Model)
}
register_and_count(["ComfyUI-Florence2", "comfyui-florence2"], florence_nodes)
# ComfyUI_bitsandbytes_NF4
nf4_nodes = {"CheckpointLoaderNF4MultiGPU": override_class(CheckpointLoaderNF4)}
register_and_count(["ComfyUI_bitsandbytes_NF4", "comfyui_bitsandbytes_nf4"], nf4_nodes)
# x-flux-comfyui
flux_controlnet_nodes = {"LoadFluxControlNetMultiGPU": override_class(LoadFluxControlNet)}
register_and_count(["x-flux-comfyui"], flux_controlnet_nodes)
# ComfyUI-MMAudio
mmaudio_nodes = {
"MMAudioModelLoaderMultiGPU": override_class(MMAudioModelLoader),
"MMAudioFeatureUtilsLoaderMultiGPU": override_class(MMAudioFeatureUtilsLoader),
"MMAudioSamplerMultiGPU": override_class(MMAudioSampler)
}
register_and_count(["ComfyUI-MMAudio", "comfyui-mmaudio"], mmaudio_nodes)
# ComfyUI-GGUF
gguf_nodes = {
"UnetLoaderGGUFDisTorchMultiGPU": override_class_with_distorch_gguf(UnetLoaderGGUF),
"UnetLoaderGGUFAdvancedDisTorchMultiGPU": override_class_with_distorch_gguf(UnetLoaderGGUFAdvanced),
"CLIPLoaderGGUFDisTorchMultiGPU": override_class_with_distorch_clip(CLIPLoaderGGUF),
"DualCLIPLoaderGGUFDisTorchMultiGPU": override_class_with_distorch_clip(DualCLIPLoaderGGUF),
"TripleCLIPLoaderGGUFDisTorchMultiGPU": override_class_with_distorch_clip_no_device(TripleCLIPLoaderGGUF),
"QuadrupleCLIPLoaderGGUFDisTorchMultiGPU": override_class_with_distorch_clip_no_device(QuadrupleCLIPLoaderGGUF),
"UnetLoaderGGUFDisTorch2MultiGPU": override_class_with_distorch_safetensor_v2(UnetLoaderGGUF),
"UnetLoaderGGUFAdvancedDisTorch2MultiGPU": override_class_with_distorch_safetensor_v2(UnetLoaderGGUFAdvanced),
"CLIPLoaderGGUFDisTorch2MultiGPU": override_class_with_distorch_safetensor_v2_clip(CLIPLoaderGGUF),
"DualCLIPLoaderGGUFDisTorch2MultiGPU": override_class_with_distorch_safetensor_v2_clip(DualCLIPLoaderGGUF),
"TripleCLIPLoaderGGUFDisTorch2MultiGPU": override_class_with_distorch_safetensor_v2_clip_no_device(TripleCLIPLoaderGGUF),
"QuadrupleCLIPLoaderGGUFDisTorch2MultiGPU": override_class_with_distorch_safetensor_v2_clip_no_device(QuadrupleCLIPLoaderGGUF),
"UnetLoaderGGUFMultiGPU": override_class(UnetLoaderGGUF),
"UnetLoaderGGUFAdvancedMultiGPU": override_class(UnetLoaderGGUFAdvanced),
"CLIPLoaderGGUFMultiGPU": override_class_clip(CLIPLoaderGGUF),
"DualCLIPLoaderGGUFMultiGPU": override_class_clip(DualCLIPLoaderGGUF),
"TripleCLIPLoaderGGUFMultiGPU": override_class_clip_no_device(TripleCLIPLoaderGGUF),
"QuadrupleCLIPLoaderGGUFMultiGPU": override_class_clip_no_device(QuadrupleCLIPLoaderGGUF)
}
register_and_count(["ComfyUI-GGUF", "comfyui-gguf"], gguf_nodes)
# PuLID_ComfyUI
pulid_nodes = {
"PulidModelLoaderMultiGPU": override_class(PulidModelLoader),
"PulidInsightFaceLoaderMultiGPU": override_class(PulidInsightFaceLoader),
"PulidEvaClipLoaderMultiGPU": override_class(PulidEvaClipLoader)
}
register_and_count(["PuLID_ComfyUI", "pulid_comfyui"], pulid_nodes)
# ComfyUI-HunyuanVideoWrapper
hunyuan_nodes = {
"HyVideoModelLoaderMultiGPU": override_class(HyVideoModelLoader),
"HyVideoVAELoaderMultiGPU": override_class(HyVideoVAELoader),
"DownloadAndLoadHyVideoTextEncoderMultiGPU": override_class(DownloadAndLoadHyVideoTextEncoder)
}
register_and_count(["ComfyUI-HunyuanVideoWrapper", "comfyui-hunyuanvideowrapper"], hunyuan_nodes)
# ComfyUI-WanVideoWrapper
wanvideo_nodes = {
"WanVideoModelLoaderMultiGPU": WanVideoModelLoader,
"WanVideoModelLoaderMultiGPU_2": WanVideoModelLoader_2,
"WanVideoVAELoaderMultiGPU": WanVideoVAELoader,
"LoadWanVideoT5TextEncoderMultiGPU": LoadWanVideoT5TextEncoder,
"LoadWanVideoClipTextEncoderMultiGPU": LoadWanVideoClipTextEncoder,
"WanVideoTextEncodeMultiGPU": WanVideoTextEncode,
"WanVideoBlockSwapMultiGPU": WanVideoBlockSwap,
"WanVideoSamplerMultiGPU": WanVideoSampler
}
register_and_count(["ComfyUI-WanVideoWrapper", "comfyui-wanvideowrapper"], wanvideo_nodes)
# Print the registration table
for item in registration_data:
logger.info(fmt_reg.format(item['name'], item['found'], str(item['count'])))
logger.info(dash_line)
# Register maintenance node
NODE_CLASS_MAPPINGS["FullCleanupMultiGPU"] = FullCleanupMultiGPU
logger.info(f"[MultiGPU] Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}")
Reference in New Issue View Git Blame Copy Permalink