Skip to content

Commit

Permalink
Make memory estimation aware of model dtype.
Browse files Browse the repository at this point in the history
  • Loading branch information
comfyanonymous committed Nov 12, 2023
1 parent dd4ba68 commit 4781819
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion comfy/model_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -161,7 +161,7 @@ def memory_required(self, input_shape):
area = input_shape[0] * input_shape[2] * input_shape[3]
if comfy.model_management.xformers_enabled() or comfy.model_management.pytorch_attention_flash_attention():
#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:
#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)
Expand Down

0 comments on commit 4781819

Please sign in to comment.