Commit 2706c0b7 authored by comfyanonymous's avatar comfyanonymous
Browse files

Some VAEs come in .pt files.

parent d133cf4f
...@@ -17,9 +17,11 @@ import comfy.samplers ...@@ -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: ...@@ -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"
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment