diff --git a/core/matcher.py b/core/matcher.py index b04baa9..1717b0d 100644 --- a/core/matcher.py +++ b/core/matcher.py @@ -283,18 +283,8 @@ def base_model_score(candidate: str, preferred: Optional[str]) -> float: # ==================== CENTRALIZED MATCHING & CONFIDENCE HELPERS ==================== -MODEL_FILE_EXTENSIONS = { - ".ckpt", - ".pt", - ".pt2", - ".bin", - ".pth", - ".safetensors", - ".pkl", - ".sft", - ".onnx", - ".gguf", -} +from .type_utils import MODEL_EXTENSIONS +MODEL_FILE_EXTENSIONS = MODEL_EXTENSIONS def strip_known_model_extension(filename: str) -> str: diff --git a/core/resolver.py b/core/resolver.py index fd63156..66f78b6 100644 --- a/core/resolver.py +++ b/core/resolver.py @@ -27,7 +27,7 @@ from .workflow_analyzer import ( should_scan_as_model_reference, ) from .matcher import find_matches -from .type_utils import as_dict, as_list +from .type_utils import as_dict, as_list, MODEL_EXTENSIONS as _MODEL_EXTENSIONS from .workflow_updater import update_workflow_nodes from .sources.civitai import resolve_urn from .sources.huggingface import parse_huggingface_url as parse_hf_url @@ -36,7 +36,7 @@ from .sources.huggingface import parse_huggingface_url as parse_hf_url URL_PATTERN = re.compile(r'(https?://(?:huggingface\.co|civitai\.com)[^\s"\'<>\)\\]+)') # Model file extensions to look for -MODEL_EXTENSIONS = (".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".onnx", ".gguf") +MODEL_EXTENSIONS = tuple(_MODEL_EXTENSIONS) from .path_utils import ( diff --git a/core/scanner.py b/core/scanner.py index 0e262a4..0014e76 100644 --- a/core/scanner.py +++ b/core/scanner.py @@ -17,6 +17,7 @@ from .log_system.log_funcs import ( ) from .path_utils import get_path_identity +from .type_utils import MODEL_EXTENSIONS # Import folder_paths lazily - it may not be available until ComfyUI is initialized try: @@ -25,21 +26,6 @@ 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 diff --git a/core/type_utils.py b/core/type_utils.py index 06d4c52..6e060c4 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -42,3 +42,17 @@ def as_list(value: Any) -> List[Any]: if isinstance(value, str): return [item.strip() for item in value.split(",") if item.strip()] return [] + + +MODEL_EXTENSIONS = { + ".ckpt", + ".pt", + ".pt2", + ".bin", + ".pth", + ".safetensors", + ".pkl", + ".sft", + ".onnx", + ".gguf", +}