Skip to content

Commit

Permalink
SD3 better memory usage estimation.
Browse files Browse the repository at this point in the history
  • Loading branch information
comfyanonymous committed Jun 12, 2024
1 parent 69c8d6d commit 694e0b4
Showing 1 changed file with 12 additions and 0 deletions.
12 changes: 12 additions & 0 deletions comfy/model_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -578,3 +578,15 @@ def extra_conds(self, **kwargs):
if cross_attn is not None:
out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
return out

def memory_required(self, input_shape):
if comfy.model_management.xformers_enabled() or comfy.model_management.pytorch_attention_flash_attention():
dtype = self.get_dtype()
if self.manual_cast_dtype is not None:
dtype = self.manual_cast_dtype
#TODO: this probably needs to be tweaked
area = input_shape[0] * input_shape[2] * input_shape[3]
return (area * comfy.model_management.dtype_size(dtype) * 0.012) * (1024 * 1024)
else:
area = input_shape[0] * input_shape[2] * input_shape[3]
return (area * 0.3) * (1024 * 1024)

0 comments on commit 694e0b4

Please sign in to comment.