- Cleaned up unnecessary whitespace and comments in model_management_mgpu.py, nodes.py, wanvideo.py, and wrappers.py for better code clarity.
- Replaced list comprehensions with direct list conversions in nodes.py for efficiency.
- Updated memory logging format in model_management_mgpu.py to streamline data capture.
- Enhanced device management in wanvideo.py by ensuring consistent device setting and loading.
- Added linting configurations in pyproject.toml to enforce code quality standards.
- Removed unused imports and optimized existing ones across multiple files.
- Replace global safetensor_allocation_store/safetensor_settings_store and create_safetensor_model_hash
with a per-model annotation (_distorch_v2_meta) stored directly on the inner model object.
- Update distorch_2 to remove global stores and hash creation; parse and consume allocation strings
from inner_model._distorch_v2_meta during model registration and loading.
- Update wrappers, checkpoint_multigpu, device_utils, and __init__ to set and read the new metadata
instead of writing/reading global stores.
- Simplify detection of DisTorch-managed models (check inner_model._distorch_v2_meta) and adjust
logging to surface inner model ids and allocation info.
- Clean up related imports and dead code paths.
Files changed: distorch_2.py, wrappers.py, checkpoint_multigpu.py, device_utils.py, model_management_mgpu.py, __init__.py
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.
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
- 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
John Pollock