390 lines
13 KiB
Python
390 lines
13 KiB
Python
import torch
|
|
|
|
|
|
ALL_CHECKPOINT_FOLDERS = "All Checkpoints"
|
|
ALL_DIFFUSION_MODEL_FOLDERS = "All Diffusion Models"
|
|
ALL_TEXT_ENCODER_FOLDERS = "All Text Encoders"
|
|
|
|
DIFFUSION_MODEL_WEIGHT_DTYPES = [
|
|
"default",
|
|
"fp8_e4m3fn",
|
|
"fp8_e4m3fn_fast",
|
|
"fp8_e5m2",
|
|
]
|
|
|
|
TEXT_ENCODER_TYPES = [
|
|
"stable_diffusion",
|
|
"stable_cascade",
|
|
"sd3",
|
|
"stable_audio",
|
|
"mochi",
|
|
"ltxv",
|
|
"pixart",
|
|
"cosmos",
|
|
"lumina2",
|
|
"wan",
|
|
"hidream",
|
|
"chroma",
|
|
"ace",
|
|
"omnigen2",
|
|
"qwen_image",
|
|
"hunyuan_image",
|
|
"flux2",
|
|
"ovis",
|
|
"longcat_image",
|
|
"cogvideox",
|
|
"lens",
|
|
"pixeldit",
|
|
"ideogram4",
|
|
"boogu",
|
|
"krea2",
|
|
"joyimage",
|
|
"mage",
|
|
"minimax",
|
|
]
|
|
|
|
|
|
def _normalize_path(value):
|
|
return str(value or "").strip().replace("\\", "/").strip("/")
|
|
|
|
|
|
def _folder_prefixes(file_names, all_folders):
|
|
prefixes = set()
|
|
for file_name in file_names:
|
|
parts = _normalize_path(file_name).split("/")[:-1]
|
|
prefixes.update("/".join(parts[:index]) for index in range(1, len(parts) + 1))
|
|
return [all_folders, *sorted(prefixes, key=str.casefold)]
|
|
|
|
|
|
def _is_in_folder(file_name, folder, all_folders):
|
|
if folder == all_folders:
|
|
return True
|
|
normalized_name = _normalize_path(file_name).casefold()
|
|
normalized_folder = _normalize_path(folder).casefold()
|
|
return bool(normalized_folder) and normalized_name.startswith(normalized_folder + "/")
|
|
|
|
|
|
def _display_name(file_name, folder, all_folders):
|
|
normalized_name = _normalize_path(file_name)
|
|
if folder == all_folders:
|
|
return normalized_name
|
|
normalized_folder = _normalize_path(folder)
|
|
return normalized_name[len(normalized_folder) + 1 :]
|
|
|
|
|
|
def _parse_allow_list(allow_list):
|
|
return {
|
|
_normalize_path(line).casefold()
|
|
for line in str(allow_list or "").splitlines()
|
|
if _normalize_path(line)
|
|
}
|
|
|
|
|
|
class _EnviralFilteredModelLoaderBase:
|
|
FOLDER_CATEGORY = None
|
|
ALL_FOLDERS = None
|
|
FILE_INPUT_NAME = None
|
|
|
|
@classmethod
|
|
def _file_names(cls):
|
|
import folder_paths
|
|
|
|
return folder_paths.get_filename_list(cls.FOLDER_CATEGORY)
|
|
|
|
@classmethod
|
|
def _folder_input(cls):
|
|
return (
|
|
"STRING,COMBO",
|
|
{
|
|
"default": cls.ALL_FOLDERS,
|
|
"widgetType": "COMBO",
|
|
"options": _folder_prefixes(cls._file_names(), cls.ALL_FOLDERS),
|
|
"tooltip": "Limits the model dropdown to this folder and its subfolders.",
|
|
},
|
|
)
|
|
|
|
@classmethod
|
|
def _file_name_input(cls):
|
|
file_names = cls._file_names()
|
|
default = file_names[0] if file_names else ""
|
|
return (
|
|
"STRING,COMBO",
|
|
{
|
|
"default": default,
|
|
"widgetType": "COMBO",
|
|
"options": file_names,
|
|
"tooltip": (
|
|
"Model name. Also accepts a linked STRING, as long as it "
|
|
"matches a model path known to ComfyUI."
|
|
),
|
|
},
|
|
)
|
|
|
|
@classmethod
|
|
def _resolve_folder(cls, folder):
|
|
folder = str(folder or "").strip()
|
|
folders = _folder_prefixes(cls._file_names(), cls.ALL_FOLDERS)
|
|
if folder in folders:
|
|
return folder
|
|
|
|
normalized_folder = _normalize_path(folder).casefold()
|
|
matches = [
|
|
candidate
|
|
for candidate in folders
|
|
if _normalize_path(candidate).casefold() == normalized_folder
|
|
]
|
|
if len(matches) == 1:
|
|
return matches[0]
|
|
raise ValueError(
|
|
f"folder value {folder!r} must be one of the current {cls.FOLDER_CATEGORY} folders"
|
|
)
|
|
|
|
@classmethod
|
|
def _resolve_file_path(cls, file_name):
|
|
import folder_paths
|
|
|
|
file_name = str(file_name or "").strip()
|
|
if not file_name:
|
|
raise ValueError(f"{cls.FILE_INPUT_NAME} must not be empty")
|
|
return folder_paths.get_full_path_or_raise(cls.FOLDER_CATEGORY, file_name)
|
|
|
|
@classmethod
|
|
def _available_files(cls, folder, allow_list=""):
|
|
allowed = _parse_allow_list(allow_list)
|
|
filtered = []
|
|
for file_name in cls._file_names():
|
|
if not _is_in_folder(file_name, folder, cls.ALL_FOLDERS):
|
|
continue
|
|
display_name = _display_name(file_name, folder, cls.ALL_FOLDERS)
|
|
if allowed and not {
|
|
_normalize_path(file_name).casefold(),
|
|
_normalize_path(display_name).casefold(),
|
|
} & allowed:
|
|
continue
|
|
filtered.append((file_name, display_name))
|
|
return filtered
|
|
|
|
@classmethod
|
|
def _validate_selection(cls, folder, file_name, allow_list=""):
|
|
folder = cls._resolve_folder(folder)
|
|
cls._resolve_file_path(file_name)
|
|
if not _is_in_folder(file_name, folder, cls.ALL_FOLDERS):
|
|
raise ValueError(f"{cls.FILE_INPUT_NAME} {file_name!r} is not inside folder {folder!r}")
|
|
available_files = {file_name for file_name, _ in cls._available_files(folder, allow_list)}
|
|
if file_name not in available_files:
|
|
raise ValueError(
|
|
f"{cls.FILE_INPUT_NAME} {file_name!r} is not included by the current allow-list"
|
|
)
|
|
return folder
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, folder, allow_list="", **kwargs):
|
|
try:
|
|
cls._validate_selection(folder, kwargs.get(cls.FILE_INPUT_NAME), allow_list)
|
|
except Exception as err:
|
|
return str(err)
|
|
return True
|
|
|
|
@classmethod
|
|
def validate_inputs(cls, folder, **kwargs):
|
|
return cls.VALIDATE_INPUTS(folder, **kwargs)
|
|
|
|
@classmethod
|
|
def _visible_file_list(cls, folder, allow_list):
|
|
return "\n".join(
|
|
display_name for _, display_name in cls._available_files(folder, allow_list)
|
|
)
|
|
|
|
|
|
class EnviralLoadCheckpointFiltered(_EnviralFilteredModelLoaderBase):
|
|
FOLDER_CATEGORY = "checkpoints"
|
|
ALL_FOLDERS = ALL_CHECKPOINT_FOLDERS
|
|
FILE_INPUT_NAME = "ckpt_name"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"folder": cls._folder_input(),
|
|
"ckpt_name": cls._file_name_input(),
|
|
},
|
|
"optional": {
|
|
"allow_list": (
|
|
"STRING",
|
|
{
|
|
"forceInput": True,
|
|
"multiline": True,
|
|
"tooltip": (
|
|
"Optional newline-separated checkpoint allow-list. Entries may "
|
|
"use displayed names or full paths."
|
|
),
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "CLIP", "VAE", "STRING")
|
|
RETURN_NAMES = ("model", "clip", "vae", "checkpoint_list")
|
|
OUTPUT_TOOLTIPS = (
|
|
"The diffusion model loaded from the checkpoint.",
|
|
"The text encoder loaded from the checkpoint.",
|
|
"The VAE loaded from the checkpoint.",
|
|
"The visible checkpoint options, one display name per line.",
|
|
)
|
|
FUNCTION = "load_checkpoint"
|
|
CATEGORY = "EnviralDesign/model"
|
|
DESCRIPTION = "Checkpoint loader with folder and text allow-list filters."
|
|
SEARCH_ALIASES = ["checkpoint", "load checkpoint", "filtered checkpoint", "model loader"]
|
|
|
|
def load_checkpoint(self, folder, ckpt_name, allow_list=""):
|
|
folder = self._validate_selection(folder, ckpt_name, allow_list)
|
|
ckpt_path = self._resolve_file_path(ckpt_name)
|
|
|
|
import folder_paths
|
|
from comfy import sd
|
|
|
|
model, clip, vae, *_ = sd.load_checkpoint_guess_config(
|
|
ckpt_path,
|
|
output_vae=True,
|
|
output_clip=True,
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
|
)
|
|
return model, clip, vae, self._visible_file_list(folder, allow_list)
|
|
|
|
|
|
class EnviralLoadDiffusionModelFiltered(_EnviralFilteredModelLoaderBase):
|
|
FOLDER_CATEGORY = "diffusion_models"
|
|
ALL_FOLDERS = ALL_DIFFUSION_MODEL_FOLDERS
|
|
FILE_INPUT_NAME = "unet_name"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"folder": cls._folder_input(),
|
|
"unet_name": cls._file_name_input(),
|
|
"weight_dtype": (
|
|
DIFFUSION_MODEL_WEIGHT_DTYPES,
|
|
{"default": "default", "advanced": True},
|
|
),
|
|
},
|
|
"optional": {
|
|
"allow_list": (
|
|
"STRING",
|
|
{
|
|
"forceInput": True,
|
|
"multiline": True,
|
|
"tooltip": (
|
|
"Optional newline-separated diffusion-model allow-list. Entries "
|
|
"may use displayed names or full paths."
|
|
),
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "STRING")
|
|
RETURN_NAMES = ("model", "diffusion_model_list")
|
|
OUTPUT_TOOLTIPS = (
|
|
"The loaded diffusion model.",
|
|
"The visible diffusion-model options, one display name per line.",
|
|
)
|
|
FUNCTION = "load_diffusion_model"
|
|
CATEGORY = "EnviralDesign/model"
|
|
DESCRIPTION = "Diffusion-model loader with folder and text allow-list filters."
|
|
SEARCH_ALIASES = ["diffusion model", "load unet", "filtered unet", "unet loader"]
|
|
|
|
def load_diffusion_model(self, folder, unet_name, weight_dtype, allow_list=""):
|
|
folder = self._validate_selection(folder, unet_name, allow_list)
|
|
model_options = {}
|
|
|
|
if weight_dtype == "fp8_e4m3fn":
|
|
model_options["dtype"] = torch.float8_e4m3fn
|
|
elif weight_dtype == "fp8_e4m3fn_fast":
|
|
model_options["dtype"] = torch.float8_e4m3fn
|
|
model_options["fp8_optimizations"] = True
|
|
elif weight_dtype == "fp8_e5m2":
|
|
model_options["dtype"] = torch.float8_e5m2
|
|
|
|
from comfy import sd
|
|
|
|
model = sd.load_diffusion_model(
|
|
self._resolve_file_path(unet_name),
|
|
model_options=model_options,
|
|
)
|
|
return model, self._visible_file_list(folder, allow_list)
|
|
|
|
|
|
class EnviralLoadTextEncoderFiltered(_EnviralFilteredModelLoaderBase):
|
|
FOLDER_CATEGORY = "text_encoders"
|
|
ALL_FOLDERS = ALL_TEXT_ENCODER_FOLDERS
|
|
FILE_INPUT_NAME = "clip_name"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"folder": cls._folder_input(),
|
|
"clip_name": cls._file_name_input(),
|
|
"type": (TEXT_ENCODER_TYPES, {"default": "stable_diffusion"}),
|
|
},
|
|
"optional": {
|
|
"device": (["default", "cpu"], {"advanced": True}),
|
|
"allow_list": (
|
|
"STRING",
|
|
{
|
|
"forceInput": True,
|
|
"multiline": True,
|
|
"tooltip": (
|
|
"Optional newline-separated text-encoder allow-list. Entries may "
|
|
"use displayed names or full paths."
|
|
),
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("CLIP", "STRING")
|
|
RETURN_NAMES = ("clip", "text_encoder_list")
|
|
OUTPUT_TOOLTIPS = (
|
|
"The loaded text encoder.",
|
|
"The visible text-encoder options, one display name per line.",
|
|
)
|
|
FUNCTION = "load_text_encoder"
|
|
CATEGORY = "EnviralDesign/model"
|
|
DESCRIPTION = "Text-encoder loader with folder and text allow-list filters."
|
|
SEARCH_ALIASES = ["text encoder", "load clip", "filtered clip", "clip loader"]
|
|
|
|
def load_text_encoder(self, folder, clip_name, type="stable_diffusion", device="default", allow_list=""):
|
|
folder = self._validate_selection(folder, clip_name, allow_list)
|
|
model_options = {}
|
|
|
|
from comfy import sd
|
|
|
|
clip_type = getattr(sd.CLIPType, type.upper(), sd.CLIPType.STABLE_DIFFUSION)
|
|
if device == "cpu":
|
|
model_options["load_device"] = model_options["offload_device"] = torch.device("cpu")
|
|
|
|
import folder_paths
|
|
|
|
clip = sd.load_clip(
|
|
ckpt_paths=[self._resolve_file_path(clip_name)],
|
|
embedding_directory=folder_paths.get_folder_paths("embeddings"),
|
|
clip_type=clip_type,
|
|
model_options=model_options,
|
|
)
|
|
return clip, self._visible_file_list(folder, allow_list)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"EnviralLoadCheckpointFiltered": EnviralLoadCheckpointFiltered,
|
|
"EnviralLoadDiffusionModelFiltered": EnviralLoadDiffusionModelFiltered,
|
|
"EnviralLoadTextEncoderFiltered": EnviralLoadTextEncoderFiltered,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"EnviralLoadCheckpointFiltered": "Enviral Load Checkpoint",
|
|
"EnviralLoadDiffusionModelFiltered": "Enviral Load Diffusion Model",
|
|
"EnviralLoadTextEncoderFiltered": "Enviral Load Text Encoder",
|
|
}
|