Fixed Memory Leaks: Added explicit cleanup of partial replicas if cloning fails mid-process
Better Error Handling: Exception handling in worker() now properly returns exceptions rather than raising, allowing cleanup of other threads
Thread Safety: Removed unused imports and ensured proper synchronization points for CUDA/XPU devices
Robust Splitting: Fixed split_batch to handle mixed lists of tensors and non-tensors correctly
Cleanup Safety: cleanup_parallel_model now handles cases where the model was partially initialized
VRAM Management: Added soft_empty_cache calls between cloning operations to prevent OOM when cloning to multiple devices
Management: Added auto balance to improve performance
Attribute Handling: Improved clone_module_simple to better handle missing attributes in layer reconstruction
Config Extraction: Added serialization check to extract_model_config to prevent storing non-serializable objects
Critical Bug Fixes:
Memory Leak: weakref.finalize now uses weakref.ref(target_model) instead of direct reference
Race Condition: Worker threads now return exceptions instead of raising them, allowing proper aggregation and cleanup
Batch Validation: get_batch_size() now validates consistent batch dimensions across tensor lists
Double Device Transfer: Removed .to(device) from split_kwargs(), now only happens in worker()
Attribute Handling: Cleaner separation of in_features vs in_channels in layer cloning
Improvements Added:
CUDA Streams: Each GPU now uses its own CUDA stream for true parallel execution
Auto VRAM Balance: New auto_vram_balance parameter adjusts splits based on available VRAM (70% user preference / 30% availability)
Gradient Checkpointing: Automatically disabled on replicas to save VRAM
Accelerate Hooks: Clears _hf_hook and hooks attributes for compatibility with accelerate offloading
Enhanced Cache Clearing: Added rope_cache and freqs_cis_cache to FLUX cleanup
Thread Safety: Proper executor context management and result validation
Performance Optimizations:
Non-blocking transfers where safe
Synchronization only on CUDA/XPU devices
Batch size auto-adjustment prevents OOM on uneven GPU setups
Fixed "'list' object has no attribute 'shape'" error when using
LTX-2 Video and other diffusion models that return tuples/lists
instead of single tensors.
Key changes:
- Added tensor container helpers: get_batch_size(), split_batch(),
move_to_device(), concatenate_results()
- parallel_forward now handles list/tuple inputs/outputs (common in
video models with temporal conditioning)
- Enhanced split_kwargs() to split list-of-tensors kwargs
(frame-wise embeddings, temporal IDs)
- Results are now properly concatenated whether model returns
Tensor, Tuple[Tensor, ...], or List[Tensor]
Tested working with: FLUX, SD3.5, LTX-2 Video, HunyuanVideo
VRAM: Maintains existing N× model replication strategy
Breaking: None (backward compatible with single-tensor outputs)
FearL0rd