Add dtype selection to T5

This commit is contained in:
City
2023-11-28 00:49:06 +01:00
parent 33e5db52a3
commit 389c16f2f5
8 changed files with 152 additions and 29 deletions
+4 -2
View File
@@ -3,6 +3,8 @@ import folder_paths
from .conf import vae_conf
from .loader import EXVAE
from ..utils.dtype import string_to_dtype, dtype_list_short
class ExtraVAELoader:
@classmethod
def INPUT_TYPES(s):
@@ -10,7 +12,7 @@ class ExtraVAELoader:
"required": {
"vae_name": (folder_paths.get_filename_list("vae"),),
"vae_type": (list(vae_conf.keys()), {"default":"kl-f8"}),
"dtype" : (["auto","fp32","fp16","bf16"],),
"dtype" : (dtype_list_short,),
}
}
RETURN_TYPES = ("VAE",)
@@ -21,7 +23,7 @@ class ExtraVAELoader:
def load_vae(self, vae_name, vae_type, dtype):
model_path = folder_paths.get_full_path("vae", vae_name)
model_conf = vae_conf[vae_type]
vae = EXVAE(model_path, model_conf, dtype)
vae = EXVAE(model_path, model_conf, string_to_dtype(dtype, "vae"))
return (vae,)
NODE_CLASS_MAPPINGS = {