- Implemented WanVideoVACEEncode class for encoding with VACE, including input parameters for width, height, number of frames, and strength.
- Enhanced WanVideoEncode class to support multi-GPU encoding, with additional parameters for noise augmentation and latent strength.
- Updated device management to ensure proper device context during encoding processes.
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.
John Pollock