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

Make memory estimation aware of model dtype.

Browse Source
This commit is contained in:
comfyanonymous 2023-11-12 04:26:16 -05:00
parent dd4ba68b6e
commit 4781819a85
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
comfy/model_base.py
View File

@ -161,7 +161,7 @@ class BaseModel(torch.nn.Module):
area = input_shape[0] * input_shape[2] * input_shape[3] area = input_shape[0] * input_shape[2] * input_shape[3]
if comfy.model_management.xformers_enabled() or comfy.model_management.pytorch_attention_flash_attention(): if comfy.model_management.xformers_enabled() or comfy.model_management.pytorch_attention_flash_attention():
#TODO: this needs to be tweaked #TODO: this needs to be tweaked
return (area / 20) * (1024 * 1024) return (area / (comfy.model_management.dtype_size(self.get_dtype()) * 10)) * (1024 * 1024)
else: else:
#TODO: this formula might be too aggressive since I tweaked the sub-quad and split algorithms to use less memory. #TODO: this formula might be too aggressive since I tweaked the sub-quad and split algorithms to use less memory.
return (((area * 0.6) / 0.9) + 1024) * (1024 * 1024) return (((area * 0.6) / 0.9) + 1024) * (1024 * 1024)