From d8616acd5e6ce7da089e4ebcbad39b170afeddc8 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sat, 4 Oct 2025 05:29:32 -0500 Subject: [PATCH 1/7] investigation --- __init__.py | 33 ++++++++++++++++++++++++--------- distorch_2.py | 4 ++-- 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/__init__.py b/__init__.py index 911087b..f9a1795 100644 --- a/__init__.py +++ b/__init__.py @@ -64,7 +64,7 @@ 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"[MultiGPU Initialization] current_text_encoder_device set to: {device}") + logger.info(f"[MultiGPU Initialization] current_text_encoder_device set to: {device}") def get_torch_device_patched(): """Return MultiGPU-aware device selection for patched mm.get_torch_device.""" @@ -77,22 +77,37 @@ def get_torch_device_patched(): logger.debug(f"[MultiGPU 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.""" +def _get_patched_text_encoder_device(): + """Internal helper to get the patched 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") + logger.info(f"[_get_patched_text_encoder_device] Condition met: accelerator not available, CPU state, or 'cpu' in current device. Returning 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})") + is_current_in_devs = str(current_text_encoder_device) in devs + device = torch.device(current_text_encoder_device) if is_current_in_devs else torch.device("cpu") + logger.info(f"[_get_patched_text_encoder_device] Available devices: {devs}. Current text encoder device: {current_text_encoder_device}. Is in list: {is_current_in_devs}. Returning: {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}") +def text_encoder_device_patched(): + """Return MultiGPU-aware text encoder device for patched mm.text_encoder_device.""" + device = _get_patched_text_encoder_device() + logger.info(f"[MultiGPU Core Patching] text_encoder_device_patched returning device: {device} (current_text_encoder_device={current_text_encoder_device})") + return device + +def text_encoder_initial_device_patched(load_device, offload_device, model_size=0): + """Return MultiGPU-aware initial text encoder device for patched mm.text_encoder_initial_device.""" + device = _get_patched_text_encoder_device() + logger.info(f"[MultiGPU Core Patching] text_encoder_initial_device_patched returning device: {device} (ignoring original args: load_device='{load_device}', offload_device='{offload_device}', model_size='{model_size}')") + return device + +logger.info(f"[MultiGPU Core Patching] Patching mm.get_torch_device, mm.text_encoder_device, and mm.text_encoder_initial_device") +logger.info(f"[MultiGPU INFO] Initial current_device: {current_device}") +logger.info(f"[MultiGPU INFO] 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 +mm.text_encoder_initial_device = text_encoder_initial_device_patched from .nodes import ( DeviceSelectorMultiGPU, @@ -289,4 +304,4 @@ for item in registration_data: logger.info(fmt_reg.format(item['name'], item['found'], str(item['count']))) logger.info(dash_line) -logger.info(f"[MultiGPU] Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}") \ No newline at end of file +logger.info(f"[MultiGPU] Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}") diff --git a/distorch_2.py b/distorch_2.py index 250a07e..22b2308 100644 --- a/distorch_2.py +++ b/distorch_2.py @@ -156,7 +156,7 @@ def register_patched_safetensor_modelpatcher(): pass if current_module_device is not None and str(current_module_device) != str(block_target_device): - logger.debug(f"[MultiGPU DisTorch V2] Moving already patched {module_name} to {block_target_device}") + logger.info(f"[MultiGPU DisTorch V2] Moving already patched {module_name} to {block_target_device}") module_object.to(block_target_device) mem_counter += module_size @@ -194,7 +194,7 @@ def register_patched_safetensor_modelpatcher(): # Step 4: Move to ultimate destination based on DisTorch assignment if block_target_device != device_to: - logger.debug(f"[MultiGPU DisTorch V2] Moving {module_name} from {device_to} to {block_target_device}") + logger.info(f"[MultiGPU DisTorch V2] Moving {module_name} from {device_to} to {block_target_device}") module_object.to(block_target_device) module_object.comfy_cast_weights = True From ee41f46beb0dfe5b221d2791d88ebce0d0b39df0 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sat, 4 Oct 2025 06:47:08 -0500 Subject: [PATCH 2/7] revert most changes --- __init__.py | 31 ++++++++----------------------- 1 file changed, 8 insertions(+), 23 deletions(-) diff --git a/__init__.py b/__init__.py index f9a1795..70cd537 100644 --- a/__init__.py +++ b/__init__.py @@ -64,7 +64,7 @@ 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.info(f"[MultiGPU Initialization] current_text_encoder_device set to: {device}") + logger.debug(f"[MultiGPU Initialization] current_text_encoder_device set to: {device}") def get_torch_device_patched(): """Return MultiGPU-aware device selection for patched mm.get_torch_device.""" @@ -77,37 +77,22 @@ def get_torch_device_patched(): logger.debug(f"[MultiGPU Core Patching] get_torch_device_patched returning device: {device} (current_device={current_device})") return device -def _get_patched_text_encoder_device(): - """Internal helper to get the patched text encoder 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") - logger.info(f"[_get_patched_text_encoder_device] Condition met: accelerator not available, CPU state, or 'cpu' in current device. Returning CPU.") else: devs = set(get_device_list()) - is_current_in_devs = str(current_text_encoder_device) in devs - device = torch.device(current_text_encoder_device) if is_current_in_devs else torch.device("cpu") - logger.info(f"[_get_patched_text_encoder_device] Available devices: {devs}. Current text encoder device: {current_text_encoder_device}. Is in list: {is_current_in_devs}. Returning: {device}") - return device - -def text_encoder_device_patched(): - """Return MultiGPU-aware text encoder device for patched mm.text_encoder_device.""" - device = _get_patched_text_encoder_device() + device = torch.device(current_text_encoder_device) if str(current_text_encoder_device) in devs else torch.device("cpu") logger.info(f"[MultiGPU Core Patching] text_encoder_device_patched returning device: {device} (current_text_encoder_device={current_text_encoder_device})") return device -def text_encoder_initial_device_patched(load_device, offload_device, model_size=0): - """Return MultiGPU-aware initial text encoder device for patched mm.text_encoder_initial_device.""" - device = _get_patched_text_encoder_device() - logger.info(f"[MultiGPU Core Patching] text_encoder_initial_device_patched returning device: {device} (ignoring original args: load_device='{load_device}', offload_device='{offload_device}', model_size='{model_size}')") - return device - -logger.info(f"[MultiGPU Core Patching] Patching mm.get_torch_device, mm.text_encoder_device, and mm.text_encoder_initial_device") -logger.info(f"[MultiGPU INFO] Initial current_device: {current_device}") -logger.info(f"[MultiGPU INFO] Initial current_text_encoder_device: {current_text_encoder_device}") +logger.info(f"[MultiGPU Core Patching] Patching mm.get_torch_device and mm.text_encoder_device") +logger.info(f"[MultiGPU DEBUG] Initial current_device: {current_device}") +logger.info(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 -mm.text_encoder_initial_device = text_encoder_initial_device_patched from .nodes import ( DeviceSelectorMultiGPU, @@ -304,4 +289,4 @@ for item in registration_data: logger.info(fmt_reg.format(item['name'], item['found'], str(item['count']))) logger.info(dash_line) -logger.info(f"[MultiGPU] Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}") +logger.info(f"[MultiGPU] Registration complete. Final mappings: {', '.join(NODE_CLASS_MAPPINGS.keys())}") \ No newline at end of file From e3750fd737c4c77e1551d15831b0d47672eb7b73 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sat, 4 Oct 2025 09:27:33 -0500 Subject: [PATCH 3/7] minor changes --- __init__.py | 2 +- distorch_2.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/__init__.py b/__init__.py index 70cd537..40b71c1 100644 --- a/__init__.py +++ b/__init__.py @@ -21,7 +21,7 @@ from .model_management_mgpu import ( ) WEB_DIRECTORY = "./web" -MGPU_MM_LOG = False +MGPU_MM_LOG = True DEBUG_LOG = False logger = logging.getLogger("MultiGPU") diff --git a/distorch_2.py b/distorch_2.py index 22b2308..ad73dc2 100644 --- a/distorch_2.py +++ b/distorch_2.py @@ -194,7 +194,7 @@ def register_patched_safetensor_modelpatcher(): # Step 4: Move to ultimate destination based on DisTorch assignment if block_target_device != device_to: - logger.info(f"[MultiGPU DisTorch V2] Moving {module_name} from {device_to} to {block_target_device}") + logger.debug(f"[MultiGPU DisTorch V2] Moving {module_name} from {device_to} to {block_target_device}") module_object.to(block_target_device) module_object.comfy_cast_weights = True From e6d19951d7b797ce5554bdd49791b550025c7723 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sat, 4 Oct 2025 12:32:37 -0500 Subject: [PATCH 4/7] Paranoia before cleanup --- distorch_2.py | 89 +++++++++++++++++------ model_management_mgpu.py | 153 --------------------------------------- wrappers.py | 84 ++++++++++++--------- 3 files changed, 118 insertions(+), 208 deletions(-) diff --git a/distorch_2.py b/distorch_2.py index ad73dc2..733d677 100644 --- a/distorch_2.py +++ b/distorch_2.py @@ -74,33 +74,78 @@ def register_patched_safetensor_modelpatcher(): original_loaded_model_memory_required = mm.LoadedModel.model_memory_required def patched_loaded_model_memory_required(self, device): - """Drive unload behavior purely by unload_distorch_model flag""" + """Truth table for memory reporting: + eject_models=0, is_distorch=0: return original + eject_models=0, is_distorch=1: return original - virtual_vram_gb_bytes + eject_models=1, is_distorch=0: mutually exclusive (shouldn't occur) + eject_models=1, is_distorch=1: return MAX memory to force eviction""" multigpu_memory_log("unload_distorch_model_memory_check", "start") - logger.mgpu_mm_log(f"[IS_DISTORCH_MODEL] Memory assessment requested for model on device: {device}") - - # Check if this is a DisTorch model with unload_distorch_model flag - is_distorch_model = hasattr(getattr(getattr(self, 'model', None), 'model', None), '_mgpu_unload_distorch_model') - model_name = type(getattr(getattr(self, 'model', None), 'model', None)).__name__ if getattr(getattr(self, 'model', None), 'model', None) else "Unknown" - logger.mgpu_mm_log(f"[IS_DISTORCH_MODEL] DisTorch model: {model_name}, is_distorch_model={is_distorch_model}") + logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] Memory assessment requested for model on device: {device}") - if is_distorch_model: - if self.model.model._mgpu_unload_distorch_model: - total_device_memory = mm.get_total_memory(device) - memory_gb = total_device_memory / (1024**3) - logger.mgpu_mm_log(f"[IS_DISTORCH_MODEL] _mgpu_unload_distorch_model=True - Reporting MAX memory ({memory_gb:.2f}GB) to force complete eviction") - return total_device_memory - else: - logger.mgpu_mm_log("[IS_DISTORCH_MODEL] _mgpu_unload_distorch_model=False - Reporting 0 bytes (prevents eviction)") - return 0 - - # Not a DisTorch model - use original behavior - logger.mgpu_mm_log("[IS_DISTORCH_MODEL] Non-DisTorch model - Using original Comfy memory calculation") + # GET ORIGINAL MEMORY REQUIREMENT original_result = original_loaded_model_memory_required(self, device) original_gb = original_result / (1024**3) if original_result else 0 - logger.mgpu_mm_log(f"[IS_DISTORCH_MODEL] Original calculation returned: {original_gb:.2f}GB") - multigpu_memory_log("keep_loaded_memory_check", "end") - return original_result + + # CHECK FOR EJECT_MODELS PROPERTY + has_eject_models = hasattr(getattr(getattr(self, 'model', None), 'model', None), '_mgpu_eject_models') + logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] Original needs: {original_gb:.2f}GB, has_eject_models={has_eject_models}") + + # CHECK IF DISTORCH MODEL WITH VIRTUAL VRAM PROPERTY + is_distorch_model = hasattr(getattr(getattr(self, 'model', None), 'model', None), '_mgpu_virtual_vram_gb') + + # TRUTH TABLE APPLICATION + if has_eject_models: + if not is_distorch_model: + logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] ERROR: eject_models=1 but not DisTorch (mutually exclusive)") + # eject_models=1, is_distorch=1: RETURN MAX MEMORY TO FORCE EVICTION + logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] eject_models=1, is_distorch={is_distorch_model} → FORCING EVICTION WITH MAX MEMORY") + + # DISABLED: Manual ejection should happen automatically when MAX memory is returned + DISABLE_MANUAL_EJECTION = True # TODO: Remove this once auto-eviction confirmed + if not DISABLE_MANUAL_EJECTION: + logger.mgpu_mm_log(f"======= DIRECT MODEL EJECTION START[{model_name}] =======") + logger.mgpu_mm_log(f"[DIRECT_EJECTION][{model_name}] Current loaded models count: {len(mm.current_loaded_models)}") + + # DIRECTLY UNLOAD ALL MODELS + models_unloaded = [] + for i, lm in enumerate(mm.current_loaded_models): + model_name_to_eject = type(getattr(lm.model, 'model', lm.model)).__name__ if lm.model else 'Unknown' + logger.mgpu_mm_log(f"[DIRECT_EJECTION][{model_name}] UNLOADING MODEL {i+1}/{len(mm.current_loaded_models)}: {model_name_to_eject}") + try: + lm.model_unload(unpatch_weights=True) + models_unloaded.append(model_name_to_eject) + logger.mgpu_mm_log(f"[DIRECT_EJECTION][{model_name}] SUCCESSFULLY UNLOADED: {model_name_to_eject}") + except Exception as e: + logger.mgpu_mm_log(f"[DIRECT_EJECTION][{model_name}] ERROR unloading {model_name_to_eject}: {e}") + + mm.current_loaded_models = [] + logger.mgpu_mm_log(f"[DIRECT_EJECTION][{model_name}] Models unloaded: {models_unloaded}") + logger.mgpu_mm_log(f"======= DIRECT MODEL EJECTION COMPLETE[{model_name}] =======") + multigpu_memory_log("eject_models_post", "complete") + + # RETURN MAX MEMORY - Should trigger auto-eviction by Comfy Core + total_device_memory = mm.get_total_memory(device) + max_gb = total_device_memory / (1024**3) + logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] Returning MAX memory ({max_gb:.2f}GB) for auto-eviction by Comfy Core") + return total_device_memory + + elif is_distorch_model: + # eject_models=0, is_distorch=1: SUBTRACT VIRTUAL VRAM FROM ORIGINAL + virtual_vram_gb = getattr(getattr(self, 'model', None), 'model', None)._mgpu_virtual_vram_gb + virtual_vram_bytes = virtual_vram_gb * (1024**3) + adjusted_result = max(0, original_result - virtual_vram_bytes) + adjusted_gb = adjusted_result / (1024**3) if adjusted_result else 0 + + logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] eject_models=0, is_distorch=1 → adjusted {original_gb:.2f}GB - {virtual_vram_gb:.2f}GB = {adjusted_gb:.2f}GB (DisTorch allocation)") + multigpu_memory_log("distorch_allocation", "reported") + return adjusted_result + + else: + # eject_models=0, is_distorch=0: RETURN ORIGINAL + logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] eject_models=0, is_distorch=0 → returning original {original_gb:.2f}GB") + multigpu_memory_log("keep_loaded_memory_check", "end") + return original_result mm.LoadedModel.model_memory_required = patched_loaded_model_memory_required diff --git a/model_management_mgpu.py b/model_management_mgpu.py index 9a890ef..4b355fb 100644 --- a/model_management_mgpu.py +++ b/model_management_mgpu.py @@ -11,35 +11,12 @@ import comfy.model_management as mm import gc from datetime import datetime, timezone import server -import weakref -import platform -import ctypes -import comfy.model_patcher from collections import defaultdict logger = logging.getLogger("MultiGPU") -# ========================================================================================== -# GC Anchor System for Model Retention -# ========================================================================================== - -# Global anchor set to prevent GC of models during selective unload -_MGPU_RETENTION_ANCHORS = set() - -def add_retention_anchor(model_patcher, reason="keep_loaded"): - """Add a model patcher to the GC anchor set to prevent premature garbage collection""" - if model_patcher is not None: - _MGPU_RETENTION_ANCHORS.add(model_patcher) - model_name = type(getattr(model_patcher, 'model', model_patcher)).__name__ - logger.mgpu_mm_log(f"[GC_ANCHOR] Added retention anchor for {model_name}, reason: {reason}, total anchors: {len(_MGPU_RETENTION_ANCHORS)}") - -def clear_all_retention_anchors(reason="manual_clear"): - """Clear all retention anchors""" - count = len(_MGPU_RETENTION_ANCHORS) - _MGPU_RETENTION_ANCHORS.clear() - logger.mgpu_mm_log(f"[GC_ANCHOR] Cleared all {count} retention anchors, reason: {reason}") # ========================================================================================== # Model Analysis and Store Management (DisTorch V1 & V2) @@ -232,133 +209,3 @@ def force_full_system_cleanup(reason="manual", force=True): summary = f"[ManagerMatch] Cleanup requested (reason={reason}) | models {pre_models}->{post_models}, cpu_delta_mb={delta_cpu_mb:.2f}" logger.mgpu_mm_log(summary) return summary - -# ========================================================================================== -# Core Patching: unload_all_models -# ========================================================================================== - -if not hasattr(mm.unload_all_models, '_mgpu_eject_distorch_patched'): - logger.info("[MultiGPU Core Patching] Patching mm.unload_all_models for DisTorch2 ejection support") - - _mgpu_original_unload_all_models = mm.unload_all_models - - def _mgpu_patched_unload_all_models(): - """Patched mm.unload_all_models with selective ejection support and comprehensive diagnostics.""" - - logger.mgpu_mm_log(f"[UNLOAD_START] Patched unload_all_models called - initial model count: {len(mm.current_loaded_models)}") - - # Check if there are any DisTorch models that want to be unloaded - has_distorch_to_unload = any( - (hasattr(lm.model, '_mgpu_unload_distorch_model') and lm.model._mgpu_unload_distorch_model) or - (hasattr(getattr(lm.model, 'model', None), '_mgpu_unload_distorch_model') and lm.model.model._mgpu_unload_distorch_model) - for lm in mm.current_loaded_models - if lm.model is not None - ) - - if not has_distorch_to_unload: - logger.mgpu_mm_log("No DisTorch models requesting unload - clearing anchors and delegating to original unload_all_models") - clear_all_retention_anchors(reason="no_selective_unload_needed") - _mgpu_original_unload_all_models() - return - - # Direct approach: iterate through loaded models and selectively unload - models_to_unload = [] - kept_models = [] - - for i, lm in enumerate(mm.current_loaded_models): - mp = lm.model # weakref call to ModelPatcher - - # DIAGNOSTIC: Log full object chain - lm_id = id(lm) - mp_id = id(mp) - inner_model = getattr(mp, 'model', None) - inner_model_id = id(inner_model) if inner_model else None - inner_model_name = type(inner_model).__name__ if inner_model else "None" - - # Format inner_model_id properly for f-string - inner_id_str = f"0x{inner_model_id:x}" if inner_model_id is not None else "None" - - logger.mgpu_mm_log(f"[OBJECT_CHAIN_READ] Model {i}: lm_id=0x{lm_id:x}, mp_id=0x{mp_id:x}, inner_model_id={inner_id_str}, inner_model_type={inner_model_name}") - - # FIX: Check flag on ModelPatcher (where it was set), not on inner model - # OLD BUG: unload_distorch_model = getattr(mp.model, '_mgpu_unload_distorch_model', False) - # NEW FIX: Check both locations to see which one has the flag - flag_on_mp = getattr(mp, '_mgpu_unload_distorch_model', None) - flag_on_inner = getattr(mp.model, '_mgpu_unload_distorch_model', None) if inner_model else None - - logger.mgpu_mm_log(f"[FLAG_CHECK] Model {i} ({inner_model_name}): flag_on_mp={flag_on_mp}, flag_on_inner={flag_on_inner}") - - # Use whichever location has the flag (for backwards compatibility during transition) - if flag_on_mp is not None: - unload_distorch_model = flag_on_mp - logger.mgpu_mm_log(f"[FLAG_SOURCE] Using flag from ModelPatcher (mp_id=0x{mp_id:x})") - elif flag_on_inner is not None: - unload_distorch_model = flag_on_inner - logger.mgpu_mm_log(f"[FLAG_SOURCE] Using flag from inner model (inner_model_id={inner_id_str})") - else: - unload_distorch_model = False - logger.mgpu_mm_log(f"[FLAG_SOURCE] No flag found - defaulting to False (keep loaded)") - - logger.mgpu_mm_log(f"[DECISION] Model {i} ({inner_model_name}): unload_distorch_model={unload_distorch_model}") - - if unload_distorch_model: - models_to_unload.append(lm) - logger.mgpu_mm_log(f"[CATEGORIZE] Model {i} ({inner_model_name}) → models_to_unload") - else: - kept_models.append(lm) - add_retention_anchor(mp, "keep_loaded_protection") - logger.mgpu_mm_log(f"[CATEGORIZE] Model {i} ({inner_model_name}) → kept_models") - - # After the kept_models/models_to_unload evaluation - logger.mgpu_mm_log(f"[CATEGORIZE_SUMMARY] kept_models: {len(kept_models)}, models_to_unload: {len(models_to_unload)}, total: {len(mm.current_loaded_models)}") - - if len(kept_models) == len(mm.current_loaded_models): - # All models are meant to be kept - no DisTorch selective unloading needed - logger.mgpu_mm_log("[DELEGATION] All models flagged to be kept - delegating to standard unload_all_models") - _mgpu_original_unload_all_models() - return - - if kept_models: - logger.mgpu_mm_log(f"[SELECTIVE_UNLOAD] Proceeding with selective unload: retaining {len(kept_models)}, unloading {len(models_to_unload)}") - - # Unload models flagged for unload - for lm in models_to_unload: - try: - model_name = type(lm.model.model).__name__ if lm.model and hasattr(lm.model, 'model') else 'Unknown' - logger.mgpu_mm_log(f"[UNLOAD_EXECUTE] Unloading model: {model_name} (lm_id=0x{id(lm):x})") - lm.model_unload(unpatch_weights=True) - except Exception as e: - logger.warning(f"[UNLOAD_ERROR] Error unloading model: {e}") - - # WEAKREF TRACKING: Attach weakref callbacks to prove if kept models are GC'd - def model_deleted_callback(ref, model_name, model_id): - logger.mgpu_mm_log(f"[WEAKREF_DELETED] Kept model GARBAGE COLLECTED: {model_name} (id=0x{model_id:x})") - - for i, lm in enumerate(kept_models): - mp = lm.model - inner_model = getattr(mp, 'model', None) - model_name = type(inner_model).__name__ if inner_model else 'Unknown' - model_id = id(lm) - weakref.ref(lm, lambda ref, name=model_name, mid=model_id: model_deleted_callback(ref, name, mid)) - logger.mgpu_mm_log(f"[WEAKREF_ATTACHED] Tracking kept model {i}: {model_name} (lm_id=0x{model_id:x}, mp_id=0x{id(mp):x})") - - # Remove unloaded models from current_loaded_models - mm.current_loaded_models = kept_models - logger.mgpu_mm_log(f"[SELECTIVE_COMPLETE] Updated mm.current_loaded_models, new count: {len(mm.current_loaded_models)}") - logger.mgpu_mm_log(f"[SELECTIVE_COMPLETE] mm.current_loaded_models id: 0x{id(mm.current_loaded_models):x}") - - # DIAGNOSTIC: Log what's remaining - for i, lm in enumerate(mm.current_loaded_models): - mp = lm.model - inner_model = getattr(mp, 'model', None) - model_name = type(inner_model).__name__ if inner_model else "None" - logger.mgpu_mm_log(f"[REMAINING_MODEL] {i}: {model_name} (lm_id=0x{id(lm):x}, mp_id=0x{id(mp):x})") - else: - logger.mgpu_mm_log("[DELEGATION] No models with keep_loaded=True found - delegating to original unload_all_models") - _mgpu_original_unload_all_models() - - mm.unload_all_models = _mgpu_patched_unload_all_models - mm.unload_all_models._mgpu_eject_distorch_patched = True - logger.info("[MultiGPU Core Patching] mm.unload_all_models patched successfully") -else: - logger.debug("[MultiGPU Core Patching] mm.unload_all_models already patched - skipping") diff --git a/wrappers.py b/wrappers.py index 483adc9..d8d70ed 100644 --- a/wrappers.py +++ b/wrappers.py @@ -37,7 +37,7 @@ def _create_distorch_safetensor_v2_override(cls, device_param_name, device_sette inputs["optional"]["virtual_vram_gb"] = ("FLOAT", {"default": 4.0, "min": 0.0, "max": 128.0, "step": 0.1}) inputs["optional"]["donor_device"] = (devices, {"default": "cpu"}) inputs["optional"]["expert_mode_allocations"] = ("STRING", {"multiline": False, "default": ""}) - inputs["optional"]["keep_loaded"] = ("BOOLEAN", {"default": True}) + inputs["optional"]["eject_models"] = ("BOOLEAN", {"default": True}) return inputs CATEGORY = "multigpu/distorch_2" @@ -45,10 +45,10 @@ def _create_distorch_safetensor_v2_override(cls, device_param_name, device_sette TITLE = f"{cls.TITLE if hasattr(cls, 'TITLE') else cls.__name__} (DisTorch2)" @classmethod - def IS_CHANGED(s, *args, virtual_vram_gb=4.0, donor_device="cpu", - expert_mode_allocations="", keep_loaded=True, **kwargs): + def IS_CHANGED(s, *args, virtual_vram_gb=4.0, donor_device="cpu", + expert_mode_allocations="", eject_models=True, **kwargs): device_value = kwargs.get(device_param_name) - settings_str = f"{device_value}{virtual_vram_gb}{donor_device}{expert_mode_allocations}{keep_loaded}" + settings_str = f"{device_value}{virtual_vram_gb}{donor_device}{expert_mode_allocations}{eject_models}" current_hash = hashlib.sha256(settings_str.encode()).hexdigest() if not hasattr(cls, '_last_hash'): @@ -60,19 +60,40 @@ def _create_distorch_safetensor_v2_override(cls, device_param_name, device_sette return current_hash def override(self, *args, virtual_vram_gb=4.0, donor_device="cpu", - expert_mode_allocations="", keep_loaded=True, **kwargs): - + expert_mode_allocations="", eject_models=True, **kwargs): + device_value = kwargs.get(device_param_name) - unload_distorch_model = not keep_loaded + + import comfy.model_management as mm + + if eject_models: + logger.mgpu_mm_log(f"[EJECT_MODELS_SETUP] eject_models=True - marking all loaded models for eviction, device target: {device_value}") + ejection_count = 0 + for i, lm in enumerate(mm.current_loaded_models): + # Set _mgpu_unload_distorch_model=True on all models to force Comfy Core eviction + model_name = type(getattr(lm.model, 'model', lm.model)).__name__ if lm.model else 'Unknown' + + if hasattr(lm.model, 'model') and lm.model.model is not None: + lm.model.model._mgpu_unload_distorch_model = True + logger.mgpu_mm_log(f"[EJECT_MARKED] Model {i}: {model_name} (id=0x{id(lm):x}) → marked for eviction") + ejection_count += 1 + elif lm.model is not None: + lm.model._mgpu_unload_distorch_model = True + logger.mgpu_mm_log(f"[EJECT_MARKED] Model {i}: {model_name} (direct patcher) → marked for eviction") + ejection_count += 1 + + logger.mgpu_mm_log(f"[EJECT_MODELS_SETUP_COMPLETE] Marked {ejection_count} models for Comfy Core eviction during load_models_gpu") + else: + logger.mgpu_mm_log(f"[EJECT_MODELS_SETUP] eject_models=False - loading without eviction") if device_value is not None: device_setter_func(device_value) - # Strip MultiGPU-specific parameters before calling original function - clean_kwargs = {k: v for k, v in kwargs.items() - if k not in [device_param_name, 'virtual_vram_gb', - 'donor_device', 'expert_mode_allocations', - 'keep_loaded']} + # Strip MultiGPU-specific parameters before calling original function (REMOVE eject_models, keep_loaded and virtual_vram_gb since we handle them above) + clean_kwargs = {k: v for k, v in kwargs.items() + if k not in [device_param_name, 'virtual_vram_gb', + 'donor_device', 'expert_mode_allocations', + 'eject_models']} if apply_device_kwarg_workaround: clean_kwargs['device'] = 'default' @@ -106,7 +127,7 @@ def _create_distorch_safetensor_v2_override(cls, device_param_name, device_sette logger.debug(f"[MultiGPU DisTorch V2] Stored allocation for model {model_hash[:8]}: {full_allocation}") logger.info(f"[MultiGPU DisTorch V2] Full allocation string: {full_allocation}") - logger.mgpu_mm_log(f"[FLAG_SET_START] Setting '_mgpu_unload_distorch_model' to: {unload_distorch_model} (keep_loaded={keep_loaded})") + logger.mgpu_mm_log(f"[MODEL_SETUP] Setting DisTorch model properties: virtual_vram_gb={virtual_vram_gb}") if hasattr(out[0], 'model'): mp = out[0] @@ -115,16 +136,19 @@ def _create_distorch_safetensor_v2_override(cls, device_param_name, device_sette inner_model_id = id(inner_model) if inner_model else None inner_model_name = type(inner_model).__name__ if inner_model else "None" inner_id_str = f"0x{inner_model_id:x}" if inner_model_id is not None else "None" - + logger.mgpu_mm_log(f"[OBJECT_CHAIN_SET] ModelPatcher: mp_id=0x{mp_id:x}, inner_model_id={inner_id_str}, inner_model_type={inner_model_name}") - - mp._mgpu_unload_distorch_model = unload_distorch_model - logger.mgpu_mm_log(f"[FLAG_SET_LOCATION] Set on ModelPatcher (mp_id=0x{mp_id:x}): mp._mgpu_unload_distorch_model = {unload_distorch_model}") - + + # SET VIRTUAL VRAM PROPERTY FOR MEMORY CALCULATION if inner_model: - inner_model._mgpu_unload_distorch_model = unload_distorch_model - logger.mgpu_mm_log(f"[FLAG_SET_COMPAT] Also set on inner model (inner_model_id=0x{inner_model_id:x}) for compatibility") - + inner_model._mgpu_virtual_vram_gb = virtual_vram_gb + logger.mgpu_mm_log(f"[VIRTUAL_VRAM_SET] Set _mgpu_virtual_vram_gb={virtual_vram_gb}GB on inner model (id=0x{inner_model_id:x}) for memory assessment") + + # SET EJECT MODELS PROPERTY IF ENABLED + if eject_models and inner_model: + inner_model._mgpu_eject_models = True + logger.mgpu_mm_log(f"[EJECT_FLAG_SET] Set _mgpu_eject_models=True on inner model (id=0x{inner_model_id:x}) - will trigger ejection during load_models_gpu") + elif hasattr(out[0], 'patcher') and hasattr(out[0].patcher, 'model'): mp = out[0].patcher mp_id = id(mp) @@ -132,19 +156,13 @@ def _create_distorch_safetensor_v2_override(cls, device_param_name, device_sette inner_model_id = id(inner_model) if inner_model else None inner_model_name = type(inner_model).__name__ if inner_model else "None" inner_id_str = f"0x{inner_model_id:x}" if inner_model_id is not None else "None" - - logger.mgpu_mm_log(f"[OBJECT_CHAIN_SET] ModelPatcher via patcher: mp_id=0x{mp_id:x}, inner_model_id={inner_id_str}, inner_model_type={inner_model_name}") - - mp._mgpu_unload_distorch_model = unload_distorch_model - logger.mgpu_mm_log(f"[FLAG_SET_LOCATION] Set on ModelPatcher (mp_id=0x{mp_id:x}): mp._mgpu_unload_distorch_model = {unload_distorch_model}") - - if inner_model: - inner_model._mgpu_unload_distorch_model = unload_distorch_model - logger.mgpu_mm_log(f"[FLAG_SET_COMPAT] Also set on inner model (inner_model_id=0x{inner_model_id:x}) for compatibility") - if unload_distorch_model: - logger.mgpu_mm_log("[FLAG_TRIGGER] unload_distorch_model=True, triggering full system cleanup") - force_full_system_cleanup(reason="policy_every_load", force=True) + logger.mgpu_mm_log(f"[OBJECT_CHAIN_SET] ModelPatcher via patcher: mp_id=0x{mp_id:x}, inner_model_id={inner_id_str}, inner_model_type={inner_model_name}") + + # SET VIRTUAL VRAM PROPERTY FOR MEMORY CALCULATION + if inner_model: + inner_model._mgpu_virtual_vram_gb = virtual_vram_gb + logger.mgpu_mm_log(f"[VIRTUAL_VRAM_SET] Set _mgpu_virtual_vram_gb={virtual_vram_gb}GB on inner model (id=0x{inner_model_id:x}) for memory assessment") return out From 72a20338ef757bf7ee251072188bef2df47ea9d3 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sat, 4 Oct 2025 17:16:51 -0500 Subject: [PATCH 5/7] feat: patch load_models_gpu for accurate memory calculations; unpatch load_models_gpu Refactor memory management in distorch_2.py to patch load_models_gpu instead of LoadedModel.model_memory_required. Implement correct memory reporting based on model flags (eject_models and is_distorch_model), ensuring proper eviction logic and improved handling of virtual VRAM. This drives behavior purely by either comfy core matching or DisTorch flag, fixing potential issues in multi-GPU setups. --- distorch_2.py | 211 +++++++++++------- web/docs/CLIPLoaderDisTorch2MultiGPU.md | 2 +- web/docs/CLIPLoaderGGUFDisTorch2MultiGPU.md | 2 +- web/docs/CLIPVisionLoaderDisTorch2MultiGPU.md | 2 +- ...CheckpointLoaderSimpleDisTorch2MultiGPU.md | 2 +- web/docs/ControlNetLoaderDisTorch2MultiGPU.md | 2 +- .../DiffControlNetLoaderDisTorch2MultiGPU.md | 2 +- web/docs/DiffusersLoaderDisTorch2MultiGPU.md | 2 +- web/docs/DualCLIPLoaderDisTorch2MultiGPU.md | 2 +- .../DualCLIPLoaderGGUFDisTorch2MultiGPU.md | 2 +- .../QuadrupleCLIPLoaderDisTorch2MultiGPU.md | 2 +- ...uadrupleCLIPLoaderGGUFDisTorch2MultiGPU.md | 2 +- web/docs/TripleCLIPLoaderDisTorch2MultiGPU.md | 2 +- .../TripleCLIPLoaderGGUFDisTorch2MultiGPU.md | 2 +- web/docs/UNETLoaderDisTorch2MultiGPU.md | 2 +- ...UnetLoaderGGUFAdvancedDisTorch2MultiGPU.md | 2 +- web/docs/UnetLoaderGGUFDisTorch2MultiGPU.md | 2 +- web/docs/VAELoaderDisTorch2MultiGPU.md | 2 +- wrappers.py | 16 +- 19 files changed, 158 insertions(+), 103 deletions(-) diff --git a/distorch_2.py b/distorch_2.py index 733d677..9b1e159 100644 --- a/distorch_2.py +++ b/distorch_2.py @@ -58,96 +58,149 @@ def register_patched_safetensor_modelpatcher(): # Patch ComfyUI's ModelPatcher if not hasattr(comfy.model_patcher.ModelPatcher, '_distorch_patched'): - # Patch LoadedModel.model_memory_required to drive behavior purely by Phase 2 = unload_distorch_model flag - from comfy.model_management import current_loaded_models - original_loaded_model_memory_required = None - for cls in current_loaded_models.__class__.__mro__: - if hasattr(cls, 'model_memory_required'): - original_loaded_model_memory_required = cls.model_memory_required - break + # PATCH load_models_gpu with correct memory calculations per model flags + original_load_models_gpu = mm.load_models_gpu - if original_loaded_model_memory_required is None: - # Global patch of LoadedModel class if available - import comfy.model_management as mm + def patched_load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None, force_full_load=False): + from comfy.model_management import cleanup_models_gc, get_free_memory, free_memory, current_loaded_models + from comfy.model_management import VRAMState, vram_state, lowvram_available, MIN_WEIGHT_MEMORY_RATIO + from comfy.model_management import minimum_inference_memory, extra_reserved_memory, is_device_cpu + + multigpu_memory_log("load_models_gpu_top_level", "start") - original_loaded_model_memory_required = mm.LoadedModel.model_memory_required + cleanup_models_gc() - def patched_loaded_model_memory_required(self, device): - """Truth table for memory reporting: - eject_models=0, is_distorch=0: return original - eject_models=0, is_distorch=1: return original - virtual_vram_gb_bytes - eject_models=1, is_distorch=0: mutually exclusive (shouldn't occur) - eject_models=1, is_distorch=1: return MAX memory to force eviction""" - multigpu_memory_log("unload_distorch_model_memory_check", "start") - model_name = type(getattr(getattr(self, 'model', None), 'model', None)).__name__ if getattr(getattr(self, 'model', None), 'model', None) else "Unknown" - logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] Memory assessment requested for model on device: {device}") + inference_memory = minimum_inference_memory() + extra_reserved_mem = extra_reserved_memory() + memory_required_total = memory_required + extra_reserved_mem + extra_mem = max(inference_memory, memory_required_total) + if minimum_memory_required is None: + minimum_memory_required = extra_mem + else: + minimum_memory_required = max(inference_memory, minimum_memory_required + extra_reserved_mem) - # GET ORIGINAL MEMORY REQUIREMENT - original_result = original_loaded_model_memory_required(self, device) - original_gb = original_result / (1024**3) if original_result else 0 + models_temp = set() + for m in models: + models_temp.add(m) + for mm_patch in m.model_patches_models(): + models_temp.add(mm_patch) - # CHECK FOR EJECT_MODELS PROPERTY - has_eject_models = hasattr(getattr(getattr(self, 'model', None), 'model', None), '_mgpu_eject_models') - logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] Original needs: {original_gb:.2f}GB, has_eject_models={has_eject_models}") + models = models_temp - # CHECK IF DISTORCH MODEL WITH VIRTUAL VRAM PROPERTY - is_distorch_model = hasattr(getattr(getattr(self, 'model', None), 'model', None), '_mgpu_virtual_vram_gb') + models_to_load = [] - # TRUTH TABLE APPLICATION - if has_eject_models: - if not is_distorch_model: - logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] ERROR: eject_models=1 but not DisTorch (mutually exclusive)") - # eject_models=1, is_distorch=1: RETURN MAX MEMORY TO FORCE EVICTION - logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] eject_models=1, is_distorch={is_distorch_model} → FORCING EVICTION WITH MAX MEMORY") - - # DISABLED: Manual ejection should happen automatically when MAX memory is returned - DISABLE_MANUAL_EJECTION = True # TODO: Remove this once auto-eviction confirmed - if not DISABLE_MANUAL_EJECTION: - logger.mgpu_mm_log(f"======= DIRECT MODEL EJECTION START[{model_name}] =======") - logger.mgpu_mm_log(f"[DIRECT_EJECTION][{model_name}] Current loaded models count: {len(mm.current_loaded_models)}") - - # DIRECTLY UNLOAD ALL MODELS - models_unloaded = [] - for i, lm in enumerate(mm.current_loaded_models): - model_name_to_eject = type(getattr(lm.model, 'model', lm.model)).__name__ if lm.model else 'Unknown' - logger.mgpu_mm_log(f"[DIRECT_EJECTION][{model_name}] UNLOADING MODEL {i+1}/{len(mm.current_loaded_models)}: {model_name_to_eject}") - try: - lm.model_unload(unpatch_weights=True) - models_unloaded.append(model_name_to_eject) - logger.mgpu_mm_log(f"[DIRECT_EJECTION][{model_name}] SUCCESSFULLY UNLOADED: {model_name_to_eject}") - except Exception as e: - logger.mgpu_mm_log(f"[DIRECT_EJECTION][{model_name}] ERROR unloading {model_name_to_eject}: {e}") - - mm.current_loaded_models = [] - logger.mgpu_mm_log(f"[DIRECT_EJECTION][{model_name}] Models unloaded: {models_unloaded}") - logger.mgpu_mm_log(f"======= DIRECT MODEL EJECTION COMPLETE[{model_name}] =======") - multigpu_memory_log("eject_models_post", "complete") - - # RETURN MAX MEMORY - Should trigger auto-eviction by Comfy Core - total_device_memory = mm.get_total_memory(device) - max_gb = total_device_memory / (1024**3) - logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] Returning MAX memory ({max_gb:.2f}GB) for auto-eviction by Comfy Core") - return total_device_memory - - elif is_distorch_model: - # eject_models=0, is_distorch=1: SUBTRACT VIRTUAL VRAM FROM ORIGINAL - virtual_vram_gb = getattr(getattr(self, 'model', None), 'model', None)._mgpu_virtual_vram_gb - virtual_vram_bytes = virtual_vram_gb * (1024**3) - adjusted_result = max(0, original_result - virtual_vram_bytes) - adjusted_gb = adjusted_result / (1024**3) if adjusted_result else 0 - - logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] eject_models=0, is_distorch=1 → adjusted {original_gb:.2f}GB - {virtual_vram_gb:.2f}GB = {adjusted_gb:.2f}GB (DisTorch allocation)") - multigpu_memory_log("distorch_allocation", "reported") - return adjusted_result + for x in models: + loaded_model = mm.LoadedModel(x) + try: + loaded_model_index = current_loaded_models.index(loaded_model) + except: + loaded_model_index = None + if loaded_model_index is not None: + loaded = current_loaded_models[loaded_model_index] + loaded.currently_used = True + models_to_load.append(loaded) else: - # eject_models=0, is_distorch=0: RETURN ORIGINAL - logger.mgpu_mm_log(f"[MEM_REPORT][{model_name}] eject_models=0, is_distorch=0 → returning original {original_gb:.2f}GB") - multigpu_memory_log("keep_loaded_memory_check", "end") - return original_result + if hasattr(x, "model"): + logging.info(f"Requested to load {x.model.__class__.__name__}") + models_to_load.append(loaded_model) - mm.LoadedModel.model_memory_required = patched_loaded_model_memory_required + for loaded_model in models_to_load: + to_unload = [] + for i in range(len(current_loaded_models)): + if loaded_model.model.is_clone(current_loaded_models[i].model): + to_unload = [i] + to_unload + for i in to_unload: + model_to_unload = current_loaded_models.pop(i) + model_to_unload.model.detach(unpatch_all=False) + model_to_unload.model_finalizer.detach() + + # DisTorch Processing + total_memory_required = {} + eject_device = None + + for loaded_model in models_to_load: + device = loaded_model.device + base_memory = loaded_model.model_memory_required(device) + + # Check DisTorch flags + is_distorch = hasattr(loaded_model.model.model, '_mgpu_virtual_vram_gb') + has_eject = hasattr(loaded_model.model.model, '_mgpu_eject_models') + + if has_eject: + eject_device = device + logger.mgpu_mm_log("DisTorch eject_models=True, is_distorch=True - MAX memory eviction") + + if is_distorch: + # is_distorch=True: use compute device allocation size + virtual_vram_gb = loaded_model.model.model._mgpu_virtual_vram_gb + virtual_vram_bytes = virtual_vram_gb * (1024**3) + adjusted_memory = max(0, base_memory - virtual_vram_bytes) + total_memory_required[device] = total_memory_required.get(device, 0) + adjusted_memory + logger.mgpu_mm_log(f"DisTorch is_distorch=True, model adjusted {(base_memory - virtual_vram_bytes)/(1024**3):.2f}GB for device {device}") + else: + # is_distorch=False: use full model size + total_memory_required[device] = total_memory_required.get(device, 0) + base_memory + logger.mgpu_mm_log(f"[LOAD_MODELS_GPU] Standard model {(base_memory)/(1024**3):.2f}GB for device {device}") + + for device in total_memory_required: + if device != torch.device("cpu"): + requested_mem = total_memory_required[device] * 1.1 + extra_mem + logger.mgpu_mm_log(f"[FREE_MEMORY_CALL] Device {device}: requesting {requested_mem/(1024**3):.2f}GB = {total_memory_required[device]/(1024**3):.2f}GB * 1.1 + {extra_mem/(1024**3):.2f}GB inference") + + + multigpu_memory_log("free_memory", "pre") + + for device in total_memory_required: + if device != torch.device("cpu"): + if device == eject_device: + total_device_memory = mm.get_total_memory(device) + logger.mgpu_mm_log(f"[LOAD_MODELS_GPU] eject_models=1, is_distorch=1 → using MAX memory ({total_device_memory/(1024**3):.2f}GB) for eviction") + free_memory(total_device_memory,device) + else: + logger.mgpu_mm_log(f"[LOAD_MODELS_GPU] eject_models=0, using Comfy Core Computed memory ({(total_memory_required[device] * 1.1 + extra_mem)/(1024**3):.2f}GB) for eviction") + free_memory(total_memory_required[device] * 1.1 + extra_mem, device) + + multigpu_memory_log("free_memory/minimum_memory_required", "post/pre") + + for device in total_memory_required: + if device != torch.device("cpu"): + free_mem = get_free_memory(device) + free_mem_gb = free_mem / (1024**3) + min_required_gb = minimum_memory_required / (1024**3) + logger.mgpu_mm_log(f"[MIN_MEMORY_CHECK] Device {device}: free={free_mem_gb:.2f}GB, required={min_required_gb:.2f}GB, will_evict={free_mem < minimum_memory_required}") + + if free_mem < minimum_memory_required: + models_l = free_memory(minimum_memory_required, device) + logger.mgpu_mm_log(f"[EVICTION] Device {device}: unloaded {len(models_l)} models due to insufficient memory") + logging.info("{} models unloaded.".format(len(models_l))) + + multigpu_memory_log("minimum_memory_required", "post") + + for loaded_model in models_to_load: + model = loaded_model.model + torch_dev = model.load_device + if is_device_cpu(torch_dev): + vram_set_state = VRAMState.DISABLED + else: + vram_set_state = vram_state + lowvram_model_memory = 0 + if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM) and not force_full_load: + loaded_memory = loaded_model.model_loaded_memory() + current_free_mem = get_free_memory(torch_dev) + loaded_memory + + lowvram_model_memory = max(128 * 1024 * 1024, (current_free_mem - minimum_memory_required), min(current_free_mem * MIN_WEIGHT_MEMORY_RATIO, current_free_mem - minimum_inference_memory())) + lowvram_model_memory = max(0.1, lowvram_model_memory - loaded_memory) + + if vram_set_state == VRAMState.NO_VRAM: + lowvram_model_memory = 0.1 + + loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights) + current_loaded_models.insert(0, loaded_model) + + # Replace the module function + mm.load_models_gpu = patched_load_models_gpu original_partially_load = comfy.model_patcher.ModelPatcher.partially_load diff --git a/web/docs/CLIPLoaderDisTorch2MultiGPU.md b/web/docs/CLIPLoaderDisTorch2MultiGPU.md index 85ba923..ca0c197 100644 --- a/web/docs/CLIPLoaderDisTorch2MultiGPU.md +++ b/web/docs/CLIPLoaderDisTorch2MultiGPU.md @@ -14,7 +14,7 @@ This node automatically detects models located in the `ComfyUI/models/clip` fold | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: false for CLIP loaders). | ## Outputs diff --git a/web/docs/CLIPLoaderGGUFDisTorch2MultiGPU.md b/web/docs/CLIPLoaderGGUFDisTorch2MultiGPU.md index 3fe43e8..fe383c4 100644 --- a/web/docs/CLIPLoaderGGUFDisTorch2MultiGPU.md +++ b/web/docs/CLIPLoaderGGUFDisTorch2MultiGPU.md @@ -14,7 +14,7 @@ This node automatically detects models located in the `ComfyUI/models/clip` and | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: false for CLIP loaders). | ## Outputs diff --git a/web/docs/CLIPVisionLoaderDisTorch2MultiGPU.md b/web/docs/CLIPVisionLoaderDisTorch2MultiGPU.md index af24618..f01fbbe 100644 --- a/web/docs/CLIPVisionLoaderDisTorch2MultiGPU.md +++ b/web/docs/CLIPVisionLoaderDisTorch2MultiGPU.md @@ -13,7 +13,7 @@ This node automatically detects models located in the `ComfyUI/models/clip_visio | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: true). | ## Outputs diff --git a/web/docs/CheckpointLoaderSimpleDisTorch2MultiGPU.md b/web/docs/CheckpointLoaderSimpleDisTorch2MultiGPU.md index dcbd5fa..2e65802 100644 --- a/web/docs/CheckpointLoaderSimpleDisTorch2MultiGPU.md +++ b/web/docs/CheckpointLoaderSimpleDisTorch2MultiGPU.md @@ -13,7 +13,7 @@ This node automatically detects models located in the `ComfyUI/models/checkpoint | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: true). | ## Outputs diff --git a/web/docs/ControlNetLoaderDisTorch2MultiGPU.md b/web/docs/ControlNetLoaderDisTorch2MultiGPU.md index c5a70fa..857d4e5 100644 --- a/web/docs/ControlNetLoaderDisTorch2MultiGPU.md +++ b/web/docs/ControlNetLoaderDisTorch2MultiGPU.md @@ -13,7 +13,7 @@ This node automatically detects models located in the `ComfyUI/models/controlnet | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: true). | ## Outputs diff --git a/web/docs/DiffControlNetLoaderDisTorch2MultiGPU.md b/web/docs/DiffControlNetLoaderDisTorch2MultiGPU.md index 9f1f091..d41e37d 100644 --- a/web/docs/DiffControlNetLoaderDisTorch2MultiGPU.md +++ b/web/docs/DiffControlNetLoaderDisTorch2MultiGPU.md @@ -13,7 +13,7 @@ This node loads ControlNet models directly from HuggingFace model repositories b | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: true). | ## Outputs diff --git a/web/docs/DiffusersLoaderDisTorch2MultiGPU.md b/web/docs/DiffusersLoaderDisTorch2MultiGPU.md index 0654562..f666024 100644 --- a/web/docs/DiffusersLoaderDisTorch2MultiGPU.md +++ b/web/docs/DiffusersLoaderDisTorch2MultiGPU.md @@ -13,7 +13,7 @@ This node loads models directly from HuggingFace model repositories by specifyin | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: true). | ## Outputs diff --git a/web/docs/DualCLIPLoaderDisTorch2MultiGPU.md b/web/docs/DualCLIPLoaderDisTorch2MultiGPU.md index 38ff506..3e5053c 100644 --- a/web/docs/DualCLIPLoaderDisTorch2MultiGPU.md +++ b/web/docs/DualCLIPLoaderDisTorch2MultiGPU.md @@ -15,7 +15,7 @@ This node automatically detects models located in the `ComfyUI/models/clip` fold | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: false for CLIP loaders). | ## Outputs diff --git a/web/docs/DualCLIPLoaderGGUFDisTorch2MultiGPU.md b/web/docs/DualCLIPLoaderGGUFDisTorch2MultiGPU.md index eb6387e..657af23 100644 --- a/web/docs/DualCLIPLoaderGGUFDisTorch2MultiGPU.md +++ b/web/docs/DualCLIPLoaderGGUFDisTorch2MultiGPU.md @@ -15,7 +15,7 @@ This node automatically detects models located in the `ComfyUI/models/clip` and | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: false for CLIP loaders). | ## Outputs diff --git a/web/docs/QuadrupleCLIPLoaderDisTorch2MultiGPU.md b/web/docs/QuadrupleCLIPLoaderDisTorch2MultiGPU.md index dd4ba62..623cc34 100644 --- a/web/docs/QuadrupleCLIPLoaderDisTorch2MultiGPU.md +++ b/web/docs/QuadrupleCLIPLoaderDisTorch2MultiGPU.md @@ -16,7 +16,7 @@ This node automatically detects models located in the `ComfyUI/models/clip` fold | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: false for CLIP loaders). | ## Outputs diff --git a/web/docs/QuadrupleCLIPLoaderGGUFDisTorch2MultiGPU.md b/web/docs/QuadrupleCLIPLoaderGGUFDisTorch2MultiGPU.md index 3623767..95af4e8 100644 --- a/web/docs/QuadrupleCLIPLoaderGGUFDisTorch2MultiGPU.md +++ b/web/docs/QuadrupleCLIPLoaderGGUFDisTorch2MultiGPU.md @@ -16,7 +16,7 @@ This node automatically detects models located in the `ComfyUI/models/clip` and | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: false for CLIP loaders). | ## Outputs diff --git a/web/docs/TripleCLIPLoaderDisTorch2MultiGPU.md b/web/docs/TripleCLIPLoaderDisTorch2MultiGPU.md index b714fb6..88aac1d 100644 --- a/web/docs/TripleCLIPLoaderDisTorch2MultiGPU.md +++ b/web/docs/TripleCLIPLoaderDisTorch2MultiGPU.md @@ -15,7 +15,7 @@ This node automatically detects models located in the `ComfyUI/models/clip` fold | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: false for CLIP loaders). | ## Outputs diff --git a/web/docs/TripleCLIPLoaderGGUFDisTorch2MultiGPU.md b/web/docs/TripleCLIPLoaderGGUFDisTorch2MultiGPU.md index e80a2fa..9c11735 100644 --- a/web/docs/TripleCLIPLoaderGGUFDisTorch2MultiGPU.md +++ b/web/docs/TripleCLIPLoaderGGUFDisTorch2MultiGPU.md @@ -15,7 +15,7 @@ This node automatically detects models located in the `ComfyUI/models/clip` and | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: false for CLIP loaders). | ## Outputs diff --git a/web/docs/UNETLoaderDisTorch2MultiGPU.md b/web/docs/UNETLoaderDisTorch2MultiGPU.md index bf79e54..dd3246f 100644 --- a/web/docs/UNETLoaderDisTorch2MultiGPU.md +++ b/web/docs/UNETLoaderDisTorch2MultiGPU.md @@ -13,7 +13,7 @@ This node automatically detects models located in the `ComfyUI/models/unet` fold | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: true). | ## Outputs diff --git a/web/docs/UnetLoaderGGUFAdvancedDisTorch2MultiGPU.md b/web/docs/UnetLoaderGGUFAdvancedDisTorch2MultiGPU.md index 601f585..1c1785d 100644 --- a/web/docs/UnetLoaderGGUFAdvancedDisTorch2MultiGPU.md +++ b/web/docs/UnetLoaderGGUFAdvancedDisTorch2MultiGPU.md @@ -16,7 +16,7 @@ This node automatically detects models located in the `ComfyUI/models/unet_gguf` | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: true). | ## Outputs diff --git a/web/docs/UnetLoaderGGUFDisTorch2MultiGPU.md b/web/docs/UnetLoaderGGUFDisTorch2MultiGPU.md index 8d37e3f..de1f3e0 100644 --- a/web/docs/UnetLoaderGGUFDisTorch2MultiGPU.md +++ b/web/docs/UnetLoaderGGUFDisTorch2MultiGPU.md @@ -13,7 +13,7 @@ This node automatically detects models located in the `ComfyUI/models/unet_gguf` | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: true). | ## Outputs diff --git a/web/docs/VAELoaderDisTorch2MultiGPU.md b/web/docs/VAELoaderDisTorch2MultiGPU.md index fb35916..9df192d 100644 --- a/web/docs/VAELoaderDisTorch2MultiGPU.md +++ b/web/docs/VAELoaderDisTorch2MultiGPU.md @@ -13,7 +13,7 @@ This node automatically detects models located in the `ComfyUI/models/vae` folde | `virtual_vram_gb` | `FLOAT` | Amount of virtual VRAM in gigabytes to allocate for distributed tensor management (default: 4.0, range: 0.0-128.0). | | `donor_device` | `STRING` | Device to donate VRAM from when allocating virtual memory (default: 'cpu'). | | `expert_mode_allocations` | `STRING` | Advanced allocation string for expert users to manually specify device/ratio distributions (e.g., 'cuda:0,50%;cpu,*'). | -| `keep_loaded` | `BOOLEAN` | Whether to keep the model loaded when triggering memory cleanup operations (default: true). | +| `eject_models` | `BOOLEAN` | Whether to unload ALL models from the target device before loading this model, enabling deterministic model eviction for testing and memory management (default: true). | ## Outputs diff --git a/wrappers.py b/wrappers.py index d8d70ed..9abf9bf 100644 --- a/wrappers.py +++ b/wrappers.py @@ -15,7 +15,7 @@ logger = logging.getLogger("MultiGPU") # DISTORCH V2 SAFETENSOR WRAPPERS (DisTorch2 for .safetensors and .gguf) # ============================================================================ -def _create_distorch_safetensor_v2_override(cls, device_param_name, device_setter_func, apply_device_kwarg_workaround): +def _create_distorch_safetensor_v2_override(cls, device_param_name, device_setter_func, apply_device_kwarg_workaround, eject_models_default=True): """Internal factory function creating DisTorch2 override class with parameterized device selection behavior.""" from .distorch_2 import ( register_patched_safetensor_modelpatcher, @@ -37,7 +37,7 @@ def _create_distorch_safetensor_v2_override(cls, device_param_name, device_sette inputs["optional"]["virtual_vram_gb"] = ("FLOAT", {"default": 4.0, "min": 0.0, "max": 128.0, "step": 0.1}) inputs["optional"]["donor_device"] = (devices, {"default": "cpu"}) inputs["optional"]["expert_mode_allocations"] = ("STRING", {"multiline": False, "default": ""}) - inputs["optional"]["eject_models"] = ("BOOLEAN", {"default": True}) + inputs["optional"]["eject_models"] = ("BOOLEAN", {"default": eject_models_default}) return inputs CATEGORY = "multigpu/distorch_2" @@ -46,7 +46,7 @@ def _create_distorch_safetensor_v2_override(cls, device_param_name, device_sette @classmethod def IS_CHANGED(s, *args, virtual_vram_gb=4.0, donor_device="cpu", - expert_mode_allocations="", eject_models=True, **kwargs): + expert_mode_allocations="", eject_models=eject_models_default, **kwargs): device_value = kwargs.get(device_param_name) settings_str = f"{device_value}{virtual_vram_gb}{donor_device}{expert_mode_allocations}{eject_models}" current_hash = hashlib.sha256(settings_str.encode()).hexdigest() @@ -60,7 +60,7 @@ def _create_distorch_safetensor_v2_override(cls, device_param_name, device_sette return current_hash def override(self, *args, virtual_vram_gb=4.0, donor_device="cpu", - expert_mode_allocations="", eject_models=True, **kwargs): + expert_mode_allocations="", eject_models=eject_models_default, **kwargs): device_value = kwargs.get(device_param_name) @@ -89,7 +89,7 @@ def _create_distorch_safetensor_v2_override(cls, device_param_name, device_sette if device_value is not None: device_setter_func(device_value) - # Strip MultiGPU-specific parameters before calling original function (REMOVE eject_models, keep_loaded and virtual_vram_gb since we handle them above) + # Strip MultiGPU-specific parameters before calling original function (REMOVE eject_models, eject_models and virtual_vram_gb since we handle them above) clean_kwargs = {k: v for k, v in kwargs.items() if k not in [device_param_name, 'virtual_vram_gb', 'donor_device', 'expert_mode_allocations', @@ -187,7 +187,8 @@ def override_class_with_distorch_safetensor_v2_clip(cls): cls, device_param_name="device", device_setter_func=set_current_text_encoder_device, - apply_device_kwarg_workaround=True + apply_device_kwarg_workaround=True, + eject_models_default=False # CLIP defaults to False ) @@ -198,7 +199,8 @@ def override_class_with_distorch_safetensor_v2_clip_no_device(cls): cls, device_param_name="device", device_setter_func=set_current_text_encoder_device, - apply_device_kwarg_workaround=False + apply_device_kwarg_workaround=False, + eject_models_default=False # CLIP defaults to False ) From f60aa6a9a72e9de57d3b4182c6f6d78d1d6c2359 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sat, 4 Oct 2025 17:42:52 -0500 Subject: [PATCH 6/7] refactor: keep_loaded --> eject_models Boolean switch. Use it to eject all other models prior to loading model for inference; helpful to maximize available latent space on device prior to UNet inference, for example So this is a change from something just newly-released in 2.5.0, but most should either see an improvement or no change to behavior. This was the weakest, and jankiest part of 2.5.0 and my decision to manage a CPU memory leak turned into a too-aggressive solution with unwanted side effects. This solution should provide a better way to manage `compute` VRAM as the most asked-for feature is a way to remove everything else from VRAM prior to main UNet inference, which this accomplishes nicely, as well as reporting back accurate information DisTorch2 on-device shard sizes. --- __init__.py | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/__init__.py b/__init__.py index 40b71c1..41b4f15 100644 --- a/__init__.py +++ b/__init__.py @@ -21,7 +21,7 @@ from .model_management_mgpu import ( ) WEB_DIRECTORY = "./web" -MGPU_MM_LOG = True +MGPU_MM_LOG = False DEBUG_LOG = False logger = logging.getLogger("MultiGPU") diff --git a/pyproject.toml b/pyproject.toml index 39f36ac..eca2803 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-multigpu" description = "Provides a suite of custom nodes to manage multiple GPUs for ComfyUI, including advanced model offloading for both GGUF and Safetensor formats with DisTorch, and bespoke MultiGPU support for WanVideoWrapper and other custom nodes." -version = "2.5.0" +version = "2.5.1" license = {file = "LICENSE"} [project.urls] From 0919fd4ccbcbdf46b21705a6f61f318efbc98fd8 Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sat, 4 Oct 2025 18:08:22 -0500 Subject: [PATCH 7/7] Minor cleanup related to 2.5.1 release --- distorch_2.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/distorch_2.py b/distorch_2.py index 9b1e159..b92bad2 100644 --- a/distorch_2.py +++ b/distorch_2.py @@ -254,7 +254,7 @@ def register_patched_safetensor_modelpatcher(): pass if current_module_device is not None and str(current_module_device) != str(block_target_device): - logger.info(f"[MultiGPU DisTorch V2] Moving already patched {module_name} to {block_target_device}") + logger.debug(f"[MultiGPU DisTorch V2] Moving already patched {module_name} to {block_target_device}") module_object.to(block_target_device) mem_counter += module_size