298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-04-16 16:43:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix potential memory leak.

Browse Source
This commit is contained in:
comfyanonymous 2024-08-26 02:07:32 -04:00
parent 9230f65823
commit c6812947e9
1 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
comfy/model_patcher.py
View File

@ -319,12 +319,21 @@ class ModelPatcher:
mem_counter = 0 mem_counter = 0
patch_counter = 0 patch_counter = 0
lowvram_counter = 0 lowvram_counter = 0
load_completely = [] loading = []
for n, m in self.model.named_modules(): for n, m in self.model.named_modules():
if hasattr(m, "comfy_cast_weights") or hasattr(m, "weight"):
loading.append((comfy.model_management.module_size(m), n, m))
load_completely = []
loading.sort(reverse=True)
for x in loading:
n = x[1]
m = x[2]
module_mem = x[0]
lowvram_weight = False lowvram_weight = False
if not full_load and hasattr(m, "comfy_cast_weights"): if not full_load and hasattr(m, "comfy_cast_weights"):
module_mem = comfy.model_management.module_size(m)
if mem_counter + module_mem >= lowvram_model_memory: if mem_counter + module_mem >= lowvram_model_memory:
lowvram_weight = True lowvram_weight = True
lowvram_counter += 1 lowvram_counter += 1
@ -356,9 +365,8 @@ class ModelPatcher:
wipe_lowvram_weight(m) wipe_lowvram_weight(m)
if hasattr(m, "weight"): if hasattr(m, "weight"):
mem_used = comfy.model_management.module_size(m) mem_counter += module_mem
mem_counter += mem_used load_completely.append((module_mem, n, m))
load_completely.append((mem_used, n, m))
load_completely.sort(reverse=True) load_completely.sort(reverse=True)
for x in load_completely: for x in load_completely: