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

Load weights that can't be lowvramed to target device.

Browse Source
This commit is contained in:
comfyanonymous 2023-12-28 21:41:10 -05:00
parent a8baa40d85
commit e1e322cf69
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_management.py
View File

@ -259,6 +259,14 @@ print("VAE dtype:", VAE_DTYPE)
current_loaded_models = [] current_loaded_models = []
def module_size(module):
module_mem = 0
sd = module.state_dict()
for k in sd:
t = sd[k]
module_mem += t.nelement() * t.element_size()
return module_mem
class LoadedModel: class LoadedModel:
def __init__(self, model): def __init__(self, model):
self.model = model self.model = model
@ -296,14 +304,14 @@ class LoadedModel:
if hasattr(m, "comfy_cast_weights"): if hasattr(m, "comfy_cast_weights"):
m.prev_comfy_cast_weights = m.comfy_cast_weights m.prev_comfy_cast_weights = m.comfy_cast_weights
m.comfy_cast_weights = True m.comfy_cast_weights = True
module_mem = 0 module_mem = module_size(m)
sd = m.state_dict()
for k in sd:
t = sd[k]
module_mem += t.nelement() * t.element_size()
if mem_counter + module_mem < lowvram_model_memory: if mem_counter + module_mem < lowvram_model_memory:
m.to(self.device) m.to(self.device)
mem_counter += module_mem mem_counter += module_mem
elif hasattr(m, "weight"): #only modules with comfy_cast_weights can be set to lowvram mode
m.to(self.device)
mem_counter += module_size(m)
print("lowvram: loaded module regularly", m)
self.model_accelerated = True self.model_accelerated = True