From 694e0b48e0f6d55ddfcbe44a5d2818774241077b Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Wed, 12 Jun 2024 00:49:00 -0400 Subject: [PATCH] SD3 better memory usage estimation. --- comfy/model_base.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/comfy/model_base.py b/comfy/model_base.py index 28458bba..7ef03440 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -578,3 +578,15 @@ class SD3(BaseModel): 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)