From 28fbb86dceb0d859cb22990e8a993cb90e20347e Mon Sep 17 00:00:00 2001 From: Will Miao Date: Sun, 13 Sep 2026 07:59:28 +0800 Subject: [PATCH] feat(backend): gate Other Models behind opt-in management toggles Other Models management is now opt-in: enable_other_models (default false) plus the enabled_other_sub_types allow-list replace the unreleased additive enabled_other_folders key. - config._get_enabled_other_folder_keys() is the single scan gate; a new refresh_other_roots() rebuilds roots and preview roots on toggle. - ModelScanner gains a _should_keep_cached_entry() hydration hook and on_library_changed(reconcile=...) so switching a sub_type off drops its entries (and hash/autov3 rows) at load time and switching it on rescans. - OtherScanner filters location-derived entries accordingly. - Other routes reject every other type while off (or a disabled sub_type) and expose an "other_disabled" page flag; download routing returns a disabled marker instead of guessing; the download manager refuses other-type downloads and default-path routing for switched-off sub_types. - Doctor / init-status / refresh-all skip the other scanner while off; the scanner stays registered so staged pending-deletes still merge. - Tests updated with explicit opt-in fixtures plus new gating coverage. --- py/config.py | 42 +++--- py/routes/base_model_routes.py | 5 + .../handlers/download_routing_handlers.py | 27 ++++ py/routes/handlers/misc_handlers.py | 31 ++++- py/routes/handlers/model_handlers.py | 12 ++ py/routes/other_routes.py | 29 ++++- py/services/download_manager.py | 20 +++ py/services/model_scanner.py | 54 ++++++-- py/services/other_scanner.py | 10 ++ py/services/settings_manager.py | 120 ++++++++++++------ py/utils/constants.py | 56 ++++++-- tests/config/test_other_paths.py | 45 ++++++- .../routes/test_download_routing_handlers.py | 45 +++++++ tests/routes/test_misc_routes.py | 6 + tests/routes/test_other_routes.py | 47 +++++++ tests/services/test_download_manager_other.py | 42 ++++++ tests/services/test_other_scanner.py | 45 +++++++ tests/services/test_recipe_scanner.py | 2 +- tests/services/test_settings_manager.py | 29 +++++ 19 files changed, 583 insertions(+), 84 deletions(-) diff --git a/py/config.py b/py/config.py index 99d589b8..4c029f74 100644 --- a/py/config.py +++ b/py/config.py @@ -18,7 +18,6 @@ import time from .utils.cache_paths import CacheType, get_cache_file_path, get_legacy_cache_paths from .utils.constants import ( - DEFAULT_OTHER_MODEL_FOLDERS, OTHER_MODEL_FOLDER_SUBTYPES, ) from .utils.settings_paths import ( @@ -1150,28 +1149,25 @@ class Config: def _get_enabled_other_folder_keys(self) -> List[str]: """Return the OTHER_MODEL_FOLDER_SUBTYPES keys that are enabled. - Default-enabled categories come from DEFAULT_OTHER_MODEL_FOLDERS; - opt-in categories (e.g. controlnet) are added via the - ``enabled_other_folders`` setting (a list of folder_paths keys). + Other Models management is opt-in: while ``enable_other_models`` is + off (the default) no other-model folder is scanned at all. When it is + on, only the folder keys of the enabled sub_types are scanned + (text_encoder merges ``text_encoders`` with the legacy ``clip`` key). """ - keys = list(DEFAULT_OTHER_MODEL_FOLDERS) try: from .services.settings_manager import get_settings_manager - extra = get_settings_manager().get("enabled_other_folders", []) + enabled_sub_types = get_settings_manager().get_enabled_other_sub_types() except Exception: - extra = [] - if isinstance(extra, str): - extra = [extra] - if isinstance(extra, Iterable): - for key in extra: - if ( - isinstance(key, str) - and key in OTHER_MODEL_FOLDER_SUBTYPES - and key not in keys - ): - keys.append(key) - return keys + enabled_sub_types = [] + if not enabled_sub_types: + return [] + allowed = set(enabled_sub_types) + return [ + key + for key, sub_type in OTHER_MODEL_FOLDER_SUBTYPES.items() + if sub_type in allowed + ] def _prepare_other_paths( self, folder_path_map: Mapping[str, Iterable[str]] @@ -1427,6 +1423,16 @@ class Config: logger.warning(f"Error initializing other model paths: {e}") return [] + def refresh_other_roots(self) -> None: + """Rebuild other-model roots after the management toggles changed. + + Called when ``enable_other_models`` / ``enabled_other_sub_types`` are + updated so the scanner immediately reflects the new folder set without + a full application restart. + """ + self.other_roots = self._init_other_paths() + self._rebuild_preview_roots() + def get_preview_static_url(self, preview_path: str) -> str: if not preview_path: return "" diff --git a/py/routes/base_model_routes.py b/py/routes/base_model_routes.py index 4d2e842d..2bac8d7c 100644 --- a/py/routes/base_model_routes.py +++ b/py/routes/base_model_routes.py @@ -149,6 +149,7 @@ class BaseModelRoutes(ABC): settings_service=self._settings, server_i18n=self._server_i18n, logger=logger, + page_context_provider=self._get_page_context_provider(), ) listing = ModelListingHandler( service=service, @@ -250,6 +251,10 @@ class BaseModelRoutes(ABC): """Get expected model types string for error messages - to be overridden by subclasses.""" return "any model type" + def _get_page_context_provider(self): + """Optional hook returning extra template context for the page view.""" + return None + def _find_model_file(self, files): """Find the appropriate model file from the files list - can be overridden by subclasses.""" return next((file for file in files if file.get("type") in MODEL_WEIGHT_FILE_TYPES and file.get("primary") is True), None) diff --git a/py/routes/handlers/download_routing_handlers.py b/py/routes/handlers/download_routing_handlers.py index 645b38a7..20ecb7ef 100644 --- a/py/routes/handlers/download_routing_handlers.py +++ b/py/routes/handlers/download_routing_handlers.py @@ -56,11 +56,38 @@ class DownloadRoutingHandler: ) if model_type.lower() in VALID_OTHER_CIVITAI_TYPES: + from ...services.settings_manager import get_settings_manager + + settings = get_settings_manager() + if not settings.is_other_models_enabled(): + # Opt-in feature is off: never auto-route, the UI falls back to + # manual folder selection and the download manager rejects it. + return web.json_response( + { + "success": True, + "root_kind": "other", + "sub_type": None, + "disabled": True, + "reason": "other_models_disabled", + } + ) + sub_type = resolve_other_download_sub_type( model_type, file_types=(str(t) for t in file_types), selected_file_type=selected_file_type, ) + if sub_type and not settings.is_other_sub_type_enabled(sub_type): + return web.json_response( + { + "success": True, + "root_kind": "other", + "sub_type": None, + "disabled": True, + "reason": "other_sub_type_disabled", + "requested_sub_type": sub_type, + } + ) return web.json_response( { "success": True, diff --git a/py/routes/handlers/misc_handlers.py b/py/routes/handlers/misc_handlers.py index 5ca0450b..711e4eb9 100644 --- a/py/routes/handlers/misc_handlers.py +++ b/py/routes/handlers/misc_handlers.py @@ -663,6 +663,17 @@ class HealthCheckHandler: "recipe": ServiceRegistry.get_recipe_scanner, } + def _active_scanner_getters( + self, + ) -> Mapping[str, Callable[[], Awaitable[Any]]]: + """Drop the opt-in other scanner while Other Models is disabled.""" + getters = self._scanner_getters + if "other" not in getters: + return getters + if get_settings_manager().is_other_models_enabled(): + return getters + return {name: getter for name, getter in getters.items() if name != "other"} + async def health_check(self, request: web.Request) -> web.Response: return web.json_response({"status": "ok"}) @@ -674,7 +685,7 @@ class HealthCheckHandler: page accepts the update and only reloads once all scanners are done. """ pending: list[str] = [] - for name, getter in self._scanner_getters.items(): + for name, getter in self._active_scanner_getters().items(): try: scanner = await getter() except Exception: @@ -764,6 +775,14 @@ class DoctorHandler: ) self._app_version_getter = app_version_getter + def _active_scanner_factories( + self, + ) -> Sequence[tuple[str, str, Callable[[], Awaitable[Any]]]]: + """Drop the opt-in other scanner while Other Models is disabled.""" + if self._settings.is_other_models_enabled(): + return self._scanner_factories + return tuple(entry for entry in self._scanner_factories if entry[0] != "other") + async def get_doctor_diagnostics(self, request: web.Request) -> web.Response: try: client_version = (request.query.get("clientVersion") or "").strip() @@ -811,7 +830,7 @@ class DoctorHandler: repaired: list[dict[str, Any]] = [] failures: list[dict[str, str]] = [] - for model_type, label, factory in self._scanner_factories: + for model_type, label, factory in self._active_scanner_factories(): try: scanner = await factory() await scanner.get_cached_data(force_refresh=True, rebuild_cache=True) @@ -843,7 +862,7 @@ class DoctorHandler: renamed: list[dict[str, Any]] = [] try: - for model_type, label, factory in self._scanner_factories: + for model_type, label, factory in self._active_scanner_factories(): try: scanner = await factory() hash_index = getattr(scanner, "_hash_index", None) @@ -1075,7 +1094,7 @@ class DoctorHandler: overall_status = "ok" summary = "All model caches look healthy." - for model_type, label, factory in self._scanner_factories: + for model_type, label, factory in self._active_scanner_factories(): try: scanner = await factory() persisted = None @@ -1160,7 +1179,7 @@ class DoctorHandler: total_conflict_groups = 0 total_conflict_files = 0 - for model_type, label, factory in self._scanner_factories: + for model_type, label, factory in self._active_scanner_factories(): # Duplicate filename detection targets LoRAs which use basename-only # syntax (). Checkpoints/embeddings reference # models via relative paths with extensions, so conflicts there would @@ -2797,6 +2816,8 @@ class ModelLibraryHandler: # Acquire the other scanner lazily so adapters without it only # fail when the payload actually contains other-type models. + # While the opt-in feature is off the scanner still exists (its + # cache is empty), so other types simply report inLibrary=False. needs_other_scanner = any( isinstance(model, dict) and str(model.get("type", "")).lower() in other_type_aliases diff --git a/py/routes/handlers/model_handlers.py b/py/routes/handlers/model_handlers.py index 7a849fde..71fd7e26 100644 --- a/py/routes/handlers/model_handlers.py +++ b/py/routes/handlers/model_handlers.py @@ -90,6 +90,7 @@ class ModelPageView: settings_service: SettingsManager, server_i18n, logger: logging.Logger, + page_context_provider: Callable[[web.Request], Dict[str, Any]] | None = None, ) -> None: self._template_env = template_env self._template_name = template_name @@ -97,6 +98,7 @@ class ModelPageView: self._settings = settings_service self._server_i18n = server_i18n self._logger = logger + self._page_context_provider = page_context_provider def _load_supporters(self) -> dict[str, Any]: """Load supporters data from JSON file.""" @@ -210,6 +212,16 @@ class ModelPageView: self._logger.error("Error loading cache data: %s", cache_error) template_context["is_initializing"] = True + if self._page_context_provider is not None: + try: + extra_context = self._page_context_provider(request) + if isinstance(extra_context, dict): + template_context.update(extra_context) + except Exception as context_error: # pragma: no cover - logging path + self._logger.error( + "Error building page context: %s", context_error + ) + rendered = self._template_env.get_template(self._template_name).render( **template_context ) diff --git a/py/routes/other_routes.py b/py/routes/other_routes.py index 3647c840..8615d97b 100644 --- a/py/routes/other_routes.py +++ b/py/routes/other_routes.py @@ -7,7 +7,11 @@ from .model_route_registrar import ModelRouteRegistrar from ..config import config from ..services.other_model_service import OtherModelService from ..services.service_registry import ServiceRegistry -from ..utils.constants import OTHER_MODEL_FOLDER_SUBTYPES, VALID_OTHER_CIVITAI_TYPES +from ..utils.constants import ( + CIVITAI_TYPE_TO_OTHER_SUB_TYPE, + OTHER_MODEL_FOLDER_SUBTYPES, + VALID_OTHER_CIVITAI_TYPES, +) logger = logging.getLogger(__name__) @@ -50,9 +54,28 @@ class OtherRoutes(BaseModelRoutes): """Validate CivitAI model type for other models. Accepts retired CivitAI types (CLIP, CLIPVision) as well — grandfathered - models on CivitAI still carry them. + models on CivitAI still carry them. Types whose sub_type is currently + disabled (or every type while the opt-in feature is off) are rejected. """ - return model_type.lower() in VALID_OTHER_CIVITAI_TYPES + normalized = (model_type or "").strip().lower() + if normalized not in VALID_OTHER_CIVITAI_TYPES: + return False + if not self._settings.is_other_models_enabled(): + return False + + sub_type = CIVITAI_TYPE_TO_OTHER_SUB_TYPE.get(normalized) + if sub_type is None: + # CivitAI "Other" has no sub_type of its own; it is only usable + # while at least one sub_type is enabled. + return bool(self._settings.get_enabled_other_sub_types()) + return self._settings.is_other_sub_type_enabled(sub_type) + + def _get_page_context_provider(self): + """Expose the opt-in feature state to the Other Models page template.""" + return self._page_context_for_other + + def _page_context_for_other(self, request: web.Request) -> Dict[str, Any]: + return {"other_disabled": not self._settings.is_other_models_enabled()} def _get_expected_model_types(self) -> str: """Get expected model types string for error messages""" diff --git a/py/services/download_manager.py b/py/services/download_manager.py index c77f28bf..0bae9e75 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -1526,6 +1526,15 @@ class DownloadManager: elif model_type_from_info == "textualinversion": model_type = "embedding" elif model_type_from_info in VALID_OTHER_CIVITAI_TYPES: + if not get_settings_manager().is_other_models_enabled(): + return { + "success": False, + "error": ( + "Other Models management is disabled. Enable it in " + "Settings > Library before downloading VAE, upscaler, " + "text encoder or CLIP files." + ), + } model_type = "other" else: return { @@ -1774,6 +1783,17 @@ class DownloadManager: default_other_roots = ( settings_manager.get("default_other_roots") or {} ) + if other_sub_type and not settings_manager.is_other_sub_type_enabled( + other_sub_type + ): + return { + "success": False, + "error": ( + f"Other-model sub-type '{other_sub_type}' is " + f"disabled in settings. Please pick a destination " + f"folder explicitly instead of using default paths." + ), + } default_path = ( default_other_roots.get(other_sub_type) if other_sub_type diff --git a/py/services/model_scanner.py b/py/services/model_scanner.py index a6869be3..368c7d21 100644 --- a/py/services/model_scanner.py +++ b/py/services/model_scanner.py @@ -210,8 +210,14 @@ class ModelScanner: """ self._cache_version += 1 - def on_library_changed(self) -> None: - """Reset caches when the active library changes.""" + def on_library_changed(self, reconcile: bool = False) -> None: + """Reset caches when the active library changes. + + When ``reconcile`` is True an incremental reconcile runs right after + the cache is re-hydrated, so newly configured roots are scanned and + entries for removed roots are purged. Used when scanner-affecting + settings (e.g. the Other Models toggles) change. + """ self._persistent_cache = get_persistent_cache() self._cache = None self._hash_index = ModelHashIndex() @@ -229,7 +235,7 @@ class ModelScanner: if loop and not loop.is_closed(): self._loop = loop self.loop = loop - loop.create_task(self.initialize_in_background()) + loop.create_task(self.initialize_in_background(reconcile=reconcile)) def _resolve_name_display_mode(self) -> str: """Return the configured display mode for name sorting.""" @@ -460,8 +466,14 @@ class ModelScanner: _, license_flags = resolve_license_info(license_source) entry['license_flags'] = license_flags - async def initialize_in_background(self) -> None: - """Initialize cache in background using thread pool""" + async def initialize_in_background(self, reconcile: bool = False) -> None: + """Initialize cache in background using thread pool + + Args: + reconcile: When True and a persisted snapshot is hydrated, run an + incremental reconcile afterwards so the cache matches the + current root configuration. + """ try: # Set initial empty cache to avoid None reference errors if self._cache is None: @@ -501,6 +513,11 @@ class ModelScanner: logger.info( f"{self.model_type.capitalize()} cache hydrated from persisted snapshot with {len(self._cache.raw_data)} models" ) + if reconcile: + # Root configuration changed (e.g. Other Models toggles): + # pick up newly enabled folders and drop rows for folders + # that are no longer managed. + await self.get_cached_data(force_refresh=True) return # Persistent load failed; fall back to a full scan @@ -663,21 +680,33 @@ class ModelScanner: if not persisted or not persisted.raw_data: return None + # Drop entries the scanner no longer manages (e.g. an other-model + # sub_type the user just disabled) before rebuilding the indexes, so + # hash/autov3 lookups cannot resolve to unmanaged files either. + kept_items = [ + item + for item in persisted.raw_data + if self._should_keep_cached_entry(item) + ] + kept_paths = { + item.get("file_path") for item in kept_items if item.get("file_path") + } + hash_index = ModelHashIndex() for sha_value, path in persisted.hash_rows: - if sha_value and path: + if sha_value and path and path in kept_paths: hash_index.add_entry(sha_value.lower(), path) # Rebuild the AutoV3 index from the persisted autov3_index rows. These # cover every known autov3 -> path mapping regardless of whether a # sha256 row also exists for the same file. for autov3_value, path in persisted.autov3_hash_rows: - if autov3_value and path: + if autov3_value and path and path in kept_paths: hash_index.add_autov3(autov3_value.lower(), path) tags_count: Dict[str, int] = {} adjusted_raw_data: List[Dict[str, Any]] = [] - for item in persisted.raw_data: + for item in kept_items: # load_cache builds a fresh dict per row, and validate_batch below # works on its own per-entry copy when auto_repair=True, so no # additional dict copy is needed here. @@ -1435,6 +1464,15 @@ class ModelScanner: """Hook for subclasses: adjust entries loaded from the persisted cache.""" return entry + def _should_keep_cached_entry(self, entry: Dict[str, Any]) -> bool: + """Hook for subclasses: decide whether a persisted entry is still managed. + + Entries rejected here are dropped (with their hash/autov3 index rows) + while hydrating the persisted cache, so a scanner whose configured + roots shrank does not surface stale models before the next reconcile. + """ + return True + def resolve_sub_type_for_path(self, file_path: Optional[str]) -> Optional[str]: """Hook for subclasses: resolve the location-derived sub_type for a file. diff --git a/py/services/other_scanner.py b/py/services/other_scanner.py index 366e2c98..0db2101a 100644 --- a/py/services/other_scanner.py +++ b/py/services/other_scanner.py @@ -454,6 +454,16 @@ class OtherScanner(ModelScanner): entry["sub_type"] = sub_type return entry + def _should_keep_cached_entry(self, entry: Dict[str, Any]) -> bool: + """Drop persisted entries whose folder is no longer a managed root. + + sub_type is location-derived and config only maps enabled roots, so a + file under a disabled sub_type - or under any other root while the + feature is off - resolves to None here and is filtered out while the + persisted cache is hydrated. + """ + return self.resolve_sub_type_for_path(entry.get("file_path")) is not None + def get_model_roots(self) -> List[str]: """Get other-model root directories""" roots: List[str] = [] diff --git a/py/services/settings_manager.py b/py/services/settings_manager.py index cdfb41fb..0d03c2d1 100644 --- a/py/services/settings_manager.py +++ b/py/services/settings_manager.py @@ -25,11 +25,13 @@ from typing import ( from platformdirs import user_config_dir from ..utils.constants import ( + DEFAULT_ENABLED_OTHER_SUB_TYPES, DEFAULT_HASH_CHUNK_SIZE_MB, DEFAULT_PRIORITY_TAG_CONFIG, - OTHER_MODEL_FOLDER_SUBTYPES, + OTHER_SUB_TYPE_FOLDER_KEYS, SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS, VALID_OTHER_SUB_TYPES, + normalize_other_sub_types, ) from ..utils.preview_selection import VALID_MATURE_BLUR_LEVELS from ..utils.settings_paths import ( @@ -86,6 +88,10 @@ DEFAULT_SETTINGS: Dict[str, Any] = { "default_unet_root": "", "default_embedding_root": "", "default_other_roots": {}, + # Other Models management is opt-in: nothing is scanned, shown or offered + # for download until the user turns the feature on. + "enable_other_models": False, + "enabled_other_sub_types": list(DEFAULT_ENABLED_OTHER_SUB_TYPES), "recipes_path": "", "base_model_path_mappings": {}, "download_path_templates": {}, @@ -680,6 +686,42 @@ class SettingsManager: normalized[sub_type] = stripped return normalized + def is_other_models_enabled(self) -> bool: + """Return True when the opt-in Other Models management is enabled.""" + return bool(self.settings.get("enable_other_models", False)) + + def get_enabled_other_sub_types(self) -> List[str]: + """Return the enabled other-model sub_types (empty when the feature is off).""" + if not self.is_other_models_enabled(): + return [] + return normalize_other_sub_types(self.settings.get("enabled_other_sub_types")) + + def is_other_sub_type_enabled(self, sub_type: Optional[str]) -> bool: + """Return True when ``sub_type`` is currently managed.""" + if not sub_type: + return False + return sub_type in self.get_enabled_other_sub_types() + + def _apply_other_model_settings_change(self) -> None: + """Rebuild other-model roots and refresh the other scanner after a toggle.""" + try: + from ..config import config # Local import to avoid circular dependency + + config.refresh_other_roots() + except Exception as exc: # pragma: no cover - defensive logging + logger.debug("Failed to refresh other-model roots: %s", exc) + + try: + from .service_registry import ServiceRegistry # pyright: ignore[reportImportCycles] + + scanner = ServiceRegistry.get_service_sync("other_scanner") + if scanner is not None and hasattr(scanner, "on_library_changed"): + # reconcile=True lets the scanner pick up newly enabled roots and + # purge rows for folders that are no longer managed. + scanner.on_library_changed(reconcile=True) + except Exception as exc: # pragma: no cover - defensive logging + logger.debug("Failed to refresh other scanner after settings change: %s", exc) + def _has_configured_paths(self, folder_paths: Any) -> bool: if not isinstance(folder_paths, Mapping): return False @@ -951,44 +993,44 @@ class SettingsManager: updated = _check_and_auto_set("unet", "default_unet_root") or updated updated = _check_and_auto_set("embeddings", "default_embedding_root") or updated - # Other-model default roots: one entry per sub_type; candidates are the - # union of that sub_type's folder_paths keys (text_encoder merges the - # legacy 'clip' key with 'text_encoders'). - sub_type_folder_keys: Dict[str, List[str]] = {} - for folder_key, sub_type in OTHER_MODEL_FOLDER_SUBTYPES.items(): - sub_type_folder_keys.setdefault(sub_type, []).append(folder_key) - + # Other-model default roots: one entry per enabled sub_type; candidates + # are the union of that sub_type's folder_paths keys (text_encoder + # merges the legacy 'clip' key with 'text_encoders'). When the opt-in + # feature is off the existing mapping is left untouched. other_roots = self._normalize_default_other_roots( self.settings.get("default_other_roots") ) - for sub_type in VALID_OTHER_SUB_TYPES: - candidates: List[str] = [] - candidate_identities: set[str] = set() - for folder_key in sub_type_folder_keys.get(sub_type, []): - for candidate in self._get_valid_root_candidates(folder_key): - identity = _normalize_root_identity(candidate) - if identity in candidate_identities: - continue - candidate_identities.add(identity) - candidates.append(candidate) - if not candidates: - continue - current = other_roots.get(sub_type, "") - if current and _normalize_root_identity(current) in candidate_identities: - continue - other_roots[sub_type] = candidates[0] - if current: - logger.info( - "Repaired stale default_other_roots[%s] from '%s' to '%s' because it is not present in primary or extra roots", - sub_type, - current, - candidates[0], - ) - else: - logger.info( - "Auto-set default_other_roots[%s] to '%s'", sub_type, candidates[0] - ) - updated = True + if self.is_other_models_enabled(): + for sub_type in self.get_enabled_other_sub_types(): + candidates: List[str] = [] + candidate_identities: set[str] = set() + for folder_key in OTHER_SUB_TYPE_FOLDER_KEYS.get(sub_type, []): + for candidate in self._get_valid_root_candidates(folder_key): + identity = _normalize_root_identity(candidate) + if identity in candidate_identities: + continue + candidate_identities.add(identity) + candidates.append(candidate) + if not candidates: + continue + current = other_roots.get(sub_type, "") + if current and _normalize_root_identity(current) in candidate_identities: + continue + other_roots[sub_type] = candidates[0] + if current: + logger.info( + "Repaired stale default_other_roots[%s] from '%s' to '%s' because it is not present in primary or extra roots", + sub_type, + current, + candidates[0], + ) + else: + logger.info( + "Auto-set default_other_roots[%s] to '%s'", + sub_type, + candidates[0], + ) + updated = True if updated: self.settings["default_other_roots"] = other_roots @@ -1699,6 +1741,10 @@ class SettingsManager: value = self.normalize_mature_blur_level(value) elif key == "default_other_roots": value = self._normalize_default_other_roots(value, strict=True) + elif key == "enabled_other_sub_types": + value = normalize_other_sub_types(value) + elif key == "enable_other_models": + value = bool(value) elif key == "recipes_path": current_recipes_dir = self._get_effective_recipes_dir() value = self._normalize_recipes_path_value(value) @@ -1735,6 +1781,8 @@ class SettingsManager: self._save_settings() if key == "recipes_path": self._notify_library_change(self.get_active_library_name()) + if key in ("enable_other_models", "enabled_other_sub_types"): + self._apply_other_model_settings_change() if portable_switch_pending: self._finalize_portable_switch() diff --git a/py/utils/constants.py b/py/utils/constants.py index bc44ea21..4f770dd0 100644 --- a/py/utils/constants.py +++ b/py/utils/constants.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, Dict, List NSFW_LEVELS = { "PG": 1, @@ -93,17 +93,51 @@ OTHER_MODEL_FOLDER_SUBTYPES = { "clip_vision": "clip_vision", "controlnet": "controlnet", } -# folder_paths keys scanned by default; anything else in -# OTHER_MODEL_FOLDER_SUBTYPES (e.g. controlnet) is opt-in via the -# "enabled_other_folders" setting. -DEFAULT_OTHER_MODEL_FOLDERS = ( - "vae", - "upscale_models", - "text_encoders", - "clip", - "clip_vision", -) VALID_OTHER_SUB_TYPES = ["vae", "upscaler", "text_encoder", "clip_vision", "controlnet"] +# Sub-types managed when the (opt-in) Other Models feature is switched on. +# The feature itself defaults to off (``enable_other_models`` = False), so +# nothing here is scanned until the user enables it. +DEFAULT_ENABLED_OTHER_SUB_TYPES: List[str] = [ + "vae", + "upscaler", + "text_encoder", + "clip_vision", +] + + +def other_sub_type_folder_keys() -> Dict[str, List[str]]: + """Invert OTHER_MODEL_FOLDER_SUBTYPES into sub_type -> folder_paths keys. + + ``text_encoder`` maps to two folder keys (``text_encoders`` and the legacy + ``clip``), so every consumer that resolves a sub_type back to folders must + merge both. + """ + mapping: Dict[str, List[str]] = {} + for folder_key, sub_type in OTHER_MODEL_FOLDER_SUBTYPES.items(): + mapping.setdefault(sub_type, []).append(folder_key) + return mapping + + +# Precomputed inverse of OTHER_MODEL_FOLDER_SUBTYPES, keeping the table order. +OTHER_SUB_TYPE_FOLDER_KEYS: Dict[str, List[str]] = other_sub_type_folder_keys() + + +def normalize_other_sub_types(value: Any) -> List[str]: + """Normalize a stored/requested enabled-sub_type list. + + Unknown values and duplicates are dropped; the result follows the + canonical VALID_OTHER_SUB_TYPES order so the stored setting and the UI + stay stable. Non-list input falls back to the defaults. + """ + if isinstance(value, str): + candidates: Any = [value] + elif isinstance(value, (list, tuple, set)): + candidates = value + else: + return list(DEFAULT_ENABLED_OTHER_SUB_TYPES) + + allowed = {item for item in candidates if isinstance(item, str)} + return [sub_type for sub_type in VALID_OTHER_SUB_TYPES if sub_type in allowed] # CivitAI model.type values accepted by the "other" page's fetch-metadata # validation (lowercased). CLIP/CLIPVision are retired upstream but still # appear on grandfathered models. diff --git a/tests/config/test_other_paths.py b/tests/config/test_other_paths.py index 9fddc3d9..e2f20baf 100644 --- a/tests/config/test_other_paths.py +++ b/tests/config/test_other_paths.py @@ -26,6 +26,14 @@ def _make_config(**overrides) -> config_module.Config: return config +@pytest.fixture(autouse=True) +def enable_other_models(): + """Other Models is opt-in; enable it for the enabled-state tests.""" + manager = get_settings_manager() + manager.set("enable_other_models", True) + yield + + class TestPrepareOtherPaths: """Unit tests for Config._prepare_other_paths.""" @@ -196,7 +204,7 @@ class TestInitOtherPaths: controlnet_dir.mkdir() self._stub_folder_paths(monkeypatch, {"controlnet": str(controlnet_dir)}) - get_settings_manager().set("enabled_other_folders", ["controlnet"]) + get_settings_manager().set("enabled_other_sub_types", ["controlnet"]) config = _make_config() roots = config._init_other_paths() @@ -207,12 +215,45 @@ class TestInitOtherPaths: == "controlnet" ) + def test_disabled_sub_type_is_not_scanned(self, monkeypatch, tmp_path): + vae_dir = tmp_path / "vae" + upscaler_dir = tmp_path / "upscale_models" + vae_dir.mkdir() + upscaler_dir.mkdir() + + self._stub_folder_paths( + monkeypatch, {"vae": str(vae_dir), "upscale_models": str(upscaler_dir)} + ) + get_settings_manager().set("enabled_other_sub_types", ["vae"]) + + config = _make_config() + roots = config._init_other_paths() + + assert roots == [_normalize(str(vae_dir))] + assert _normalize(str(upscaler_dir)) not in config.other_root_subtypes + + def test_feature_disabled_scans_nothing(self, monkeypatch, tmp_path): + vae_dir = tmp_path / "vae" + vae_dir.mkdir() + + self._stub_folder_paths(monkeypatch, {"vae": str(vae_dir)}) + get_settings_manager().set("enable_other_models", False) + + config = _make_config() + roots = config._init_other_paths() + + assert roots == [] + assert config.other_root_subtypes == {} + assert config.other_folder_roots == {} + def test_unknown_opt_in_keys_are_ignored(self, monkeypatch, tmp_path): vae_dir = tmp_path / "vae" vae_dir.mkdir() self._stub_folder_paths(monkeypatch, {"vae": str(vae_dir)}) - get_settings_manager().set("enabled_other_folders", ["not_a_real_key", 42]) + get_settings_manager().set( + "enabled_other_sub_types", ["vae", "not_a_real_key", 42] + ) config = _make_config() roots = config._init_other_paths() diff --git a/tests/routes/test_download_routing_handlers.py b/tests/routes/test_download_routing_handlers.py index 260a9e11..4e1474b6 100644 --- a/tests/routes/test_download_routing_handlers.py +++ b/tests/routes/test_download_routing_handlers.py @@ -5,6 +5,22 @@ import json import pytest from py.routes.handlers.download_routing_handlers import DownloadRoutingHandler +from py.services.settings_manager import get_settings_manager + + +@pytest.fixture(autouse=True) +def enable_other_models(): + """Other Models is opt-in; enable every sub_type for the routing tests.""" + manager = get_settings_manager() + manager.settings["enable_other_models"] = True + manager.settings["enabled_other_sub_types"] = [ + "vae", + "upscaler", + "text_encoder", + "clip_vision", + "controlnet", + ] + yield class FakeRequest: @@ -149,3 +165,32 @@ async def test_other_invalid_selected_file_type_rejected(): FakeRequest({"model_type": "VAE", "selected_file_type": 123}) ) assert response.status == 400 + + +@pytest.mark.asyncio +async def test_other_routing_disabled_when_feature_off(): + get_settings_manager().settings["enable_other_models"] = False + + handler = DownloadRoutingHandler() + response = await handler.get_download_routing( + FakeRequest({"model_type": "VAE", "file_types": ["Model"]}) + ) + payload = json.loads(response.text) + assert payload["sub_type"] is None + assert payload["disabled"] is True + assert payload["reason"] == "other_models_disabled" + + +@pytest.mark.asyncio +async def test_other_routing_disabled_for_switched_off_sub_type(): + get_settings_manager().settings["enabled_other_sub_types"] = ["vae"] + + handler = DownloadRoutingHandler() + response = await handler.get_download_routing( + FakeRequest({"model_type": "Upscaler", "file_types": ["Model"]}) + ) + payload = json.loads(response.text) + assert payload["sub_type"] is None + assert payload["disabled"] is True + assert payload["reason"] == "other_sub_type_disabled" + assert payload["requested_sub_type"] == "upscaler" diff --git a/tests/routes/test_misc_routes.py b/tests/routes/test_misc_routes.py index 9ada287d..62908587 100644 --- a/tests/routes/test_misc_routes.py +++ b/tests/routes/test_misc_routes.py @@ -61,6 +61,12 @@ class DummySettings: def get(self, key, default=None): return self.data.get(key, default) + def is_other_models_enabled(self): + return bool(self.data.get("enable_other_models", False)) + + def get_enabled_other_sub_types(self): + return list(self.data.get("enabled_other_sub_types") or []) + def set(self, key, value): self.data[key] = value diff --git a/tests/routes/test_other_routes.py b/tests/routes/test_other_routes.py index c7085858..bfa2cd86 100644 --- a/tests/routes/test_other_routes.py +++ b/tests/routes/test_other_routes.py @@ -30,6 +30,20 @@ def routes(): return handler +@pytest.fixture(autouse=True) +def enable_other_models(): + """Other Models is opt-in; these tests exercise the enabled state.""" + from py.services.settings_manager import get_settings_manager + + manager = get_settings_manager() + manager.set("enable_other_models", True) + manager.set( + "enabled_other_sub_types", + ["vae", "upscaler", "text_encoder", "clip_vision", "controlnet"], + ) + yield + + def test_common_and_specific_routes_registered(): """Registration smoke test: /api/lm/other/* surface plus the /other page.""" app = web.Application() @@ -64,6 +78,39 @@ def test_validate_civitai_model_type_rejects_foreign_types(model_type): assert OtherRoutes()._validate_civitai_model_type(model_type) is False +def test_validate_rejects_everything_when_feature_disabled(): + from py.services.settings_manager import get_settings_manager + + get_settings_manager().set("enable_other_models", False) + + handler = OtherRoutes() + for model_type in ("VAE", "Upscaler", "TextEncoder", "CLIPVision", "Other"): + assert handler._validate_civitai_model_type(model_type) is False + + +def test_validate_rejects_switched_off_sub_type(): + from py.services.settings_manager import get_settings_manager + + get_settings_manager().set("enabled_other_sub_types", ["vae"]) + + handler = OtherRoutes() + assert handler._validate_civitai_model_type("VAE") is True + assert handler._validate_civitai_model_type("Upscaler") is False + + +def test_page_context_reports_feature_state(): + from py.services.settings_manager import get_settings_manager + + manager = get_settings_manager() + handler = OtherRoutes() + provider = handler._get_page_context_provider() + + assert provider(None) == {"other_disabled": False} + + manager.set("enable_other_models", False) + assert provider(None) == {"other_disabled": True} + + def test_get_expected_model_types_mentions_supported_types(): expected = OtherRoutes()._get_expected_model_types() for name in ("VAE", "Upscaler", "TextEncoder", "CLIPVision", "Controlnet"): diff --git a/tests/services/test_download_manager_other.py b/tests/services/test_download_manager_other.py index 532b54da..12367523 100644 --- a/tests/services/test_download_manager_other.py +++ b/tests/services/test_download_manager_other.py @@ -43,6 +43,14 @@ def isolate_settings(monkeypatch, tmp_path): "text_encoder": str(tmp_path / "text_encoders"), "clip_vision": str(tmp_path / "clip_vision"), }, + "enable_other_models": True, + "enabled_other_sub_types": [ + "vae", + "upscaler", + "text_encoder", + "clip_vision", + "controlnet", + ], "download_path_templates": { "lora": "{base_model}/{first_tag}", "checkpoint": "{base_model}/{first_tag}", @@ -231,6 +239,40 @@ async def test_download_rejects_unknown_model_type( assert result["error"].startswith("Model type") +@pytest.mark.asyncio +async def test_download_rejects_other_when_feature_disabled( + monkeypatch, scanners, metadata_provider, tmp_path +): + """The opt-in feature is off: no other-type download is accepted.""" + metadata_provider.payload = _other_payload("VAE") + get_settings_manager().settings["enable_other_models"] = False + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=99, save_dir=str(tmp_path) + ) + + assert result["success"] is False + assert "disabled" in result["error"].lower() + + +@pytest.mark.asyncio +async def test_default_paths_reject_switched_off_sub_type( + monkeypatch, scanners, metadata_provider, tmp_path +): + """A disabled sub_type refuses default-path routing (manual pick still works).""" + metadata_provider.payload = _other_payload("VAE") + get_settings_manager().settings["enabled_other_sub_types"] = ["upscaler"] + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=99, use_default_paths=True + ) + + assert result["success"] is False + assert "disabled" in result["error"].lower() + + @pytest.mark.asyncio async def test_early_gate_checks_other_scanner( monkeypatch, scanners, metadata_provider, tmp_path diff --git a/tests/services/test_other_scanner.py b/tests/services/test_other_scanner.py index 14882acd..6c42a600 100644 --- a/tests/services/test_other_scanner.py +++ b/tests/services/test_other_scanner.py @@ -173,6 +173,51 @@ class TestOtherScannerRoots: assert result["sub_type"] == "upscaler" +class TestOtherScannerHydrationFilter: + """Persisted entries for roots that are no longer managed are dropped.""" + + def test_keeps_entries_under_enabled_roots(self, other_config): + scanner = _make_scanner() + assert ( + scanner._should_keep_cached_entry( + {"file_path": f"{other_config['vae']}/model.safetensors"} + ) + is True + ) + + def test_drops_entries_under_disabled_root(self, other_config, monkeypatch): + scanner = _make_scanner() + # Only vae stays managed; the upscaler root disappeared from the map. + monkeypatch.setattr( + config_module.config, + "other_root_subtypes", + {other_config["vae"]: "vae"}, + ) + + assert ( + scanner._should_keep_cached_entry( + {"file_path": f"{other_config['upscaler']}/model.safetensors"} + ) + is False + ) + assert ( + scanner._should_keep_cached_entry( + {"file_path": f"{other_config['vae']}/model.safetensors"} + ) + is True + ) + + def test_drops_everything_when_feature_off(self, monkeypatch): + monkeypatch.setattr(config_module.config, "other_root_subtypes", {}) + scanner = _make_scanner() + assert ( + scanner._should_keep_cached_entry( + {"file_path": "/models/vae/model.safetensors"} + ) + is False + ) + + class TestOtherScannerLazyHash: """Lazy hashing: pending by default, singleflight on-demand calculation.""" diff --git a/tests/services/test_recipe_scanner.py b/tests/services/test_recipe_scanner.py index e2bad6c2..88b6f47b 100644 --- a/tests/services/test_recipe_scanner.py +++ b/tests/services/test_recipe_scanner.py @@ -2513,7 +2513,7 @@ async def test_on_library_changed_bumps_cache_version(tmp_path: Path, monkeypatc scanner = DummyScanner(str(tmp_path)) assert scanner.cache_version == 0 - async def _noop_initialize() -> None: + async def _noop_initialize(reconcile: bool = False) -> None: pass monkeypatch.setattr(scanner, "initialize_in_background", _noop_initialize) diff --git a/tests/services/test_settings_manager.py b/tests/services/test_settings_manager.py index 9d4dc2be..0cbb2166 100644 --- a/tests/services/test_settings_manager.py +++ b/tests/services/test_settings_manager.py @@ -1222,7 +1222,33 @@ def test_default_other_roots_stay_empty_without_other_folders(manager): assert manager.get("default_other_roots") == {} +def test_other_models_disabled_by_default(manager): + assert manager.is_other_models_enabled() is False + assert manager.get_enabled_other_sub_types() == [] + assert manager.is_other_sub_type_enabled("vae") is False + + +def test_auto_set_default_other_roots_skipped_when_feature_off(manager): + manager.settings["enable_other_models"] = False + manager.settings["default_other_roots"] = {} + manager.settings["folder_paths"] = {"vae": ["/vae"]} + + manager._auto_set_default_roots() + + assert manager.get("default_other_roots") == {} + + +def test_set_enabled_other_sub_types_normalizes(manager): + manager.settings["enable_other_models"] = True + manager.set("enabled_other_sub_types", ["controlnet", "vae", "nope", "vae", 42]) + + assert manager.get("enabled_other_sub_types") == ["vae", "controlnet"] + assert manager.is_other_sub_type_enabled("vae") is True + assert manager.is_other_sub_type_enabled("upscaler") is False + + def test_auto_set_default_other_roots(manager): + manager.settings["enable_other_models"] = True manager.settings["default_other_roots"] = {} manager.settings["folder_paths"] = { "vae": ["/vae"], @@ -1243,6 +1269,7 @@ def test_auto_set_default_other_roots(manager): def test_auto_set_default_other_roots_text_encoder_dual_key_union(manager): """text_encoder candidates merge text_encoders and the legacy clip key.""" + manager.settings["enable_other_models"] = True manager.settings["default_other_roots"] = {} manager.settings["folder_paths"] = { "clip": ["/legacy-clip"], @@ -1261,6 +1288,7 @@ def test_auto_set_default_other_roots_text_encoder_dual_key_union(manager): def test_auto_set_default_other_roots_repairs_stale(manager): + manager.settings["enable_other_models"] = True manager.settings["default_other_roots"] = {"vae": "/stale-vae"} manager.settings["folder_paths"] = {"vae": ["/vae"]} @@ -1270,6 +1298,7 @@ def test_auto_set_default_other_roots_repairs_stale(manager): def test_auto_set_default_other_roots_uses_extra_folder_paths(manager): + manager.settings["enable_other_models"] = True manager.settings["default_other_roots"] = {} manager.settings["folder_paths"] = {"vae": []} manager.settings["extra_folder_paths"] = {"vae": ["/extra-vae"]}