298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-01-25 15:55:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

Some VAEs come in .pt files.

Browse Source
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 import comfy.sd
supported_ckpt_extensions = ['.ckpt'] supported_ckpt_extensions = ['.ckpt']
supported_pt_extensions = ['.ckpt', '.pt']
try: try:
import safetensors.torch import safetensors.torch
supported_ckpt_extensions += ['.safetensors'] supported_ckpt_extensions += ['.safetensors']
supported_pt_extensions += ['.safetensors']
except: except:
print("Could not import safetensors, safetensors support disabled.") print("Could not import safetensors, safetensors support disabled.")
@ -132,7 +134,7 @@ class VAELoader:
vae_dir = os.path.join(models_dir, "vae") vae_dir = os.path.join(models_dir, "vae")
@classmethod @classmethod
def INPUT_TYPES(s): 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",) RETURN_TYPES = ("VAE",)
FUNCTION = "load_vae" FUNCTION = "load_vae"