From e1635e999678a00afefb0ab31a7bafa95f22ccbf Mon Sep 17 00:00:00 2001 From: John Pollock Date: Tue, 9 Sep 2025 12:00:56 -0500 Subject: [PATCH] Improve memory handling for safetensor models in corner cases - Added preemptive model unloading and cache clearing in register_patched_safetensor_modelpatcher() to resolve potential memory issues when allocations are unavailable, prompting the usage of the standard loaders. --- distorch_2.py | 4 ++++ pyproject.toml | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/distorch_2.py b/distorch_2.py index 66e6cbc..6a54ed1 100644 --- a/distorch_2.py +++ b/distorch_2.py @@ -67,10 +67,14 @@ def register_patched_safetensor_modelpatcher(): allocations = safetensor_allocation_store.get(debug_hash) if not hasattr(self.model, '_distorch_high_precision_loras') or not allocations: + mm.unload_all_models() + soft_empty_cache_multigpu(logger) result = original_partially_load(self, device_to, extra_memory, force_patch_weights) if hasattr(self, '_distorch_block_assignments'): del self._distorch_block_assignments return result + + mm.unload_all_models() soft_empty_cache_multigpu(logger) mem_counter = 0 diff --git a/pyproject.toml b/pyproject.toml index 49327bc..09d9e5c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-multigpu" description = "Provides a suite of custom nodes to manage multiple GPUs for ComfyUI, including advanced model offloading for both GGUF and Safetensor formats with DisTorch, and bespoke MultiGPU support for WanVideoWrapper and other custom nodes." -version = "2.4.2" +version = "2.4.3" license = {file = "LICENSE"} [project.urls]