From ca82ade7652c80727b402f51a115feb5df4ad27a Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Mon, 4 Dec 2023 11:52:06 -0500
Subject: [PATCH] Use .itemsize to get dtype size for fp8.

---
 comfy/model_management.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/comfy/model_management.py b/comfy/model_management.py
index 18d15f9d..94d59696 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -430,6 +430,13 @@ def dtype_size(dtype):
     dtype_size = 4
     if dtype == torch.float16 or dtype == torch.bfloat16:
         dtype_size = 2
+    elif dtype == torch.float32:
+        dtype_size = 4
+    else:
+        try:
+            dtype_size = dtype.itemsize
+        except: #Old pytorch doesn't have .itemsize
+            pass
     return dtype_size
 
 def unet_offload_device():