qnloft
/
ComfyUI
mirror of https://github.com/comfyanonymous/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Some VAEs come in .pt files.

This commit is contained in:
comfyanonymous 2023-01-28 12:28:29 -05:00
parent d133cf4f06
commit 2706c0b7a5
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
nodes.py
View File

@ -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"