From 2706c0b7a5dfc70f357ebd366ee1431d3e83dc42 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sat, 28 Jan 2023 12:28:29 -0500 Subject: [PATCH] Some VAEs come in .pt files. --- nodes.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/nodes.py b/nodes.py index aaec64f8..7e737d0c 100644 --- a/nodes.py +++ b/nodes.py @@ -17,9 +17,11 @@ import comfy.samplers import comfy.sd supported_ckpt_extensions = ['.ckpt'] +supported_pt_extensions = ['.ckpt', '.pt'] try: import safetensors.torch supported_ckpt_extensions += ['.safetensors'] + supported_pt_extensions += ['.safetensors'] except: print("Could not import safetensors, safetensors support disabled.") @@ -132,7 +134,7 @@ class VAELoader: vae_dir = os.path.join(models_dir, "vae") @classmethod def INPUT_TYPES(s): - return {"required": { "vae_name": (filter_files_extensions(os.listdir(s.vae_dir), supported_ckpt_extensions), )}} + return {"required": { "vae_name": (filter_files_extensions(os.listdir(s.vae_dir), supported_pt_extensions), )}} RETURN_TYPES = ("VAE",) FUNCTION = "load_vae"