Re add BF16 for VAE

This commit is contained in:
City
2023-12-15 00:33:25 +01:00
parent 169a9aea03
commit 3782b16606
2 changed files with 3 additions and 1 deletions
+1 -1
View File
@@ -12,7 +12,7 @@ class ExtraVAELoader:
"required": { "required": {
"vae_name": (folder_paths.get_filename_list("vae"),), "vae_name": (folder_paths.get_filename_list("vae"),),
"vae_type": (list(vae_conf.keys()), {"default":"kl-f8"}), "vae_type": (list(vae_conf.keys()), {"default":"kl-f8"}),
"dtype" : (dtype_list_short,), "dtype" : ([*dtype_list_short, "BF16"],),
} }
} }
RETURN_TYPES = ("VAE",) RETURN_TYPES = ("VAE",)
+2
View File
@@ -38,6 +38,8 @@ def string_to_dtype(s="none", mode=None):
return None return None
elif s in ["fp32", "float32", "float"]: elif s in ["fp32", "float32", "float"]:
return torch.float32 return torch.float32
elif s in ["bf16", "bfloat16"]:
return torch.bfloat16
elif s in ["fp16", "float16", "half"]: elif s in ["fp16", "float16", "half"]:
return torch.float16 return torch.float16
elif "fp8" in s or "float8" in s: elif "fp8" in s or "float8" in s: