Dariusz L
32d437df50
Add support for additional ComfyUI model types including diffusers, gligen, style models (latent/regular upscale), audio encoders, background removal, frame interpolation, geometry estimation, detection, model patches/photomaker, and optical flow models. Backend changes: - Extend CATEGORY_MAP, URN_TYPE_MAP, and node type hints in workflow analyzer to cover dozens of new loader node types and singular/plural aliases. - Add MODEL_OUTPUT_TYPE_TO_CATEGORY and get_node_output_category_hint() so loader output types can contribute category hints. - Make widget name matching normalize underscores, hyphens, and spaces. - Improve directory discovery to enumerate folders from folder_paths while filtering unsupported categories. - Include folders as model entries in directory scanning (with a diffusers guard requiring model_index.json). Frontend changes: - Update download target category maps, display names, token name maps, and per-node widget index lists to match the expanded backend coverage.
257 lines
8.5 KiB
Python
257 lines
8.5 KiB
Python
"""
|
|
Directory Scanner Module
|
|
|
|
Scans configured model directories and finds available model files.
|
|
"""
|
|
|
|
import os
|
|
import time
|
|
from typing import List, Dict, Tuple, Optional
|
|
|
|
from .log_system.log_funcs import (
|
|
log_debug,
|
|
log_info,
|
|
log_warn,
|
|
log_error,
|
|
log_exception,
|
|
)
|
|
|
|
# Import folder_paths lazily - it may not be available until ComfyUI is initialized
|
|
try:
|
|
import folder_paths
|
|
except ImportError:
|
|
folder_paths = None
|
|
log_warn("Model Resolver: folder_paths not available yet - will retry later")
|
|
|
|
# Model file extensions to look for
|
|
# This matches folder_paths.supported_pt_extensions
|
|
MODEL_EXTENSIONS = {
|
|
".ckpt",
|
|
".pt",
|
|
".pt2",
|
|
".bin",
|
|
".pth",
|
|
".safetensors",
|
|
".pkl",
|
|
".sft",
|
|
".onnx",
|
|
".gguf",
|
|
}
|
|
|
|
_MODEL_FILES_CACHE: Optional[List[Dict[str, str]]] = None
|
|
_MODEL_FILES_CACHE_AT: float = 0.0
|
|
_MODEL_FILES_CACHE_TTL_SECONDS = 2.0
|
|
|
|
|
|
def get_model_directories() -> Dict[str, Tuple[List[str], set]]:
|
|
"""
|
|
Get all configured model directories from folder_paths.
|
|
|
|
Returns:
|
|
Dictionary mapping category name to a tuple. ComfyUI may provide either:
|
|
- (paths, extensions), or
|
|
- (paths, extensions, recursive_flag)
|
|
"""
|
|
global folder_paths
|
|
|
|
if folder_paths is None:
|
|
# Try to import again
|
|
try:
|
|
import folder_paths as fp
|
|
|
|
folder_paths = fp
|
|
except ImportError:
|
|
log_error("Model Resolver: folder_paths still not available")
|
|
return {}
|
|
|
|
return folder_paths.folder_names_and_paths.copy()
|
|
|
|
|
|
def scan_directory(
|
|
directory: str, extensions: set, category: str
|
|
) -> List[Dict[str, str]]:
|
|
"""
|
|
Recursively scan a single directory for model files.
|
|
|
|
Args:
|
|
directory: Absolute path to directory to scan
|
|
extensions: Set of file extensions to look for
|
|
category: Model category name (e.g., 'checkpoints', 'loras')
|
|
|
|
Returns:
|
|
List of dictionaries with model information:
|
|
{
|
|
'filename': 'model.safetensors',
|
|
'path': 'absolute/path/to/model.safetensors',
|
|
'relative_path': 'subfolder/model.safetensors' or 'model.safetensors',
|
|
'category': 'checkpoints',
|
|
'base_directory': 'absolute/path/to/base'
|
|
}
|
|
"""
|
|
models = []
|
|
|
|
if not os.path.exists(directory) or not os.path.isdir(directory):
|
|
# log_debug(f"Directory does not exist or is not accessible: {directory}")
|
|
return models
|
|
|
|
try:
|
|
# Get absolute path and normalize
|
|
base_directory = os.path.abspath(directory)
|
|
|
|
# Walk through directory recursively
|
|
for root, dirs, files in os.walk(base_directory, followlinks=True):
|
|
# Skip hidden directories
|
|
dirs[:] = [d for d in dirs if not d.startswith(".")]
|
|
|
|
if "folder" in (extensions or set()) and root != base_directory:
|
|
if category != "diffusers" or "model_index.json" in files:
|
|
try:
|
|
relative_path = os.path.relpath(root, base_directory)
|
|
except ValueError:
|
|
relative_path = os.path.basename(root)
|
|
|
|
models.append(
|
|
{
|
|
"filename": os.path.basename(root),
|
|
"path": root,
|
|
"relative_path": relative_path,
|
|
"category": category,
|
|
"base_directory": base_directory,
|
|
}
|
|
)
|
|
|
|
for filename in files:
|
|
# Check if file has a model extension
|
|
file_ext = os.path.splitext(filename)[1].lower()
|
|
|
|
# For categories with empty extension set, accept all files
|
|
# Otherwise, check if extension matches
|
|
# Accept if:
|
|
# - no explicit extensions configured, or
|
|
# - matches configured extensions, or
|
|
# - matches our known model extensions
|
|
if (
|
|
len(extensions or set()) == 0
|
|
or file_ext in extensions
|
|
or file_ext in MODEL_EXTENSIONS
|
|
):
|
|
full_path = os.path.join(root, filename)
|
|
|
|
# Calculate relative path from base directory
|
|
# IMPORTANT: Use OS-native path separators (backslashes on Windows)
|
|
# This matches ComfyUI's recursive_search format for get_filename_list
|
|
try:
|
|
relative_path = os.path.relpath(full_path, base_directory)
|
|
# DO NOT normalize - keep OS-native separators to match ComfyUI
|
|
# ComfyUI's get_filename_list uses os.path.relpath which returns
|
|
# backslashes on Windows, forward slashes on Unix
|
|
except ValueError:
|
|
# If paths are on different drives (Windows), use filename only
|
|
relative_path = filename
|
|
|
|
models.append(
|
|
{
|
|
"filename": filename,
|
|
"path": full_path,
|
|
"relative_path": relative_path,
|
|
"category": category,
|
|
"base_directory": base_directory,
|
|
}
|
|
)
|
|
except (OSError, PermissionError) as e:
|
|
log_warn(f"Error scanning directory {directory}: {e}")
|
|
|
|
return models
|
|
|
|
|
|
def scan_all_directories() -> List[Dict[str, str]]:
|
|
"""
|
|
Scan all configured model directories and return list of available models.
|
|
|
|
Returns:
|
|
List of dictionaries with model information (same format as scan_directory)
|
|
"""
|
|
all_models = []
|
|
directories = get_model_directories()
|
|
|
|
for category, value in directories.items():
|
|
# Skip categories that aren't typically model directories
|
|
if category in ["custom_nodes", "configs"]:
|
|
continue
|
|
|
|
# Unpack folder_paths value flexibly: (paths, extensions) or (paths, extensions, recursive)
|
|
paths = []
|
|
extensions = set()
|
|
try:
|
|
if isinstance(value, (list, tuple)):
|
|
if len(value) >= 2:
|
|
paths = value[0] or []
|
|
raw_exts = value[1]
|
|
else:
|
|
# Unexpected format; treat value as paths
|
|
paths = list(value)
|
|
raw_exts = []
|
|
elif isinstance(value, dict):
|
|
paths = value.get("paths") or value.get("path") or []
|
|
raw_exts = value.get("extensions") or []
|
|
else:
|
|
# Unknown format; skip category
|
|
log_debug(
|
|
f"Unexpected folder_paths format for category {category}: {type(value)}"
|
|
)
|
|
continue
|
|
|
|
# Normalize extensions to a set[str]
|
|
if isinstance(raw_exts, (list, tuple, set)):
|
|
extensions = {str(e).lower() for e in raw_exts}
|
|
elif raw_exts:
|
|
extensions = {str(raw_exts).lower()}
|
|
except Exception as e:
|
|
log_warn(f"Error interpreting folder_paths entry for {category}: {e}")
|
|
continue
|
|
|
|
for directory_path in paths:
|
|
try:
|
|
models = scan_directory(directory_path, extensions, category)
|
|
all_models.extend(models)
|
|
#log_debug(f"Found {len(models)} models in {category}/{directory_path}")
|
|
except Exception as e:
|
|
log_warn(f"Error scanning {category} directory {directory_path}: {e}")
|
|
|
|
return all_models
|
|
|
|
|
|
def invalidate_model_files_cache() -> None:
|
|
"""Clear the in-memory model file cache."""
|
|
global _MODEL_FILES_CACHE, _MODEL_FILES_CACHE_AT
|
|
_MODEL_FILES_CACHE = None
|
|
_MODEL_FILES_CACHE_AT = 0.0
|
|
|
|
|
|
def get_model_files(force_rescan: bool = False) -> List[Dict[str, str]]:
|
|
"""
|
|
Get list of all available model files with metadata.
|
|
|
|
This is the main entry point for getting model files.
|
|
|
|
Args:
|
|
force_rescan: If True, bypass the short-lived cache and rescan directories
|
|
|
|
Returns:
|
|
List of model dictionaries (same format as scan_directory)
|
|
"""
|
|
global _MODEL_FILES_CACHE, _MODEL_FILES_CACHE_AT
|
|
|
|
now = time.monotonic()
|
|
if (
|
|
not force_rescan
|
|
and _MODEL_FILES_CACHE is not None
|
|
and (now - _MODEL_FILES_CACHE_AT) < _MODEL_FILES_CACHE_TTL_SECONDS
|
|
):
|
|
return _MODEL_FILES_CACHE
|
|
|
|
models = scan_all_directories()
|
|
_MODEL_FILES_CACHE = models
|
|
_MODEL_FILES_CACHE_AT = now
|
|
return models
|