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
Advanced Checkpoint Loaders allow users to map each of the elements of the checkpoint to a different device, or in the case of DisTorch2, shard the UNet and CLiP .safetensors arbitrarily whilst ensuring actual computation remains on selected `compute` device.
Added example workflow for standard and DisTorch2 MultGPU checkpoint loaders.
This commit introduces two main improvements: refactoring the checkpoint loading mechanism and fixing the initial device placement for the text encoder (CLIP).
1. **Fix Text Encoder Device Handling:**
- A new patch is applied to `mm.text_encoder_initial_device` to gain control over the device used when the text encoder is first loaded.
- The `CLIPLoader` override now forces `device='default'` to ensure ComfyUI's patching mechanism is triggered correctly, preventing the text encoder from being incorrectly placed on the wrong GPU.
2. **Refactor Checkpoint Loaders:**
- Removed the global stores (`checkpoint_dtype_store`, `checkpoint_half_store`, `checkpoint_config_store`).
- The `CheckpointLoaderSimpleMultiGPU` and `AdvCheckpointLoaderMultiGPU` nodes now use arguments and ComfyUI's internal defaults directly. This simplifies the logic, reduces global state, and makes the code easier to follow.
Additionally, log message prefixes have been updated to be more descriptive, aiding in debugging.
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
- Adding a new `fantasyportrait_model` input to support FantasyPortrait models.
- Renaming the `vace_model` input to a more generic `extra_model` to allow loading other auxiliary models like VACE or MTV Crafter.
- Correcting the node type for the `fantasytalking_model` from `FANTASYTALKINGMODEL` to `FANTASYTALKMODEL`.
Fix compute device inclusion in expert mode allocations
Include compute device in vram_string when expert_mode_allocations
is set but virtual_vram_gb is 0. This ensures the compute device is
properly specified in the full allocation string for expert mode
configurations without virtual VRAM.
Bump version to 2.2.1
This commit refactors the DisTorch safetensor loading and allocation logic for improved performance and correctness.
The main changes to the block assignment are:
- The primary compute device is now included in the pool of "donor" devices, allowing for more holistic memory quota calculation across all available GPUs.
- Unassigned "orphan" blocks are now allocated to the compute device instead of the CPU. This keeps more of the model in VRAM, reducing potential bottlenecks.
Additionally, this commit:
- Fixes a bug in the byte expert string parser where the wildcard `*` was incorrectly checked in the device name instead of the value.
- Standardizes variable names like `allocations_string` for better code clarity and consistency.
Change the fallback assignment for tensor blocks that do not fit within any donor device's memory quota. These blocks are now assigned to the CPU instead of the primary compute device.
This prevents potential VRAM Out-Of-Memory errors on the main GPU, improving the stability of the model loading process, especially under tight memory constraints.
Additionally, the allocation log is now updated to include all available system devices, even those with zero allocation, to provide a more comprehensive report.
This commit refactors several aspects of the DisTorch2 device allocation logic to make it more robust, predictable, and easier to debug.
Key changes:
- Rework the byte-based allocation string parser (`calculate_fraction_from_byte_expert_string`). The new implementation correctly respects the user-defined device order and more cleanly handles the wildcard (`*`) for assigning remaining model parts.
-revert the "improvements" to the analyze safetensor loading routine causing it to catestrophically fail
This commit enhances the device handling logic within `analyze_safetensor_loading` for greater robustness and better user feedback.
Key changes:
- Dynamically discovers all available devices using `get_device_list` instead of only using devices from the allocation string. This prevents potential `KeyError` crashes when analyzing devices that are not part of the distribution plan.
- Changes the fallback device for unallocated model blocks from the primary compute device to the CPU. This is a safer default that prevents unexpectedly overloading the main GPU.
- Adds a warning log when a block falls back to the CPU, alerting the user to a possible misconfiguration in their allocation string.
Update the "Bytes" mode documentation in the README to specify that the CPU acts as the default wildcard device.
This change clarifies that if no `*` is explicitly used in the device allocation string, the remainder of the model will be automatically assigned to the CPU. This helps users better understand the default behavior and prevent confusion.
Update the documentation to include the new 'bytes' and 'ratio' expert modes for model allocation.
These new modes provide more intuitive, model-driven ways for users to control how models are split across multiple devices.
- Adds 'bytes' mode for direct allocation in GB/MB, similar to Huggingface's `device_map`.
- Adds 'ratio' mode for proportional splitting, inspired by llama.cpp.
- Rebrands the original expert mode as 'fraction' mode for clarity.
- Provides clear examples for all three expert modes.
Introduces a new expert allocation mode allowing users to define model distribution using absolute memory values (e.g., "8g", "512m"). This provides more direct and predictable control over how a model is split across devices compared to the percentage-based method.
The new allocation string format is `device,size;device,size;...`, for example: `"cuda:0,8g;cuda:1,4g;cpu*,2g"`.
Key features:
- A wildcard `*` designates a device to receive any remaining unallocated model parts.
- If requested allocations exceed the model size, they are pro-rated down.
- A new `parse_memory_string` utility handles flexible memory unit parsing (g, m, k, b).
Additionally, the device allocation summary table has been improved to be more descriptive, now showing total VRAM, percentage of device VRAM used, absolute model GB allocated, and the model distribution percentage.
- Remove redundant logging and counters in safetensor model patcher
- Add model original dtype detection for better precision handling
- Streamline FP8 casting conditions and remove verbose debug logs
- Improve static allocation parsing and device assignment flow
- Remove redundant references to GGUF patterns in comments for clarity
- Update logging prefixes from '[MULTIGPU_DISTORCHV2]' to '[MultiGPU_DisTorch2]' for consistency
- Reorganize code in `register_patched_safetensor_modelpatcher` to streamline allocation checks and device assignments
- Add force_patch_weights parameter to new_partially_load signature for better control
- Implement check for _distorch_high_precision_loras with fallback to original loading behavior
- Include cleanup for _distorch_block_assignments attribute
- Comment out debug logging statements to minimize noise during execution
Add logic to detect and unpatch weights for modules with comfy_cast_weights, introduce memory and patch counters, and integrate static device assignments from analyze_safetensor_loading to improve distributed safetensor loading efficiency.
Remove redundant comments and simplify compute device determination by importing and using `current_device` directly, improving code readability and streamlining the analysis logic for better efficiency in model allocation handling.
Simplify the analyze_safetensor_loading function by replacing manual model module iteration with ComfyUI's built-in _load_list() for calculating total memory and building block lists. This improves efficiency, reduces redundant code, and enhances compatibility with ComfyUI's internal mechanisms while maintaining accurate memory reporting and threshold filtering.
This commit removes several utility modules used for debugging, memory inspection, and hardware information gathering. These tools are no longer required and their removal simplifies the codebase.
The following files have been deleted:
- `debug_utils.py`
- `device_memory_audit.py`
- `hardware_info.py`
- `model_sig.py`
Additionally, the call to log memory usage on startup has been removed from `__init__.py`.
Update the README to specify that the "up to 10% faster GGUF inference" claim for DisTorch2 is a direct comparison against the previous DisTorch V1 implementation.
This clarification helps manage user expectations and provides a more accurate performance context.
Update the `get_device_list` function to detect and include the 'mps' (Metal Performance Shaders) backend if it's available through PyTorch.
This allows users on Apple Silicon hardware to see and select their GPU for accelerated computations.
This commit introduces a major reorganization of the `examples` directory to improve clarity and discoverability. Workflows are now grouped into subdirectories based on the features they demonstrate (e.g., `distorch`, `distorch2`, `gguf`, `multiGPU`).
Key changes:
- Moved existing example JSON files into new categorized folders.
- Added several new and updated workflows, particularly for DisTorch2.
- Removed outdated or redundant example files.
- Renamed an internal function from `..._gguf_v2` to `..._safetensor_v2` to better reflect its broader functionality in DisTorch2.
John Pollock