Skip to content

Commit e1635e9

Browse files
committed
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.
1 parent 803cf54 commit e1635e9

2 files changed

Lines changed: 5 additions & 1 deletion

File tree

distorch_2.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,10 +67,14 @@ def new_partially_load(self, device_to, extra_memory=0, full_load=False, force_p
6767
allocations = safetensor_allocation_store.get(debug_hash)
6868

6969
if not hasattr(self.model, '_distorch_high_precision_loras') or not allocations:
70+
mm.unload_all_models()
71+
soft_empty_cache_multigpu(logger)
7072
result = original_partially_load(self, device_to, extra_memory, force_patch_weights)
7173
if hasattr(self, '_distorch_block_assignments'):
7274
del self._distorch_block_assignments
7375
return result
76+
77+
mm.unload_all_models()
7478
soft_empty_cache_multigpu(logger)
7579

7680
mem_counter = 0

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
[project]
22
name = "comfyui-multigpu"
33
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."
4-
version = "2.4.2"
4+
version = "2.4.3"
55
license = {file = "LICENSE"}
66

77
[project.urls]

0 commit comments

Comments
 (0)