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.
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.
Add comprehensive guide for accessing and using documentation on core MultiGPU and DisTorch2 nodes, covering 36+ nodes with detailed parameters, outputs, and usage examples. This enhances user experience by providing easy reference for standard ComfyUI loaders and DisTorch2 features, while clarifying coverage excludes third-party nodes.
Add a new constant WEB_DIRECTORY set to "./web" to define the directory path for web-related assets during package initialization. This improves organization by centralizing the path configuration. Additionally, removed trailing newline at file end to maintain consistent code formatting.
Update active context documentation to reflect v2.5.0 release candidate status.
Major achievements documented:
- DisTorch2 allocation refactoring (-179 lines, unified UNET/CLIP logic)
- Production cleanup removing debug instrumentation (-40 lines)
- Verified selective unload system working with production logs
- Architecture status showing all core files production-ready
- Updated memory management pipeline with verification details
Reorganized content to prioritize recent session achievements (2025-09-30)
and production readiness status. Total code reduction: 219 lines through
consolidation and cleanup while maintaining full functionality.
This commit represents a significant architectural refactoring to improve code organization,
eliminate redundancy, and fix a critical bug in wrapper functions. Net reduction of 531 lines
while improving maintainability and fixing functionality.
## wrappers.py (NEW FILE: +531 lines)
- Created dedicated module for ALL node wrapper/override functions
- Consolidated 10 wrapper types from 3 different files into single location:
* DisTorch V2 SafeTensor wrappers (factory + 3 implementations)
* DisTorch V1 legacy wrappers (4 GGUF/CLIP wrappers, rewritten to call V2 backend)
* Standard MultiGPU wrappers (3 device selection wrappers)
- CRITICAL FIX: All wrappers now strip MultiGPU-specific parameters before calling
original ComfyUI functions (fixes CheckpointLoaderSimple TypeError)
- Improved architecture: clear separation between wrapper UI and backend logic
## distorch.py (DELETED: -529 lines)
- Removed entire legacy DisTorch V1 file
- All V1 wrapper functions moved to wrappers.py and rewritten to call V2 backend
- Backend allocation functions no longer needed (V2 backend handles all cases)
- Eliminates code duplication and maintenance burden
## distorch_2.py (-409 lines)
- Removed duplicate _create_distorch_safetensor_v2_override factory function
(was incorrectly present in both distorch_2.py and wrappers.py)
- Removed 3 wrapper export functions (moved to wrappers.py)
- File now contains ONLY backend logic:
* register_patched_safetensor_modelpatcher()
* analyze_safetensor_loading() and analyze_safetensor_loading_clip()
* calculate_safetensor_vvram_allocation()
* Allocation stores and model hash functions
- Added clear documentation comment about wrapper migration
## __init__.py (-230 lines)
- Removed 3 local wrapper function definitions (moved to wrappers.py)
- Removed soft_empty_cache_distorch2_patched (moved to device_utils.py)
- Removed all distorch.py imports (file deleted)
- Added imports from new wrappers.py module (10 wrapper functions)
- Updated imports from distorch_2.py (backend functions only, no wrappers)
- Improved architecture: __init__.py now focused on initialization and registration
## device_utils.py (+68 lines)
- Moved soft_empty_cache_distorch2_patched() from __init__.py
- Added comprehensive memory management patch in architecturally correct location
- Patch includes:
* DisTorch2 detection and multi-device VRAM management
* Adaptive CPU memory threshold checking
* Force flag support for executor cache reset (Manager parity)
- Applied patch at module level: mm.soft_empty_cache = soft_empty_cache_distorch2_patched
- Behavior preserved: patch still executes when device_utils is imported by __init__.py
## nodes.py (-30 lines)
- Removed unused wrapper function imports
- Cleaned up import statements to reflect new architecture
## Impact Summary
- Improved architecture: Clear separation between wrappers (UI) and backend (logic)
- Eliminated distorch.py: Reduced from 3 files to 2 (wrappers.py + distorch_2.py)
- Net code reduction: 531 lines removed while adding functionality
- Better maintainability: Single source of truth for all wrapper functions
- Preserved behavior: All patches execute correctly, no functional changes
## Breaking Changes
None - this is a pure refactor with no API or behavioral changes.
Problem: Models correctly categorized as "keep loaded" during selective
unload were disappearing before the next cleanup cycle. After reassigning
mm.current_loaded_models = kept_models, Python's garbage collector would
clear the models because the list was their only remaining strong reference.
Solution: Implement GC anchor system using a global set to hold strong
references to ModelPatcher objects that must survive cleanup cycles.
Changes:
- Add _MGPU_RETENTION_ANCHORS global set and helper functions
- Add early delegation check: if no DisTorch models want unload, clear
anchors and delegate to original unload_all_models
- Add retention anchor when categorizing kept models
- Clear anchors before delegating to allow normal cleanup
Result: Self-contained, reversible protection mechanism. Models with
keep_loaded=True survive automatic cleanup but can be cleared with
explicit "Clear All Models" button. Tested on both keep_loaded=True
and keep_loaded=False scenarios.
Refs: memory-bank/distorch_selective_unload_solution.md
Add comprehensive documentation for the DisTorch selective model unload
solution that addresses Python garbage collection issues. The document
details:
- Problem: Models with keep_loaded=False were being prematurely garbage
collected despite being added to current_loaded_models list
- Root cause: Reassigning current_loaded_models created the only strong
reference, making models vulnerable to GC between assignment and next
access
- Solution: Global GC anchor set (_MGPU_RETENTION_ANCHORS) maintains
strong references to ModelPatcher objects that should survive cleanup
- Implementation: Early delegation check, anchor protection during
categorization, and explicit lifecycle management
- Testing results: Confirms Flux unloads while VAE/CLIP remain protected
This solution ensures DisTorch models can selectively unload while
keeping VAE and CLIP models loaded, preventing memory management race
conditions with Python's garbage collector.
Update documentation to reflect that Phase 3 selective ejection is now
fully deployed in code, including per-model flags, patched unload_all_models,
and Manager parity via force_full_system_cleanup.
Clarify understanding of remaining bug:
- Remove incorrect hypothesis that "all-kept delegation" causes the issue
- Document that delegation to original unload_all_models is intentional
behavior needed for post-execution cleanup
- Simplify bug description to acknowledge root cause is still unknown
- Focus on observable symptom: retained models (keep_loaded=True) are
still being ejected despite selective unload logic being present
This commit corrects misleading documentation and removes false leads
to help focus investigation on the actual unknown root cause.
- Implement caching of the last computed settings hash using a class attribute `_last_hash`
- Compare current hash against the cached one to detect changes
- Add logging to indicate first call or when settings have changed, using shortened hash for brevity
- Applied consistently across `override_class_with_distorch_safetensor_v2`, `_v2_clip`, and `_v2_clip_no_device`
- Improves efficiency by avoiding redundant change detection and aids debugging of settings modifications
- Renamed `keep_loaded` variable to `should_retain` for improved clarity
- Simplified assignment by directly retrieving `_mgpu_keep_loaded` attribute with default False
- Updated logging accordingly; may alter behavior for non-DisTorch models to no longer retain automatically
- Modified condition to retain models lacking `_mgpu_keep_loaded` attribute or with `keep_loaded=True`
- Improves reliability of unloading by distinguishing DisTorch and non-DisTorch models
- Addresses potential premature unloading of intended persistent models in multi-GPU setups
feat: add configurable multi-GPU memory cleanup policies
- Add MULTIGPU_CLEANUP_POLICY environment variable with options: off, threshold, every_load, every_load+threshold
- Add MULTIGPU_CPU_RESET_THRESHOLD for memory threshold-based cleanup (default 0.85)
- Add MULTIGPU_MALLOC_TRIM toggle to control malloc trimming behavior
- Implement cleanup triggers in load_models_gpu based on configured policy
- Make malloc trim conditional in soft_empty_cache_distorch2_patched
- Add configuration logging for better observability
This allows users to customize memory management behavior for multi-GPU setups through
Establish comprehensive development guidelines including project overview,
memory bank documentation requirements, technical patterns, and current
status for ComfyUI-MultiGPU contributors. Includes critical CPU memory
leak investigation details and mandated development philosophy.
- Introduce MGPU_MM_LOG flag and logger.mgpu_mm_log(...) to gate and
prefix MultiGPU Model Management logs (disabled by default)
- Replace ad-hoc logger.info("[MultiGPU ...]") calls with mgpu_mm_log
in DisTorch2 cache-clearing and delegation paths to reduce noise
- In load_models_gpu, parse safetensor allocation strings to infer
incoming_compute_device and incoming_compute_planned_bytes (supports
hash#device;GB and expert fraction syntax); track required bytes
- Remove coarse large-model threshold heuristic in favor of allocation-
informed planning
Why: centralize and quiet verbose MGPU logs by default, and enable
smarter, data-driven device selection and memory planning for multi-GPU
model loading.
- Introduce MEMORY_LOG flag and logger.memory method to gate high-volume memory logs
- Demote device setter logs from info to debug to reduce noise
- Clarify patch announcement (remove text_encoder_initial_device mention)
- Update soft_empty_cache patch log to emphasize multi-device allocation/clearing; delegate to original when DisTorch2 is inactive
- Rework load_models_gpu preflight for large DisTorch2 models:
- more robust ModelPatcher detection (direct or via .patcher)
- track allowed devices and incoming model names
- improved large-model detection and proactive unload/clearing on donor/offload devices
- mitigates OOM during large model (e.g., UNet) swaps
- Minor cleanup of verbose comments and wording in logs
- Add comfyui_memory_load and create_model_identifier utilities (device_utils)
- Log GPU memory before/after UNet, VAE, and CLIP construction and after UNet weight load
- Include model identifiers in logs to correlate memory to specific patchers
- Guard logging calls with try/except to avoid impacting load flow
- Improves observability of memory usage for multi-GPU checkpoints and aids OOM/debugging
multi-GPU cache clear + proactive unload to prevent OOM
- Patch mm.soft_empty_cache to clear caches on all GPUs when DisTorch2 models are active; otherwise delegate to original ComfyUI behavior. Uses safetensor allocation store and model hashes to detect DisTorch2 models; adds soft_empty_cache_multigpu import.
- Patch mm.load_models_gpu (guarded to apply once) to proactively unload large, unneeded models (>2GB) before loading large DisTorch2 models. Frees compute and donor device memory to prevent UNet OOM during model swaps.
- Preserve original functions for fallback, validate inputs, and log clearly to reduce risk during reloads and unexpected usage.
- Remove current_text_encoder_initial_device and its updates
- Delete text_encoder_initial_device_patched and stop overriding mm.text_encoder_initial_device
- Simplify set_current_text_encoder_device and logging to track only current_text_encoder_device
Bump revision to 2.4.7
John Pollock