Add dtype selection to T5

VAE/loader.py

@@ -4,20 +4,13 @@ import comfy.utils
 from comfy import model_management
 from comfy import diffusers_convert
 
-vae_dtype_dict = {
-	"auto" : model_management.vae_device(),
-	"fp32" : torch.float32,
-	"fp16" : torch.float16,
-	"bf16" : torch.bfloat16,
-}
-
 class EXVAE(comfy.sd.VAE):
-	def __init__(self, model_path, model_conf, dtype=None):
+	def __init__(self, model_path, model_conf, dtype=torch.float32):
 		self.latent_dim   = model_conf["embed_dim"]
 		self.latent_scale = model_conf["embed_scale"]
 		self.device = model_management.vae_device()
 		self.offload_device = model_management.vae_offload_device()
-		self.vae_dtype = vae_dtype_dict.get(dtype, "auto")
+		self.vae_dtype = dtype
 
 		sd = comfy.utils.load_torch_file(model_path)
 		model = None
@@ -104,7 +97,7 @@ class EXVAE(comfy.sd.VAE):
 			free_memory = model_management.get_free_memory(self.device)
 			batch_number = int(free_memory / memory_used)
 			batch_number = max(1, batch_number)
-			samples = torch.empty((pixel_samples.shape[0], self.first_stage_model.embed_dim, round(pixel_samples.shape[2] // self.latent_scale), round(pixel_samples.shape[3] // self.latent_scale)), device="cpu")
+			samples = torch.empty((pixel_samples.shape[0], self.latent_dim, round(pixel_samples.shape[2] // self.latent_scale), round(pixel_samples.shape[3] // self.latent_scale)), device="cpu")
 			for x in range(0, pixel_samples.shape[0], batch_number):
 				pixels_in = (2. * pixel_samples[x:x+batch_number] - 1.).to(self.vae_dtype).to(self.device)
 				samples[x:x+batch_number] = self.first_stage_model.encode(pixels_in).cpu().float()

VAE/nodes.py

@@ -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 = {