40 lines
1.5 KiB
Python
40 lines
1.5 KiB
Python
# combined_diffusers_loader.py
|
|
from .utils import DiffusersUtils
|
|
from .unet_loader import DiffusersUNETLoader
|
|
from .clip_loader import DiffusersClipLoader
|
|
from .vae_loader import DiffusersVAELoader
|
|
|
|
class CombinedDiffusersLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
base_path = DiffusersUtils.get_base_path()
|
|
model_directories = DiffusersUtils.get_model_directories(base_path)
|
|
return {
|
|
"required": {
|
|
"sub_directory": (model_directories,),
|
|
"clip_type": (["stable_diffusion", "stable_cascade"],),
|
|
"clip_parts": (["all", "part_1", "part_2"],),
|
|
"transformer_parts": (["all", "part_1", "part_2", "part_3"],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
|
FUNCTION = "load_models"
|
|
CATEGORY = "DiffusersLoader/Combined"
|
|
|
|
@classmethod
|
|
def load_models(cls, sub_directory, clip_type="stable_diffusion", clip_parts="all", transformer_parts="all"):
|
|
|
|
unet_model = DiffusersUNETLoader.load_model(sub_directory, transformer_parts)
|
|
|
|
|
|
clip_model = DiffusersClipLoader.load_model(sub_directory, clip_type, clip_parts)
|
|
|
|
|
|
vae_model = DiffusersVAELoader.load_model(sub_directory)
|
|
|
|
return unet_model, clip_model, vae_model
|
|
|
|
# Ensure the node is registered properly
|
|
NODE_CLASS_MAPPINGS = {"CombinedDiffusersLoader": CombinedDiffusersLoader}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"CombinedDiffusersLoader": "Combined Diffusers Loader"} |