Added VRAM purge toggles to INPUT_TYPES:
purge_cache: Clears CUDA cache during cleanup (default: True)
purge_models: Aggressively unloads all models from VRAM using comfy.model_management.unload_all_models() (default: False)
Updated setup_parallel method signature to accept purge_cache and purge_models parameters
Store purge preferences on the model: target_model._parallel_purge_cache and target_model._parallel_purge_models
Enhanced cleanup_parallel_model function to:
Read the purge flags from the model attributes
Conditionally run comfy.model_management.unload_all_models() if purge_models is True
Conditionally clear CUDA cache and soft empty cache based on purge_cache setting
Multi-GPU cache clearing (iterates through all CUDA devices)
The purge settings persist with the model and are applied when the parallel processing resources are cleaned up (either when the node is re-run with new settings or when the model is garbage collected).
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
FearL0rd