From 4781819a85847a8cf180a41d0ee4cdf99979e5be Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Sun, 12 Nov 2023 04:26:16 -0500
Subject: [PATCH] Make memory estimation aware of model dtype.

---
 comfy/model_base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/comfy/model_base.py b/comfy/model_base.py
index f6de0b25..37bf24bb 100644
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@@ -161,7 +161,7 @@ class BaseModel(torch.nn.Module):
         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)