diff --git a/core/matcher.py b/core/matcher.py index 260559d..4bab19d 100644 --- a/core/matcher.py +++ b/core/matcher.py @@ -284,7 +284,7 @@ def base_model_score(candidate: str, preferred: Optional[str]) -> float: # ==================== CENTRALIZED MATCHING & CONFIDENCE HELPERS ==================== -from .type_utils import MODEL_EXTENSIONS +from .type_utils import MODEL_EXTENSIONS, PRECISION_FORMAT_SUFFIXES MODEL_FILE_EXTENSIONS = MODEL_EXTENSIONS @@ -392,6 +392,19 @@ def calculate_candidate_rank( return matches, confidence + score +def clean_filename_for_search(filename: str) -> str: + """ + Clean up filename for better search results. + Remove common suffixes that might prevent matches. + """ + if not filename: + return "" + basename = os.path.basename(filename).strip() + base = os.path.splitext(basename)[0] + pattern = r"[-_]?(" + "|".join(PRECISION_FORMAT_SUFFIXES) + r").*$" + return re.sub(pattern, "", base, flags=re.IGNORECASE).strip(" -_") + + def build_filename_search_queries(filename: str) -> List[str]: """ Build a list of candidate search queries from a filename by stripping @@ -405,7 +418,7 @@ def build_filename_search_queries(filename: str) -> List[str]: stem_lower = stem.lower() # Pattern to match common precision/format suffixes - pattern = r"[-_]?(fp16|fp32|fp8|fp4|bf16|e4m3fn|mixed|scaled|pruned|emaonly|q4|q8)" + pattern = r"[-_]?(" + "|".join(PRECISION_FORMAT_SUFFIXES) + r")" # 1. Strip suffix followed by any other characters (like clean_filename_for_search) cleaned = re.sub(pattern + r".*$", "", stem, flags=re.IGNORECASE).strip(" -_") diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index f42399a..6d98ffb 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -19,7 +19,7 @@ log = create_module_logger(__name__) from ..path_utils import write_json_atomic, METADATA_DIR, read_json_safe from ..type_utils import check_credential_http, parse_size_header as _parse_size_header, fetch_remote_file_size, fetch_remote_file_size_cached, clear_remote_size_cache, extract_file_size -from ..matcher import build_filename_search_queries +from ..matcher import build_filename_search_queries, clean_filename_for_search HF_API_URL = "https://huggingface.co/api" HF_AUTHOR_FALLBACKS = ["Comfy-Org"] @@ -358,23 +358,6 @@ def _fetch_remote_file_size_bytes( return fetch_remote_file_size_cached(probe_url, headers=headers, timeout=timeout) - -def clean_filename_for_search(filename: str) -> str: - """ - Clean up filename for better search results. - Remove common suffixes that might prevent matches. - """ - base = os.path.splitext(filename)[0] - # Remove common precision/format suffixes - base = re.sub( - r"[-_]?(fp16|fp32|fp8|fp4|bf16|e4m3fn|scaled|pruned|emaonly|mixed|q4|q8).*$", - "", - base, - flags=re.IGNORECASE, - ) - return base - - def _build_huggingface_result( repo_id: str, file_path: str, diff --git a/core/sources/lora_manager_archive.py b/core/sources/lora_manager_archive.py index 5182d10..af87b08 100644 --- a/core/sources/lora_manager_archive.py +++ b/core/sources/lora_manager_archive.py @@ -18,7 +18,7 @@ from ..matcher import ( base_model_score as _base_model_score, calculate_candidate_rank, ) -from ..type_utils import select_primary_model_file +from ..type_utils import select_primary_model_file, get_generic_filename_tokens from ..progress import report_progress, get_progress_reporter from ..log_system import create_module_logger log = create_module_logger(__name__) @@ -29,22 +29,7 @@ DEFAULT_DB_RELATIVE_PATH = os.path.join("civitai", "civitai.sqlite") _search_cache: Dict[str, Any] = {} _db_path_cache: Optional[str] = None -GENERIC_FILENAME_TOKENS = { - "safetensors", - "ckpt", - "pt", - "pth", - "bin", - "onnx", - "gguf", - "fp16", - "fp32", - "bf16", - "e4m3fn", - "scaled", - "pruned", - "emaonly", -} +GENERIC_FILENAME_TOKENS = get_generic_filename_tokens() def clear_search_cache(): diff --git a/core/type_utils.py b/core/type_utils.py index 80e8e94..67ec3fe 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -60,6 +60,29 @@ MODEL_EXTENSIONS = { } +PRECISION_FORMAT_SUFFIXES = [ + "fp16", + "fp32", + "fp8", + "fp4", + "bf16", + "e4m3fn", + "mixed", + "scaled", + "pruned", + "emaonly", + "q4", + "q8", +] + + +def get_generic_filename_tokens() -> set: + """Return a set of generic filename tokens (extensions and suffixes) in lowercase.""" + dotless_extensions = {ext.lstrip(".") for ext in MODEL_EXTENSIONS} + return dotless_extensions | set(PRECISION_FORMAT_SUFFIXES) + + + def first_non_empty(*values: Any, default: Any = "") -> Any: """ Return the first value that is not None, not an empty/whitespace-only string,