qnloft
/
ComfyUI
mirror of https://github.com/comfyanonymous/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Use function to calculate model size in model patcher.

This commit is contained in:
comfyanonymous 2023-12-28 21:46:20 -05:00
parent e1e322cf69
commit 12e822c6c8
1 changed files with 2 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
comfy/model_patcher.py
View File

@ -28,13 +28,9 @@ class ModelPatcher:
if self.size > 0: if self.size > 0:
return self.size return self.size
model_sd = self.model.state_dict() model_sd = self.model.state_dict()
size = 0 self.size = comfy.model_management.module_size(self.model)
for k in model_sd:
t = model_sd[k]
size += t.nelement() * t.element_size()
self.size = size
self.model_keys = set(model_sd.keys()) self.model_keys = set(model_sd.keys())
return size return self.size
def clone(self): def clone(self):
n = ModelPatcher(self.model, self.load_device, self.offload_device, self.size, self.current_device, weight_inplace_update=self.weight_inplace_update) n = ModelPatcher(self.model, self.load_device, self.offload_device, self.size, self.current_device, weight_inplace_update=self.weight_inplace_update)