Some fixes to the lowvram system.

This commit is contained in:
comfyanonymous 2024-11-22 16:40:04 -05:00
parent 94323a26a7
commit bc6be6c11e

View File

@ -374,9 +374,14 @@ class ModelPatcher:
loading = []
for n, m in self.model.named_modules():
params = []
skip = False
for name, param in m.named_parameters(recurse=False):
params.append(name)
if hasattr(m, "comfy_cast_weights") or len(params) > 0:
for name, param in m.named_parameters(recurse=True):
if name not in params:
skip = True # skip random weights in non leaf modules
break
if not skip and (hasattr(m, "comfy_cast_weights") or len(params) > 0):
loading.append((comfy.model_management.module_size(m), n, m, params))
load_completely = []
@ -420,6 +425,7 @@ class ModelPatcher:
if m.comfy_cast_weights:
wipe_lowvram_weight(m)
if full_load or mem_counter + module_mem < lowvram_model_memory:
mem_counter += module_mem
load_completely.append((module_mem, n, m, params))