39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
# vae_loader.py
|
|
import os
|
|
import comfy.sd
|
|
import comfy.utils
|
|
from .base_loader import DiffusersLoaderBase
|
|
from .utils import DiffusersUtils
|
|
|
|
class DiffusersVAELoader(DiffusersLoaderBase):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"sub_directory": (DiffusersUtils.get_model_directories(DiffusersUtils.get_base_path()),),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VAE",)
|
|
FUNCTION = "load_vae"
|
|
CATEGORY = "DiffusersLoader"
|
|
|
|
@classmethod
|
|
def load_vae(cls, sub_directory):
|
|
return (cls.load_model(sub_directory),)
|
|
|
|
@classmethod
|
|
def load_model(cls, sub_directory):
|
|
base_path = DiffusersUtils.get_base_path()
|
|
sub_dir_path = os.path.join(base_path, sub_directory)
|
|
|
|
vae_folder = os.path.join(sub_dir_path, "vae")
|
|
vae_path = DiffusersUtils.find_model_file(vae_folder)
|
|
print(f"DiffusersVAELoader: Attempting to load VAE model from: {vae_path}")
|
|
|
|
try:
|
|
vae_sd = comfy.utils.load_torch_file(vae_path)
|
|
return comfy.sd.VAE(sd=vae_sd)
|
|
except Exception as e:
|
|
print(f"DiffusersVAELoader: Error loading VAE model: {e}")
|
|
raise |