From 3782b166066f61270668e32fff1c34703ee9a71d Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 15 Dec 2023 00:33:25 +0100
Subject: [PATCH] Re add BF16 for VAE

---
 VAE/nodes.py   | 2 +-
 utils/dtype.py | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/VAE/nodes.py b/VAE/nodes.py
index c5c10f9..b5ad1ca 100644
--- a/VAE/nodes.py
+++ b/VAE/nodes.py
@@ -12,7 +12,7 @@ class ExtraVAELoader:
 			"required": {
 				"vae_name": (folder_paths.get_filename_list("vae"),),
 				"vae_type": (list(vae_conf.keys()), {"default":"kl-f8"}),
-				"dtype"   : (dtype_list_short,),
+				"dtype"   : ([*dtype_list_short, "BF16"],),
 			}
 		}
 	RETURN_TYPES = ("VAE",)
diff --git a/utils/dtype.py b/utils/dtype.py
index a5bf5fb..03121cd 100644
--- a/utils/dtype.py
+++ b/utils/dtype.py
@@ -38,6 +38,8 @@ def string_to_dtype(s="none", mode=None):
 		return None
 	elif s in ["fp32", "float32", "float"]:
 		return torch.float32
+	elif s in ["bf16", "bfloat16"]:
+		return torch.bfloat16
 	elif s in ["fp16", "float16", "half"]:
 		return torch.float16
 	elif "fp8" in s or "float8" in s: