feat: add ModelMapper with folder_paths introspection
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -66,3 +66,116 @@ class NodePackageMapper:
|
||||
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
|
||||
|
||||
self._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:
|
||||
self._folder_files[folder_type] = frozenset(files)
|
||||
|
||||
# Reverse map: filename -> folder_type (last write wins on collision)
|
||||
self._reverse = {}
|
||||
for folder_type, files in self._folder_files.items():
|
||||
for f in files:
|
||||
self._reverse[f] = folder_type
|
||||
|
||||
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 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()
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
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))
|
||||
|
||||
return results
|
||||
|
||||
def invalidate(self):
|
||||
"""Force rebuild on next access."""
|
||||
self._folder_files = None
|
||||
self._reverse = None
|
||||
|
||||
Reference in New Issue
Block a user