- 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
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.
- 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.
- 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.
- Extend module docstring to include inspection capabilities
- Add create_model_identifier() to generate unique hashes from model type and size
- Add analyze_tensor_locations() to analyze tensor device placement and memory usage
- Include imports for hashlib, psutil, and comfy.model_management to support new features
These utilities enable end-to-end tracking of model state and placement for better debugging and management in multi-GPU setups.
Add comprehensive memory cache clearing aligned with ComfyUI patterns to improve stability and reduce OOM incidents in multi-device scenarios.
**Addresses Memory/Garbage Collection Issues:**
- Created `soft_empty_cache_multigpu()` function in device_utils.py
- Replicates ComfyUI's cache clearing for all devices (CUDA, MPS, XPU, NPU, MLU)
- Includes CUDA IPC collect optimization like ComfyUI
- Strategically placed calls before major memory allocations
**Addresses CLIP loading issues:**
- Fixed DisTorch2 device device varibale management before text encoder operations
**`soft_empty_cache_multigpu()` implementation Aligned with ComfyUI's Patterns:**
- Called after GC operations
- Placed before major memory allocations
- Matches ComfyUI's proven memory management strategy
- Same device clearing logic for multi-device scenarios
Refactor device detection into dedicated utility module
- Extract device enumeration and compatibility checks to device_utils.py
- Add support for additional device types (NPU, MLU, DirectML, CoreX)
- Update all modules to use centralized device utilities
- Implement caching for device list to improve performance
- Reduce code duplication across distorch, nodes, and wanvideo modules
John Pollock