refactor(services): share weight-file extension set between rematch and find_matching_models

This commit is contained in:
Will Miao
2026-08-19 21:54:22 +08:00
parent 8a16034135
commit e57e11897e
2 changed files with 30 additions and 9 deletions
+25 -2
View File
@@ -25,6 +25,28 @@ from .cache_health_monitor import CacheHealthMonitor, CacheHealthStatus
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Canonical set of weight-file extensions stripped when normalizing model
# names for matching (ModelScanner.find_matching_models and the recipe rematch
# filename key share this set). It is the union of the LoRA scanner set
# ({".safetensors"}) and the Checkpoint scanner set (ComfyUI's
# supported_pt_extensions plus ".gguf") so type-blind lookups (lora +
# checkpoint merged) cover every format either scanner indexes. ".safebin"
# is deliberately absent — no scanner indexes it, so a recipe entry
# "model.safebin" must not be bound to a local "model.safetensors".
WEIGHT_FILE_EXTENSIONS = frozenset(
{
".safetensors",
".ckpt",
".pt",
".pt2",
".bin",
".pth",
".pkl",
".sft",
".gguf",
}
)
def _is_excluded_dir(name: str) -> bool: def _is_excluded_dir(name: str) -> bool:
"""Return True when a directory entry must be skipped during model walks. """Return True when a directory entry must be skipped during model walks.
@@ -2155,10 +2177,11 @@ class ModelScanner:
``base_model`` is given, confident mismatches are rejected while ``base_model`` is given, confident mismatches are rejected while
unknowns on either side stay eligible (lenient guard). unknowns on either side stay eligible (lenient guard).
``extensions`` should be the scanner's own ``file_extensions`` so ``extensions`` should be the scanner's own ``file_extensions`` so
suffix stripping only covers formats the scanner actually indexes. suffix stripping only covers formats the scanner actually indexes;
when omitted, the shared :data:`WEIGHT_FILE_EXTENSIONS` set is used.
""" """
# Longest first so overlapping suffixes strip correctly. # Longest first so overlapping suffixes strip correctly.
exts = sorted(extensions or (".safetensors", ".ckpt", ".pt", ".bin"), key=len, reverse=True) exts = sorted(extensions or WEIGHT_FILE_EXTENSIONS, key=len, reverse=True)
normalized_name = str(name).replace("\\", "/").casefold() normalized_name = str(name).replace("\\", "/").casefold()
for ext in exts: for ext in exts:
+5 -7
View File
@@ -15,6 +15,7 @@ from ..config import config
from ..utils.constants import VALID_CHECKPOINT_SUB_TYPES, VALID_LORA_TYPES from ..utils.constants import VALID_CHECKPOINT_SUB_TYPES, VALID_LORA_TYPES
from ..utils.file_utils import calculate_autov3 from ..utils.file_utils import calculate_autov3
from ..utils.recipe_open_stats import RecipeOpenStats from ..utils.recipe_open_stats import RecipeOpenStats
from .model_scanner import WEIGHT_FILE_EXTENSIONS
from .recipe_cache import RecipeCache from .recipe_cache import RecipeCache
from .recipes.errors import RecipeNotFoundError, RecipePersistenceError from .recipes.errors import RecipeNotFoundError, RecipePersistenceError
from natsort import natsorted from natsort import natsorted
@@ -36,11 +37,6 @@ logger = logging.getLogger(__name__)
# explicitly to "diffusion_model" (mirrors Oracle R2-F1). # explicitly to "diffusion_model" (mirrors Oracle R2-F1).
_CHECKPOINT_MODEL_TYPE_ALIASES = {"diffusionmodel": "diffusion_model"} _CHECKPOINT_MODEL_TYPE_ALIASES = {"diffusionmodel": "diffusion_model"}
# Known weight-file extensions stripped by _normalize_filename_key. Names are
# stored extensionless on both sides, so splitext would misread dotted stems
# ("my.mix" -> "my") and silently collide distinct models.
_WEIGHT_FILE_EXTS = (".safetensors", ".ckpt", ".pt", ".pth", ".gguf", ".bin", ".safebin", ".sft")
class RecipeScanner: class RecipeScanner:
"""Service for scanning and managing recipe images""" """Service for scanning and managing recipe images"""
@@ -179,13 +175,15 @@ class RecipeScanner:
Only known weight-file extensions are stripped — names are stored Only known weight-file extensions are stripped — names are stored
extensionless on both sides, so splitext would misread dotted stems extensionless on both sides, so splitext would misread dotted stems
("my.mix" -> "my") and collide distinct models. ("my.mix" -> "my") and collide distinct models. The extension set is
shared with ModelScanner.find_matching_models, and is iterated longest
first to keep the strip ordering identical to that function.
""" """
if not name: if not name:
return "" return ""
basename = os.path.basename(name.replace("\\", "/")) basename = os.path.basename(name.replace("\\", "/"))
lower = basename.lower() lower = basename.lower()
for ext in _WEIGHT_FILE_EXTS: for ext in sorted(WEIGHT_FILE_EXTENSIONS, key=len, reverse=True):
if lower.endswith(ext): if lower.endswith(ext):
basename = basename[: -len(ext)] basename = basename[: -len(ext)]
break break