- 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.
This commit introduces DisTorch v2.0.0, a major overhaul that extends multi-device model distribution to standard `.safetensors` models.
Key changes include:
- **Universal `.safetensors` Support:** The core distribution logic is no longer limited to GGUF models. It now fully supports `.safetensors`, allowing any UNet supported by native Comfy loaders to have its layers distributed across multiple devices (GPUs and CPU/RAM).
This commit refactors the model loading logic to properly integrate with ComfyUI's caching system.
- Implemented the `IS_CHANGED` class method, which creates a hash of the DisTorch-specific settings (e.g., `compute_device`, `virtual_vram_gb`).
- This allows ComfyUI to automatically detect when settings have changed and trigger a model reload, invalidating the cache correctly.
- Removed the previous manual and less reliable logic for unloading and reloading the model from within the `override` function.
- Set the default log level to "Engineering" to provide more detailed output during development.
This commit introduces a major architectural refactoring, laying the groundwork for DisTorch V2. The changes focus on improving modularity, memory management, and diagnostics.
Key changes include:
- Renaming `distorch_safetensor.py` to `distorch_2.py` to house the new core logic.
- Deleting the legacy `block_swap.py` module.
- Adding `device_memory_audit.py` for more sophisticated analysis of GPU memory usage.
- Implementing a centralized and configurable logging system in `__init__.py` to provide standardized and level-controlled (DEBUG/INFO) output for better debugging.
- Add distorch_safetensor.py module with safetensor allocation and model hashing utilities
- Update all DisTorch2 nodes to use new override_class_with_distorch_safetensor_v2
- Add FLUX-specific DisTorch2 nodes for checkpoint and UNET loading
- Import new safetensor functions for VRAM allocation analysis and model patching
This commit introduces a new block swapping mechanism specifically for WanVideo models to enable running them on GPUs with limited VRAM.
A new `WanVideoBlockSwapManager` is implemented which uses a pre-allocation or "shell" strategy. Instead of moving entire blocks between CPU and GPU, this approach:
1. Pre-allocates a single "shell" block on the GPU, sized to match the largest block in the model.
2. Offloads designated model blocks to the CPU.
3. Patches the `forward` method of these offloaded blocks.
4. During inference, the patched method copies the weights (`state_dict`) from the CPU block into the GPU shell just before execution.
This method avoids the overhead of allocating and deallocating GPU memory for each block, reducing memory fragmentation and potentially improving performance and or corruption copying potentially modified blocks back to the swap space.
This commit introduces block swapping functionality for Qwen models, enabling them to run on systems with limited VRAM by offloading layers to a swap device (e.g., CPU RAM).
Key changes:
- A new `QwenBlockSwapManager` class is implemented to handle the patching of Qwen transformer blocks.
- The `apply_block_swap` function is extended to detect Qwen models and apply the swapping logic to their `transformer_blocks`.
- A model signature for Qwen is added to `model_sig.py` to correctly identify the swappable modules.
- A new diagnostic function, `log_unsupported_model_analysis`, is added to log the structure of unsupported models, aiding future development.
The memory analysis function, `analyze_safetensor_distorch`, has been improved to provide a more accurate and detailed report.
Instead of estimating the number of swapped blocks based on an average size, the function now receives the actual list of blocks being swapped. It generates a per-block table detailing each block's ID, type, size, and its final assignment (COMPUTE or SWAP).
This provides users with a precise breakdown of the memory offload, reflecting the actual state of the model rather than a theoretical calculation.
Additionally, the unused `log_memory_usage` helper function has been removed.
This commit completely rewrites the block swapping implementation for improved stability, correctness, and code structure.
Key changes:
- Replaces the fragile monkey-patching of the `forward` method with the standard PyTorch `register_forward_pre_hook`.
- Introduces a `BlockSwapManager` class to encapsulate all swapping logic, separating it from the ComfyUI node.
- Implements a "Sequential Swapping" strategy: the previously active block is offloaded before the current block is loaded, ensuring only one block is on the active device at a time.
- Adds a `cleanup` method to properly remove hooks after execution, preventing state leakage between runs.
- Fixes a critical bug where the hook signature was incorrect.
- Adds a memory logging utility for easier debugging.
This commit removes the unused `reserved_swap_gb` parameter from the `analyze_safetensor_distorch` function and its call sites. This simplifies the function's signature and cleans up the analysis output by removing the "Reserve" metric, which was always zero.
Additionally, the `.gitignore` file is simplified by removing entries for the `binaries/` directory, which are no longer needed.
This commit introduces a major update, "DisTorch v2", which integrates the new `BlockSwap` system for more efficient and dynamic memory management across multiple GPUs.
Key changes:
- **BlockSwap Integration:** GGUF model loading is completely refactored to use `BlockSwap`, enabling more intelligent VRAM allocation based on tensor analysis.
- **Node Renaming:** All SafeTensor loader nodes are renamed from `...DisTorchMultiGPU` to `...DisTorch2MultiGPU` to clearly distinguish the new implementation from the old one.
- **Legacy Support:** The previous GGUF loader is preserved as a legacy option for backward compatibility.
- **Improved Memory Calculation:** A more accurate memory calculation function (`get_total_memory_v2`) is implemented and used by the new system.
This commit refactors the codebase by extracting major components from the main `__init__.py` file into their own dedicated modules. This improves code organization, readability, and maintainability.
- **`distorch.py`**: New file containing the `DisTorch` class, which manages multi-GPU device patching and distribution logic.
- **`block_swap.py`**: New file containing the generic `BlockSwap` class for UNet block swapping to manage VRAM.
- **`wanvideo.py`**: New file containing the `WanVideoBlockSwap` class, a specialized implementation for WanVideo models.
- **`__init__.py`**: Simplified to handle node registration and imports from the new modules.
- Added comprehensive file logging with timestamps
- Log file created in logs/distorch_TIMESTAMP.log
- Debug logging for all major operations:
- Model structure analysis
- Memory calculations
- Block partitioning strategy
- Device movements and swaps
- Hook installations
- Performance metrics
- Console output for important INFO messages
- Full exception tracebacks captured
- Simplified node naming from DisTorchBlockSwap to DisTorch
- Cleaned up accidentally added main_branch directory
- Updated all references in __init__.py and core/blockswap.py
- Created comprehensive architecture documentation (ARCHITECTURE_V2.0.0.md)
- Added DOE optimization planning document (DOE_OPTIMIZATION.md)
- Implemented DisTorchBlockSwap node for safetensor models
- Created core/blockswap.py with BlockSwapManager
- Unified VirtualVRAM interface design
- Based on analysis of WanVideo's block swap methodology
Populate 'type' options by sourcing from core nodes to avoid drift:\n- CLIPLoaderGGUF now derives 'type' from nodes.CLIPLoader.INPUT_TYPES()\n- DualCLIPLoaderGGUF now derives 'type' from nodes.DualCLIPLoader.INPUT_TYPES()\nThis fixes missing or outdated 'type' options in GGUF Single and Dual CLIP loaders.\n\nchore: bump version to 1.8.2
Add guarded Intel XPU support alongside CUDA:
- get_device_list now includes xpu:N when available
- device selection (model/text encoder) considers CUDA or XPU and validates devices
- DisTorch donor/offload selection includes xpu devices
Also: remove unused MergeFluxLoRAs node and mapping; delete tools/ and precompiled_binaries/; bump project version to 1.8.1.
John Pollock