From 39ee5100be07f021d8a58aa2da8bfc2ff558ccc1 Mon Sep 17 00:00:00 2001 From: Geundol222 Date: Sun, 14 Dec 2025 22:31:08 +0900 Subject: [PATCH] change single selector --- checkpoint_loader.py | 44 +++++++++++++---------------- diffusion_model_loader.py | 32 +++++++++------------ lora_loader.py | 51 +++++++++++++++------------------ tree_utils.py | 23 +++++++++++++-- vae_loader.py | 59 ++++++++++++++++++++++----------------- 5 files changed, 110 insertions(+), 99 deletions(-) diff --git a/checkpoint_loader.py b/checkpoint_loader.py index f785b3d..a5dd232 100644 --- a/checkpoint_loader.py +++ b/checkpoint_loader.py @@ -1,27 +1,21 @@ from typing_extensions import override -from comfy_api.latest import ComfyExtension, io, _io +from comfy_api.latest import ComfyExtension, io import folder_paths -from .tree_utils import attach_tree_metadata, build_tree, list_dirs, list_files, resolve_selected_path, sanitize_rel_dir +from .tree_utils import attach_tree_metadata, build_tree, list_files, resolve_selected_path -def build_folder_file_input(folder_id: str, file_id: str, folder_type: str, tooltip: str | None = None) -> _io.DynamicCombo.Input: - options: list[_io.DynamicCombo.Option] = [] - for rel_dir in list_dirs(folder_type): - file_options = list_files(folder_type, rel_dir) - if not file_options: - continue - child_id = f"{file_id}__{sanitize_rel_dir(rel_dir)}" - inputs = [io.Combo.Input(child_id, options=file_options, tooltip=tooltip)] - options.append(_io.DynamicCombo.Option(rel_dir if rel_dir else "root", inputs)) +def build_file_input(input_id: str, folder_type: str, tooltip: str | None = None) -> io.Combo.Input: + options: list[str] = [] + for rel_path in list_files(folder_type, ""): + options.append(rel_path) if not options: - child_id = f"{file_id}__none" - options.append(_io.DynamicCombo.Option("No files found", [io.Combo.Input(child_id, options=[""], tooltip="No files found")])) - combo = _io.DynamicCombo.Input(folder_id, options=options, tooltip="Select folder") - tree = build_tree(folder_type, file_id) - return attach_tree_metadata(combo, tree, tooltip="Select folder") + options = [""] + combo = io.Combo.Input(input_id, options=sorted(set(options)), tooltip=tooltip or "Select file") + tree = build_tree(folder_type, input_id) + return attach_tree_metadata(combo, tree, tooltip=tooltip or "Select file") class CheckpointLoader(io.ComfyNode): @@ -35,7 +29,7 @@ class CheckpointLoader(io.ComfyNode): category="SK Loader/Advanced", inputs=[ io.Combo.Input("config_name", options=folder_paths.get_filename_list("configs")), - build_folder_file_input("ckpt_folder", "ckpt_name", "checkpoints"), + build_file_input("ckpt", "checkpoints"), ], outputs=[ io.Model.Output(), @@ -45,11 +39,11 @@ class CheckpointLoader(io.ComfyNode): ) @classmethod - def execute(cls, config_name: str, ckpt_folder: dict) -> io.NodeOutput: + def execute(cls, config_name: str, ckpt: dict | str) -> io.NodeOutput: import comfy.sd config_path = folder_paths.get_full_path("configs", config_name) - ckpt_path = resolve_selected_path("checkpoints", ckpt_folder, "ckpt_folder", "ckpt_name") + ckpt_path = resolve_selected_path("checkpoints", ckpt, "ckpt", "ckpt") loaded = comfy.sd.load_checkpoint( config_path, ckpt_path, @@ -70,7 +64,7 @@ class CheckpointLoaderSimple(io.ComfyNode): display_name="[SK] Checkpoint", category="SK Loader", inputs=[ - build_folder_file_input("ckpt_folder", "ckpt_name", "checkpoints", tooltip="The checkpoint (model) to load."), + build_file_input("ckpt", "checkpoints", tooltip="The checkpoint (model) to load."), ], outputs=[ io.Model.Output(), @@ -81,10 +75,10 @@ class CheckpointLoaderSimple(io.ComfyNode): ) @classmethod - def execute(cls, ckpt_folder: dict) -> io.NodeOutput: + def execute(cls, ckpt: dict | str) -> io.NodeOutput: import comfy.sd - ckpt_path = resolve_selected_path("checkpoints", ckpt_folder, "ckpt_folder", "ckpt_name") + ckpt_path = resolve_selected_path("checkpoints", ckpt, "ckpt", "ckpt") out = comfy.sd.load_checkpoint_guess_config( ckpt_path, output_vae=True, @@ -104,7 +98,7 @@ class unCLIPCheckpointLoader(io.ComfyNode): display_name="[SK] unCLIP Checkpoint", category="SK Loader", inputs=[ - build_folder_file_input("ckpt_folder", "ckpt_name", "checkpoints"), + build_file_input("ckpt", "checkpoints"), ], outputs=[ io.Model.Output(), @@ -115,10 +109,10 @@ class unCLIPCheckpointLoader(io.ComfyNode): ) @classmethod - def execute(cls, ckpt_folder: dict) -> io.NodeOutput: + def execute(cls, ckpt: dict | str) -> io.NodeOutput: import comfy.sd - ckpt_path = resolve_selected_path("checkpoints", ckpt_folder, "ckpt_folder", "ckpt_name") + ckpt_path = resolve_selected_path("checkpoints", ckpt, "ckpt", "ckpt") out = comfy.sd.load_checkpoint_guess_config( ckpt_path, output_vae=True, diff --git a/diffusion_model_loader.py b/diffusion_model_loader.py index 9701ebf..4696f47 100644 --- a/diffusion_model_loader.py +++ b/diffusion_model_loader.py @@ -1,27 +1,21 @@ from typing_extensions import override -from comfy_api.latest import ComfyExtension, io, _io +from comfy_api.latest import ComfyExtension, io import folder_paths -from .tree_utils import attach_tree_metadata, build_tree, list_dirs, list_files, resolve_selected_path, sanitize_rel_dir +from .tree_utils import attach_tree_metadata, build_tree, list_files, resolve_selected_path -def build_folder_file_input(folder_id: str, file_id: str, folder_type: str, tooltip: str | None = None) -> _io.DynamicCombo.Input: - options: list[_io.DynamicCombo.Option] = [] - for rel_dir in list_dirs(folder_type): - file_options = list_files(folder_type, rel_dir) - if not file_options: - continue - child_id = f"{file_id}__{sanitize_rel_dir(rel_dir)}" - inputs = [io.Combo.Input(child_id, options=file_options, tooltip=tooltip)] - options.append(_io.DynamicCombo.Option(rel_dir if rel_dir else "root", inputs)) +def build_file_input(input_id: str, folder_type: str, tooltip: str | None = None) -> io.Combo.Input: + options: list[str] = [] + for rel_path in list_files(folder_type, ""): + options.append(rel_path) if not options: - child_id = f"{file_id}__none" - options.append(_io.DynamicCombo.Option("No files found", [io.Combo.Input(child_id, options=[""], tooltip="No files found")])) - combo = _io.DynamicCombo.Input(folder_id, options=options, tooltip="Select folder") - tree = build_tree(folder_type, file_id) - return attach_tree_metadata(combo, tree, tooltip="Select folder") + options = [""] + combo = io.Combo.Input(input_id, options=sorted(set(options)), tooltip=tooltip or "Select file") + tree = build_tree(folder_type, input_id) + return attach_tree_metadata(combo, tree, tooltip=tooltip or "Select file") class UNETLoader(io.ComfyNode): @@ -34,7 +28,7 @@ class UNETLoader(io.ComfyNode): display_name="[SK] Diffusion Model", category="SK Loader/Advanced", inputs=[ - build_folder_file_input("unet_folder", "unet_name", "diffusion_models"), + build_file_input("unet", "diffusion_models"), io.Combo.Input( "weight_dtype", options=["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"], @@ -46,7 +40,7 @@ class UNETLoader(io.ComfyNode): ) @classmethod - def execute(cls, unet_folder: dict, weight_dtype: str) -> io.NodeOutput: + def execute(cls, unet: dict | str, weight_dtype: str) -> io.NodeOutput: import torch import comfy.sd @@ -59,7 +53,7 @@ class UNETLoader(io.ComfyNode): elif weight_dtype == "fp8_e5m2": model_options["dtype"] = torch.float8_e5m2 - unet_path = resolve_selected_path("diffusion_models", unet_folder, "unet_folder", "unet_name") + unet_path = resolve_selected_path("diffusion_models", unet, "unet", "unet") model = comfy.sd.load_diffusion_model(unet_path, model_options=model_options) return io.NodeOutput(model) diff --git a/lora_loader.py b/lora_loader.py index 05263e2..32c967e 100644 --- a/lora_loader.py +++ b/lora_loader.py @@ -1,25 +1,20 @@ from typing_extensions import override -from comfy_api.latest import ComfyExtension, io, _io +from comfy_api.latest import ComfyExtension, io -from .tree_utils import attach_tree_metadata, build_tree, list_dirs, list_files, resolve_selected_path, sanitize_rel_dir +from .tree_utils import attach_tree_metadata, build_tree, list_files, resolve_selected_path -def build_folder_file_input(folder_id: str, file_id: str, folder_type: str, tooltip: str | None = None) -> _io.DynamicCombo.Input: - options: list[_io.DynamicCombo.Option] = [] - for rel_dir in list_dirs(folder_type): - file_options = list_files(folder_type, rel_dir) - if not file_options: - continue - child_id = f"{file_id}__{sanitize_rel_dir(rel_dir)}" - inputs = [io.Combo.Input(child_id, options=file_options, tooltip=tooltip)] - options.append(_io.DynamicCombo.Option(rel_dir if rel_dir else "root", inputs)) +def build_file_input(input_id: str, folder_type: str, tooltip: str | None = None) -> io.Combo.Input: + """Single combo with tree metadata over all files.""" + options: list[str] = [] + for rel_path in list_files(folder_type, ""): + options.append(rel_path) if not options: - child_id = f"{file_id}__none" - options.append(_io.DynamicCombo.Option("No files found", [io.Combo.Input(child_id, options=[""], tooltip="No files found")])) - combo = _io.DynamicCombo.Input(folder_id, options=options, tooltip="Select folder") - tree = build_tree(folder_type, file_id) - return attach_tree_metadata(combo, tree, tooltip="Select folder") + options = [""] + combo = io.Combo.Input(input_id, options=sorted(set(options)), tooltip=tooltip or "Select file") + tree = build_tree(folder_type, input_id) + return attach_tree_metadata(combo, tree, tooltip=tooltip or "Select file") def build_lora_slot_inputs(idx: int) -> list: @@ -28,7 +23,7 @@ def build_lora_slot_inputs(idx: int) -> list: label = f"LoRA #{idx}" return [ io.Boolean.Input(f"{prefix}_enabled", default=False, tooltip=f"Enable {label}"), - build_folder_file_input(f"{prefix}_folder", f"{prefix}_name", "loras", tooltip=f"Select {label}"), + build_file_input(prefix, "loras", tooltip=f"Select {label}"), io.Float.Input( f"{prefix}_strength_model", default=1.0, @@ -60,7 +55,7 @@ class LoraLoader(io.ComfyNode): inputs=[ io.Model.Input("model", tooltip="The diffusion model the LoRA will be applied to."), io.Clip.Input("clip", tooltip="The CLIP model the LoRA will be applied to."), - build_folder_file_input("lora_folder", "lora_name", "loras", tooltip="The name of the LoRA."), + build_file_input("lora", "loras", tooltip="The name of the LoRA."), io.Float.Input( "strength_model", default=1.0, @@ -86,20 +81,20 @@ class LoraLoader(io.ComfyNode): ) @classmethod - def _apply_lora(cls, model, clip, selection: dict, strength_model: float, strength_clip: float): + def _apply_lora(cls, model, clip, selection: dict | str, strength_model: float, strength_clip: float): import comfy.sd import comfy.utils if strength_model == 0 and strength_clip == 0: return model, clip - lora_path = resolve_selected_path("loras", selection, "lora_folder", "lora_name") + lora_path = resolve_selected_path("loras", selection, "lora", "lora") loaded = comfy.utils.load_torch_file(lora_path, safe_load=True) return comfy.sd.load_lora_for_models(model, clip, loaded, strength_model, strength_clip) @classmethod - def execute(cls, model, clip, lora_folder: dict, strength_model: float, strength_clip: float) -> io.NodeOutput: - model_lora, clip_lora = cls._apply_lora(model, clip, lora_folder, strength_model, strength_clip) + def execute(cls, model, clip, lora: dict | str, strength_model: float, strength_clip: float) -> io.NodeOutput: + model_lora, clip_lora = cls._apply_lora(model, clip, lora, strength_model, strength_clip) return io.NodeOutput(model_lora, clip_lora) @@ -114,7 +109,7 @@ class LoraLoaderModelOnly(LoraLoader): category="SK Loader", inputs=[ io.Model.Input("model"), - build_folder_file_input("lora_folder", "lora_name", "loras"), + build_file_input("lora", "loras"), io.Float.Input( "strength_model", default=1.0, @@ -129,8 +124,8 @@ class LoraLoaderModelOnly(LoraLoader): ) @classmethod - def execute(cls, model, lora_folder: dict, strength_model: float) -> io.NodeOutput: - model_lora, _ = cls._apply_lora(model, None, lora_folder, strength_model, 0) + def execute(cls, model, lora: dict | str, strength_model: float) -> io.NodeOutput: + model_lora, _ = cls._apply_lora(model, None, lora, strength_model, 0) return io.NodeOutput(model_lora) @@ -170,8 +165,8 @@ class PowerLoraLoader(io.ComfyNode): if not enabled: continue - selection = kwargs.get(f"lora_{idx}_folder") - if not isinstance(selection, dict): + selection = kwargs.get(f"lora_{idx}") + if not isinstance(selection, (dict, str)): continue strength_model = float(kwargs.get(f"lora_{idx}_strength_model", 1.0)) @@ -180,7 +175,7 @@ class PowerLoraLoader(io.ComfyNode): continue try: - lora_path = resolve_selected_path("loras", selection, f"lora_{idx}_folder", f"lora_{idx}_name") + lora_path = resolve_selected_path("loras", selection, f"lora_{idx}", f"lora_{idx}") except FileNotFoundError: continue diff --git a/tree_utils.py b/tree_utils.py index 343ecce..6af92e8 100644 --- a/tree_utils.py +++ b/tree_utils.py @@ -105,11 +105,30 @@ def build_tree(folder_type: str, file_id: str, extra_roots: list[dict[str, Any]] return tree -def resolve_selected_path(folder_type: str, selection: dict, folder_id: str, file_id: str) -> str: - """Resolve a selection map (combo or tree UI) to an absolute path under the given folder type.""" +def resolve_selected_path(folder_type: str, selection: dict | str, folder_id: str | None = None, file_id: str | None = None) -> str: + """Resolve a selection (string or legacy dict) to an absolute path under the given folder type.""" + # Simple string selection: try absolute, then relative to the folder_type roots. + if isinstance(selection, str): + rel = selection.strip() + if rel in ("", ""): + raise FileNotFoundError("No file selected") + rel = rel.replace("\\", "/") + if os.path.isabs(rel) and os.path.exists(rel): + return rel + rel_no_prefix = rel[len(folder_type) + 1 :] if rel.startswith(f"{folder_type}/") else rel + for candidate_rel in (rel_no_prefix, rel): + for base in folder_paths.get_folder_paths(folder_type): + candidate = os.path.join(base, candidate_rel) + if os.path.exists(candidate): + return candidate + return rel # last resort + if not isinstance(selection, dict): raise ValueError("Invalid selection") + folder_id = folder_id or "" + file_id = file_id or "" + def _valid(val: Any) -> bool: return isinstance(val, str) and val not in ("", "") diff --git a/vae_loader.py b/vae_loader.py index c6f73a2..fad0d03 100644 --- a/vae_loader.py +++ b/vae_loader.py @@ -3,7 +3,7 @@ from typing import Any from typing_extensions import override -from comfy_api.latest import ComfyExtension, io, _io +from comfy_api.latest import ComfyExtension, io import folder_paths from .tree_utils import ALLOWED_EXT, attach_tree_metadata, list_dirs, list_files, sanitize_rel_dir, _ensure_branch @@ -66,34 +66,43 @@ def build_vae_tree(file_id: str) -> list[dict[str, Any]]: return tree -def build_vae_input(folder_id: str, file_id: str) -> _io.DynamicCombo.Input: - options: list[_io.DynamicCombo.Option] = [] +def build_vae_input(input_id: str) -> io.Combo.Input: + """Single select with tree metadata; options = builtins + all vae/vae_approx files.""" + options: list[str] = [] # built-in options - options.append( - _io.DynamicCombo.Option( - "builtins", - [io.Combo.Input(f"{file_id}__builtins", options=BUILTIN_VAES, tooltip="Built-in VAE/TAES")], - ) - ) + options.extend(BUILTIN_VAES) for folder_type in ("vae", "vae_approx"): - for rel_dir in list_dirs(folder_type): - file_options = list_files(folder_type, rel_dir) - if not file_options: - continue - label = f"{folder_type}/{rel_dir}" if rel_dir else folder_type - child_id = f"{file_id}__{sanitize_rel_dir(label)}" - inputs = [io.Combo.Input(child_id, options=file_options, tooltip="Select VAE file")] - options.append(_io.DynamicCombo.Option(label, inputs)) + for rel_path in list_files(folder_type, ""): + options.append(f"{folder_type}/{rel_path}") + if not options: - options.append(_io.DynamicCombo.Option("No files found", [io.Combo.Input(f"{file_id}__none", options=[""], tooltip="No files found")])) - combo = _io.DynamicCombo.Input(folder_id, options=options, tooltip="Select folder") - tree = build_vae_tree(file_id) - return attach_tree_metadata(combo, tree, tooltip="Select folder") + options = [""] + + combo = io.Combo.Input(input_id, options=sorted(set(options)), tooltip="Select VAE") + tree = build_vae_tree(input_id) + return attach_tree_metadata(combo, tree, tooltip="Select VAE") -def resolve_selected_path(selection: dict, folder_id: str, file_id: str) -> str: +def resolve_selected_path(selection: dict | str, folder_id: str = "vae_folder", file_id: str = "vae_name") -> str: + """Resolve tree-aware single select or legacy two-combo selection to a usable path or builtin.""" + if isinstance(selection, str): + rel = selection.strip() + if rel in BUILTIN_VAES: + return rel + rel = rel.replace("\\", "/") + if os.path.isabs(rel) and os.path.exists(rel): + return rel + for folder_type in ("vae", "vae_approx"): + rel_no_prefix = rel[len(folder_type) + 1 :] if rel.startswith(f"{folder_type}/") else rel + for candidate_rel in (rel_no_prefix, rel): + for base in folder_paths.get_folder_paths(folder_type): + candidate = os.path.join(base, candidate_rel) + if os.path.exists(candidate): + return candidate + return rel # last resort + if not isinstance(selection, dict): raise ValueError("Invalid selection") folder_sel = selection.get(folder_id, "builtins") @@ -182,7 +191,7 @@ class VAELoader(io.ComfyNode): display_name="[SK] VAE", category="SK Loader", inputs=[ - build_vae_input("vae_folder", "vae_name"), + build_vae_input("vae"), ], outputs=[ io.Vae.Output(), @@ -191,12 +200,12 @@ class VAELoader(io.ComfyNode): # TODO: scale factor? @classmethod - def execute(cls, vae_folder: dict) -> io.NodeOutput: + def execute(cls, vae: dict | str) -> io.NodeOutput: import torch import comfy.sd import comfy.utils - resolved = resolve_selected_path(vae_folder, "vae_folder", "vae_name") + resolved = resolve_selected_path(vae, "vae_folder", "vae_name") # Check if it's a builtin VAE if resolved in ("pixel_space", "taesd", "taesdxl", "taesd3", "taef1"):