255 lines
9.0 KiB
Python
255 lines
9.0 KiB
Python
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
|