import logging import os logger = logging.getLogger(__name__) class NodePackageMapper: """Maps node class_type names to their source package.""" def __init__(self): self._map = None def _build_map(self): import nodes self._map = {} for class_type, node_cls in nodes.NODE_CLASS_MAPPINGS.items(): module = getattr(node_cls, "RELATIVE_PYTHON_MODULE", None) if module: # "custom_nodes.PackageName" -> "PackageName" # "comfy_extras.nodes_xyz" -> "__builtin__" # "comfy_api_nodes.xyz" -> "__builtin__" parts = module.split(".") if parts[0] == "custom_nodes" and len(parts) > 1: self._map[class_type] = parts[1] else: self._map[class_type] = "__builtin__" else: self._map[class_type] = "__builtin__" @property def mapping(self): if self._map is None: self._build_map() return self._map def get_package(self, class_type): return self.mapping.get(class_type, "__unknown__") def get_all_packages(self): """Return set of all known package names, including zero-node packages.""" packages = set(self.mapping.values()) try: import nodes import folder_paths # Get all custom_nodes directories to filter LOADED_MODULE_DIRS custom_node_dirs = set() for d in folder_paths.get_folder_paths("custom_nodes"): custom_node_dirs.add(os.path.normpath(d)) # LOADED_MODULE_DIRS contains ALL modules (custom nodes, comfy_extras, # comfy_api_nodes). We only want custom node packages, identified by # their directory being directly inside a custom_nodes directory. for module_name, module_dir in nodes.LOADED_MODULE_DIRS.items(): parent_dir = os.path.normpath(os.path.dirname(module_dir)) if parent_dir in custom_node_dirs: packages.add(os.path.basename(module_dir)) except Exception: logger.warning("Could not read LOADED_MODULE_DIRS", exc_info=True) packages.discard("__builtin__") return packages def invalidate(self): """Force rebuild on next access (e.g. after node reload).""" self._map = None # Folder types that are not model files and should not be tracked EXCLUDED_FOLDER_TYPES = { "loras", "configs", "custom_nodes", "temp", "output", "input", "annotators", "assets", } class ModelMapper: """Tracks which folder_paths model types exist and resolves filenames to types.""" def __init__(self): self._folder_files = None # {folder_type: frozenset(filenames)} self._reverse = None # {filename: folder_type} def _build(self): try: import folder_paths folder_files = {} for folder_type in folder_paths.folder_names_and_paths: if folder_type in EXCLUDED_FOLDER_TYPES: continue try: files = folder_paths.get_filename_list(folder_type) except Exception: files = [] if files: folder_files[folder_type] = frozenset(files) # Build reverse map: filename -> folder_type (last write wins on collision) reverse = {} for folder_type, files in folder_files.items(): for f in files: reverse[f] = folder_type # Assign atomically only on success self._folder_files = folder_files self._reverse = reverse except Exception: logger.warning("ModelMapper: failed to build model map", exc_info=True) self._folder_files = {} self._reverse = {} def _ensure(self): if self._folder_files is None: self._build() def get_model_type(self, filename): """Return the folder type for a filename, or None if not tracked.""" self._ensure() return self._reverse.get(filename) def get_all_models(self): """Return {folder_type: [filename, ...]} for all tracked types.""" self._ensure() return {k: sorted(v) for k, v in self._folder_files.items()} def delete_models(self, models): """Delete verified model files and return per-model results. Each requested file must still be in the mapper's known model list and resolve beneath the configured root for its model type. This keeps the HTTP endpoint from being usable to remove arbitrary files. """ self._ensure() try: import folder_paths except ImportError: return [{"status": "error", "message": "folder_paths unavailable"}] results = [] seen = set() for item in models: if not isinstance(item, dict): results.append({"status": "error", "message": "invalid model entry"}) continue model_type = item.get("model_type") model_name = item.get("model_name") if not isinstance(model_type, str) or not isinstance(model_name, str): results.append({ "model_type": model_type, "model_name": model_name, "status": "error", "message": "invalid model entry", }) continue key = (model_type, model_name) if key in seen: continue seen.add(key) result = {"model_type": model_type, "model_name": model_name} known = self._folder_files.get(model_type, frozenset()) if model_name not in known: result.update(status="error", message="model is not installed") results.append(result) continue try: path = folder_paths.get_full_path(model_type, model_name) roots = folder_paths.get_folder_paths(model_type) path = os.path.abspath(path) if path else None allowed = any( os.path.commonpath([path, os.path.abspath(root)]) == os.path.abspath(root) for root in roots ) if path else False except (OSError, TypeError, ValueError): path = None allowed = False if not path or not allowed or not os.path.isfile(path): result.update(status="error", message="model file not found in its configured folder") results.append(result) continue try: os.remove(path) result["status"] = "deleted" except OSError as exc: result.update(status="error", message=str(exc)) results.append(result) # The model-folder listing is now stale whether any deletion succeeded # or not (another process could have changed it), so rebuild next time. self.invalidate() return results def extract_models_from_prompt(self, prompt): """Scan a prompt dict and return (model_name, model_type) pairs. For each node, inspects INPUT_TYPES() to find list-type (folder dropdown) inputs, then resolves the selected value against the folder_paths reverse map. """ self._ensure() try: import nodes as comfy_nodes except ImportError: return [] seen = set() results = [] for node_data in prompt.values(): class_type = node_data.get("class_type") node_inputs = node_data.get("inputs", {}) if not class_type or not node_inputs: continue node_cls = comfy_nodes.NODE_CLASS_MAPPINGS.get(class_type) if node_cls is None: continue try: input_types = node_cls.INPUT_TYPES() for category in ("required", "optional"): for input_name, input_def in input_types.get(category, {}).items(): if not isinstance(input_def, (list, tuple)) or not input_def: continue # ComfyUI folder dropdowns have a list as their type if not isinstance(input_def[0], list): continue value = node_inputs.get(input_name) if not isinstance(value, str) or value in seen: continue model_type = self.get_model_type(value) if model_type: seen.add(value) results.append((value, model_type)) except Exception: continue return results def invalidate(self): """Force rebuild on next access.""" self._folder_files = None self._reverse = None