From 3152023fbc4f8ee6598a863314ca98d48ea9c2e6 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 11 Dec 2023 23:50:38 -0500 Subject: [PATCH] Use inference dtype for unet memory usage estimation. --- comfy/model_base.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/comfy/model_base.py b/comfy/model_base.py index 412c8379..a7582b33 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -177,9 +177,12 @@ class BaseModel(torch.nn.Module): 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 needs to be tweaked area = input_shape[0] * input_shape[2] * input_shape[3] - return (area * comfy.model_management.dtype_size(self.get_dtype()) / 50) * (1024 * 1024) + return (area * comfy.model_management.dtype_size(dtype) / 50) * (1024 * 1024) else: #TODO: this formula might be too aggressive since I tweaked the sub-quad and split algorithms to use less memory. area = input_shape[0] * input_shape[2] * input_shape[3]