From 378c3274063c7e5f76ace71fb42134dae3a29c27 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Wed, 8 Apr 2026 15:27:12 +0200 Subject: [PATCH] Add subdirectory selector --- __init__.py | 4 +- nodes/downloader_nodes.py | 231 +++++++++++++++++++++++++++++++++++++- nodes/downloader_utils.py | 9 +- 3 files changed, 235 insertions(+), 9 deletions(-) diff --git a/__init__.py b/__init__.py index a8a2706..f78e99f 100644 --- a/__init__.py +++ b/__init__.py @@ -30,7 +30,7 @@ from .nodes.lora_resize import ( LoRAMergeToModel ) from .nodes.downloader_nodes import ( - CheckpointDownloader, LoRADownloader, EmbeddingDownloader, + CheckpointDownloader, DiffusionModelDownloader, LoRADownloader, EmbeddingDownloader, VAEDownloader, ControlNetDownloader, ManualPathDownloader ) @@ -65,7 +65,7 @@ class ModelUtilsExtension(ComfyExtension): # LoRA Merge To Model LoRAMergeToModel, # Downloaders - CheckpointDownloader, LoRADownloader, EmbeddingDownloader, + CheckpointDownloader, DiffusionModelDownloader, LoRADownloader, EmbeddingDownloader, VAEDownloader, ControlNetDownloader, ManualPathDownloader, ] diff --git a/nodes/downloader_nodes.py b/nodes/downloader_nodes.py index bc13b70..6099656 100644 --- a/nodes/downloader_nodes.py +++ b/nodes/downloader_nodes.py @@ -1,7 +1,31 @@ +import os import folder_paths from comfy_api.latest import io from .downloader_utils import scan_and_process +def get_subdirectories(folder_name: str) -> list[str]: + """Helper to get a list of subdirectories for a given model folder name.""" + paths = folder_paths.get_folder_paths(folder_name) + if not paths: + return ["/"] + + subdirs = set() + for base_path in paths: + if not os.path.exists(base_path): + continue + for root, dirs, _ in os.walk(base_path): + # Ignore hidden directories directly during walk + dirs[:] = [d for d in dirs if not d.startswith('.')] + + rel_path = os.path.relpath(root, base_path) + if rel_path == ".": + subdirs.add("/") + else: + # Store with forward slashes for consistency + subdirs.add("/" + rel_path.replace("\\", "/")) + + return sorted(list(subdirs)) + class BaseDownloaderNode(io.ComfyNode): CATEGORY = "ModelUtils/Downloader" OUTPUT_NODE = True @@ -50,29 +74,226 @@ class BaseDownloaderNode(io.ComfyNode): return [] class CheckpointDownloader(BaseDownloaderNode): + @classmethod + def define_schema(cls): + inputs = list(cls.get_common_inputs().values()) + inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("checkpoints"), default="/")) + return io.Schema( + node_id=cls.__name__, + display_name="Checkpoint Downloader", + category=cls.CATEGORY, + inputs=inputs, + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) + + @classmethod + def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + return super().execute( + recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples, + threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs + ) + @classmethod def get_scan_dirs(cls, **kwargs): - return folder_paths.get_folder_paths("checkpoints") + base_paths = folder_paths.get_folder_paths("checkpoints") + if not base_paths: + return [] + sub = kwargs.get("subdirectory", "/") + if sub == "/": + return base_paths + + target_paths = [] + for bp in base_paths: + target_path = os.path.join(bp, sub.lstrip("/")) + if os.path.exists(target_path): + target_paths.append(target_path) + return target_paths + +class DiffusionModelDownloader(BaseDownloaderNode): + @classmethod + def define_schema(cls): + inputs = list(cls.get_common_inputs().values()) + inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("diffusion_models"), default="/")) + return io.Schema( + node_id=cls.__name__, + display_name="Diffusion Model Downloader", + category=cls.CATEGORY, + inputs=inputs, + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) + + @classmethod + def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + return super().execute( + recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples, + threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs + ) + + @classmethod + def get_scan_dirs(cls, **kwargs): + base_paths = folder_paths.get_folder_paths("diffusion_models") + if not base_paths: + return [] + sub = kwargs.get("subdirectory", "/") + if sub == "/": + return base_paths + + target_paths = [] + for bp in base_paths: + target_path = os.path.join(bp, sub.lstrip("/")) + if os.path.exists(target_path): + target_paths.append(target_path) + return target_paths class LoRADownloader(BaseDownloaderNode): + @classmethod + def define_schema(cls): + inputs = list(cls.get_common_inputs().values()) + inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("loras"), default="/")) + return io.Schema( + node_id=cls.__name__, + display_name="LoRA Downloader", + category=cls.CATEGORY, + inputs=inputs, + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) + + @classmethod + def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + return super().execute( + recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples, + threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs + ) + @classmethod def get_scan_dirs(cls, **kwargs): - return folder_paths.get_folder_paths("loras") + base_paths = folder_paths.get_folder_paths("loras") + if not base_paths: + return [] + sub = kwargs.get("subdirectory", "/") + if sub == "/": + return base_paths + + target_paths = [] + for bp in base_paths: + target_path = os.path.join(bp, sub.lstrip("/")) + if os.path.exists(target_path): + target_paths.append(target_path) + return target_paths class EmbeddingDownloader(BaseDownloaderNode): + @classmethod + def define_schema(cls): + inputs = list(cls.get_common_inputs().values()) + inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("embeddings"), default="/")) + return io.Schema( + node_id=cls.__name__, + display_name="Embedding Downloader", + category=cls.CATEGORY, + inputs=inputs, + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) + + @classmethod + def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + return super().execute( + recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples, + threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs + ) + @classmethod def get_scan_dirs(cls, **kwargs): - return folder_paths.get_folder_paths("embeddings") + base_paths = folder_paths.get_folder_paths("embeddings") + if not base_paths: + return [] + sub = kwargs.get("subdirectory", "/") + if sub == "/": + return base_paths + + target_paths = [] + for bp in base_paths: + target_path = os.path.join(bp, sub.lstrip("/")) + if os.path.exists(target_path): + target_paths.append(target_path) + return target_paths class VAEDownloader(BaseDownloaderNode): + @classmethod + def define_schema(cls): + inputs = list(cls.get_common_inputs().values()) + inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("vae"), default="/")) + return io.Schema( + node_id=cls.__name__, + display_name="VAE Downloader", + category=cls.CATEGORY, + inputs=inputs, + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) + + @classmethod + def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + return super().execute( + recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples, + threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs + ) + @classmethod def get_scan_dirs(cls, **kwargs): - return folder_paths.get_folder_paths("vae") + base_paths = folder_paths.get_folder_paths("vae") + if not base_paths: + return [] + sub = kwargs.get("subdirectory", "/") + if sub == "/": + return base_paths + + target_paths = [] + for bp in base_paths: + target_path = os.path.join(bp, sub.lstrip("/")) + if os.path.exists(target_path): + target_paths.append(target_path) + return target_paths class ControlNetDownloader(BaseDownloaderNode): + @classmethod + def define_schema(cls): + inputs = list(cls.get_common_inputs().values()) + inputs.insert(0, io.Combo.Input("subdirectory", options=get_subdirectories("controlnet"), default="/")) + return io.Schema( + node_id=cls.__name__, + display_name="ControlNet Downloader", + category=cls.CATEGORY, + inputs=inputs, + outputs=[io.String.Output(display_name="status")], + is_output_node=True, + ) + + @classmethod + def execute(cls, subdirectory, recursive, nsfw_level, max_examples, threads, api_key, **kwargs): + return super().execute( + recursive=recursive, nsfw_level=nsfw_level, max_examples=max_examples, + threads=threads, api_key=api_key, subdirectory=subdirectory, **kwargs + ) + @classmethod def get_scan_dirs(cls, **kwargs): - return folder_paths.get_folder_paths("controlnet") + base_paths = folder_paths.get_folder_paths("controlnet") + if not base_paths: + return [] + sub = kwargs.get("subdirectory", "/") + if sub == "/": + return base_paths + + target_paths = [] + for bp in base_paths: + target_path = os.path.join(bp, sub.lstrip("/")) + if os.path.exists(target_path): + target_paths.append(target_path) + return target_paths class ManualPathDownloader(BaseDownloaderNode): @classmethod diff --git a/nodes/downloader_utils.py b/nodes/downloader_utils.py index 3ae40dd..7694525 100644 --- a/nodes/downloader_utils.py +++ b/nodes/downloader_utils.py @@ -548,10 +548,15 @@ def scan_and_process(scan_dirs, recursive=True, db_path=None, nsfw_level="All", path = Path(scan_dir) if recursive: for ext in model_exts: - files_to_process.extend(path.rglob(f"*{ext}")) + # rglob includes all directories, so we must filter out hidden ones + for f in path.rglob(f"*{ext}"): + if not any(part.startswith('.') for part in f.parts if part != '.' and part != '..'): + files_to_process.append(f) else: for ext in model_exts: - files_to_process.extend(path.glob(f"*{ext}")) + for f in path.glob(f"*{ext}"): + if not any(part.startswith('.') for part in f.parts if part != '.' and part != '..'): + files_to_process.append(f) nsfw_threshold = NSFW_LEVELS.get(nsfw_level, NSFW_LEVELS["All"])