From f2a7297cb99feb7639bf7d8df0a387f65c398703 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Sat, 12 Sep 2026 15:56:47 +0800 Subject: [PATCH] feat(backend): CivitAI download support for other model types with subtype routing --- .../handlers/download_routing_handlers.py | 26 +- py/routes/handlers/misc_handlers.py | 26 + py/routes/other_routes.py | 34 +- py/services/download_manager.py | 93 ++- py/services/download_routing.py | 54 +- py/services/settings_manager.py | 114 ++++ py/utils/constants.py | 15 + py/utils/example_images_download_manager.py | 12 + py/utils/models.py | 3 +- settings.json.example | 1 + .../routes/test_download_routing_handlers.py | 58 ++ tests/routes/test_misc_routes.py | 34 +- tests/routes/test_other_routes.py | 47 ++ tests/services/test_download_manager_other.py | 565 ++++++++++++++++++ tests/services/test_download_routing.py | 102 ++++ .../test_example_images_other_dispatch.py | 168 ++++++ tests/services/test_other_scanner.py | 22 +- tests/services/test_settings_manager.py | 134 +++++ 18 files changed, 1492 insertions(+), 16 deletions(-) create mode 100644 tests/services/test_download_manager_other.py create mode 100644 tests/services/test_example_images_other_dispatch.py diff --git a/py/routes/handlers/download_routing_handlers.py b/py/routes/handlers/download_routing_handlers.py index 0d3c7a12..645b38a7 100644 --- a/py/routes/handlers/download_routing_handlers.py +++ b/py/routes/handlers/download_routing_handlers.py @@ -7,7 +7,11 @@ import logging from aiohttp import web -from ...services.download_routing import is_diffusion_model_download +from ...services.download_routing import ( + is_diffusion_model_download, + resolve_other_download_sub_type, +) +from ...utils.constants import VALID_OTHER_CIVITAI_TYPES logger = logging.getLogger(__name__) @@ -31,6 +35,7 @@ class DownloadRoutingHandler: model_type = payload.get("model_type", "") base_model = payload.get("base_model") or "" file_types = payload.get("file_types") or [] + selected_file_type = payload.get("selected_file_type") if not isinstance(model_type, str) or not model_type: return web.json_response( @@ -44,6 +49,25 @@ class DownloadRoutingHandler: }, status=400, ) + if selected_file_type is not None and not isinstance(selected_file_type, str): + return web.json_response( + {"success": False, "error": "selected_file_type must be a string"}, + status=400, + ) + + if model_type.lower() in VALID_OTHER_CIVITAI_TYPES: + sub_type = resolve_other_download_sub_type( + model_type, + file_types=(str(t) for t in file_types), + selected_file_type=selected_file_type, + ) + return web.json_response( + { + "success": True, + "root_kind": "other", + "sub_type": sub_type, + } + ) is_diffusion = is_diffusion_model_download( model_type, diff --git a/py/routes/handlers/misc_handlers.py b/py/routes/handlers/misc_handlers.py index c461841d..5ca0450b 100644 --- a/py/routes/handlers/misc_handlers.py +++ b/py/routes/handlers/misc_handlers.py @@ -53,6 +53,7 @@ from ...utils.constants import ( PREVIEW_EXTENSIONS, SUPPORTED_MEDIA_EXTENSIONS, VALID_LORA_TYPES, + VALID_OTHER_CIVITAI_TYPES, ) from .hf_handlers import HfHandler from .agent_handlers import AgentHandler @@ -2068,6 +2069,7 @@ class ServiceRegistryAdapter: get_embedding_scanner: Callable[[], Awaitable[Any]] get_downloaded_version_history_service: Callable[[], Awaitable[Any]] get_backup_service: Callable[[], Awaitable[Any]] = _noop_backup_service + get_other_scanner: Callable[[], Awaitable[Any]] = ServiceRegistry.get_other_scanner class ModelLibraryHandler: @@ -2789,12 +2791,30 @@ class ModelLibraryHandler: model_type.lower() for model_type in CIVITAI_USER_MODEL_TYPES } lora_type_aliases = {model_type.lower() for model_type in VALID_LORA_TYPES} + other_type_aliases = { + model_type.lower() for model_type in VALID_OTHER_CIVITAI_TYPES + } + + # Acquire the other scanner lazily so adapters without it only + # fail when the payload actually contains other-type models. + needs_other_scanner = any( + isinstance(model, dict) + and str(model.get("type", "")).lower() in other_type_aliases + for model in models + ) + other_scanner = None + if needs_other_scanner: + other_scanner = await self._service_registry.get_other_scanner() type_scanner_map: Dict[str, Any] = { **{alias: lora_scanner for alias in lora_type_aliases}, "checkpoint": checkpoint_scanner, "textualinversion": embedding_scanner, } + if other_scanner is not None: + type_scanner_map.update( + {alias: other_scanner for alias in other_type_aliases} + ) versions: list[dict[str, Any]] = [] history_service = await self._get_download_history_service() @@ -2818,12 +2838,17 @@ class ModelLibraryHandler: "embedding", model_ids, ) + other_downloaded = await history_service.get_downloaded_version_ids_bulk( + "other", + model_ids, + ) downloaded_version_map: Dict[str, Dict[int, set[int]]] = { "lora": lora_downloaded, "locon": lora_downloaded, "dora": lora_downloaded, "checkpoint": checkpoint_downloaded, "textualinversion": embedding_downloaded, + **{alias: other_downloaded for alias in VALID_OTHER_CIVITAI_TYPES}, } for model in models: if not isinstance(model, dict): @@ -3982,6 +4007,7 @@ def build_service_registry_adapter() -> ServiceRegistryAdapter: get_lora_scanner=ServiceRegistry.get_lora_scanner, get_checkpoint_scanner=ServiceRegistry.get_checkpoint_scanner, get_embedding_scanner=ServiceRegistry.get_embedding_scanner, + get_other_scanner=ServiceRegistry.get_other_scanner, get_downloaded_version_history_service=ServiceRegistry.get_downloaded_version_history_service, get_backup_service=ServiceRegistry.get_backup_service, ) diff --git a/py/routes/other_routes.py b/py/routes/other_routes.py index be36e066..3647c840 100644 --- a/py/routes/other_routes.py +++ b/py/routes/other_routes.py @@ -1,12 +1,13 @@ import logging -from typing import Any, Dict +from typing import Any, Dict, List from aiohttp import web from .base_model_routes import BaseModelRoutes 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 VALID_OTHER_CIVITAI_TYPES +from ..utils.constants import OTHER_MODEL_FOLDER_SUBTYPES, VALID_OTHER_CIVITAI_TYPES logger = logging.getLogger(__name__) @@ -41,6 +42,9 @@ class OtherRoutes(BaseModelRoutes): """Setup Other-model-specific routes""" # Other-model info by name registrar.add_prefixed_route('GET', '/api/lm/{prefix}/info/{name}', prefix, self.get_other_model_info) + # Other-model roots grouped by sub_type (text_encoders + legacy clip + # are aggregated under text_encoder) + registrar.add_prefixed_route('GET', '/api/lm/{prefix}/roots_by_subtype', prefix, self.get_roots_by_subtype) def _validate_civitai_model_type(self, model_type: str) -> bool: """Validate CivitAI model type for other models. @@ -58,6 +62,32 @@ class OtherRoutes(BaseModelRoutes): """Parse other-model-specific parameters (none in Phase 1).""" return {} + async def get_roots_by_subtype(self, request: web.Request) -> web.Response: + """Return other-model roots grouped by sub_type. + + Aggregates the per-folder_paths-key roots from config + (``text_encoders`` and the legacy ``clip`` key both land under + ``text_encoder``). + """ + try: + roots_by_subtype: Dict[str, List[str]] = {} + for key, roots in (config.other_folder_roots or {}).items(): + sub_type = OTHER_MODEL_FOLDER_SUBTYPES.get(key) + if not sub_type: + continue + bucket = roots_by_subtype.setdefault(sub_type, []) + for root in roots: + if root and root not in bucket: + bucket.append(root) + return web.json_response( + {"success": True, "roots_by_subtype": roots_by_subtype} + ) + except Exception as e: + logger.error(f"Error getting other roots by sub_type: {e}", exc_info=True) + return web.json_response( + {"success": False, "error": str(e)}, status=500 + ) + async def get_other_model_info(self, request: web.Request) -> web.Response: """Get detailed information for a specific other model by name""" try: diff --git a/py/services/download_manager.py b/py/services/download_manager.py index 2f52f337..c77f28bf 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -17,12 +17,18 @@ from dataclasses import dataclass, field import uuid from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, cast from urllib.parse import urlparse -from ..utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata +from ..utils.models import ( + LoraMetadata, + CheckpointMetadata, + EmbeddingMetadata, + OtherModelMetadata, +) from ..utils.constants import ( CARD_PREVIEW_WIDTH, MODEL_WEIGHT_FILE_TYPES, SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS, VALID_LORA_TYPES, + VALID_OTHER_CIVITAI_TYPES, ) from ..utils.civitai_utils import normalize_civitai_download_url, rewrite_preview_url from ..utils.file_utils import calculate_sha256, calculate_autov3 @@ -31,7 +37,7 @@ from ..utils.utils import sanitize_folder_name from ..utils.exif_utils import ExifUtils from ..utils.metadata_manager import MetadataManager from .service_registry import ServiceRegistry -from .download_routing import is_diffusion_model_download +from .download_routing import is_diffusion_model_download, resolve_other_download_sub_type from .settings_manager import get_settings_manager from .metadata_service import get_default_metadata_provider, get_metadata_provider from .downloader import get_downloader, DownloadProgress, DownloadStreamControl @@ -228,12 +234,21 @@ class DownloadManager: return False async def _get_scanner_for_model_type(self, model_type: str): - """Return the scanner responsible for the given model type.""" + """Return the scanner responsible for the given model type. + + Every supported type resolves explicitly — an unknown type must never + fall through to the lora scanner (an "other" download would silently + dedupe against the lora library). + """ if model_type == "checkpoint": return await self._get_checkpoint_scanner() if model_type == "embedding": return await ServiceRegistry.get_embedding_scanner() - return await self._get_lora_scanner() + if model_type == "other": + return await ServiceRegistry.get_other_scanner() + if model_type == "lora": + return await self._get_lora_scanner() + raise ValueError(f'Unknown model type "{model_type}"') @staticmethod def _resolve_target_file( @@ -978,6 +993,8 @@ class DownloadManager: return CheckpointMetadata.from_civitai_info(version_info, file_info, save_path) if model_type == "embedding": return EmbeddingMetadata.from_civitai_info(version_info, file_info, save_path) + if model_type == "other": + return OtherModelMetadata.from_civitai_info(version_info, file_info, save_path) return LoraMetadata.from_civitai_info(version_info, file_info, save_path) def _resolve_save_path_from_persisted_record(self, record: Dict[str, Any]) -> Optional[str]: @@ -1438,6 +1455,7 @@ class DownloadManager: lora_scanner = await self._get_lora_scanner() checkpoint_scanner = await self._get_checkpoint_scanner() embedding_scanner = await ServiceRegistry.get_embedding_scanner() + other_scanner = await ServiceRegistry.get_other_scanner() # Check lora scanner first if await lora_scanner.check_model_version_exists(model_version_id): @@ -1462,6 +1480,13 @@ class DownloadManager: "error": "Model version already exists in embedding library", } + # Check other scanner + if await other_scanner.check_model_version_exists(model_version_id): + return { + "success": False, + "error": "Model version already exists in other library", + } + # Use CivArchive provider directly when source is 'civarchive' # This prioritizes CivArchive metadata (with mirror availability info) over Civitai if source == "civarchive": @@ -1500,6 +1525,8 @@ class DownloadManager: model_type = "lora" elif model_type_from_info == "textualinversion": model_type = "embedding" + elif model_type_from_info in VALID_OTHER_CIVITAI_TYPES: + model_type = "other" else: return { "success": False, @@ -1686,6 +1713,13 @@ class DownloadManager: "success": False, "error": "Model version already exists in embedding library", } + elif model_type == "other": + other_scanner = await ServiceRegistry.get_other_scanner() + if await other_scanner.check_model_version_exists(version_id): + return { + "success": False, + "error": "Model version already exists in other library", + } # Handle use_default_paths if use_default_paths: @@ -1725,6 +1759,45 @@ class DownloadManager: "error": "Default embedding root path not set in settings", } save_dir = default_path + elif model_type == "other": + other_sub_type = resolve_other_download_sub_type( + model_type_from_info, + file_types=( + f.get("type", "") + for f in version_info.get("files", []) + if isinstance(f, dict) + ), + selected_file_type=( + target_file.get("type") if explicit_file else None + ), + ) + default_other_roots = ( + settings_manager.get("default_other_roots") or {} + ) + default_path = ( + default_other_roots.get(other_sub_type) + if other_sub_type + else None + ) + if not isinstance(default_path, str) or not default_path: + if other_sub_type: + detail = ( + f"No default root configured for other-model " + f"sub-type '{other_sub_type}'" + ) + else: + detail = ( + "Could not determine the other-model sub-type " + "from the model metadata" + ) + return { + "success": False, + "error": ( + f"{detail}. Please pick a destination folder " + f"explicitly instead of using default paths." + ), + } + save_dir = default_path # Calculate relative path using template relative_path = self._calculate_relative_path(version_info, model_type) @@ -1921,6 +1994,11 @@ class DownloadManager: version_info, file_info, save_path ) logger.info(f"Creating EmbeddingMetadata for {file_name}") + elif model_type == "other": + metadata = OtherModelMetadata.from_civitai_info( + version_info, file_info, save_path + ) + logger.info(f"Creating OtherModelMetadata for {file_name}") else: return { "success": False, @@ -2133,6 +2211,8 @@ class DownloadManager: scanner = await self._get_checkpoint_scanner() elif model_type == "embedding": scanner = await ServiceRegistry.get_embedding_scanner() + elif model_type == "other": + scanner = await ServiceRegistry.get_other_scanner() except Exception as exc: logger.debug("Failed to acquire scanner for %s models: %s", model_type, exc) @@ -2629,6 +2709,9 @@ class DownloadManager: elif model_type == "embedding": scanner = await ServiceRegistry.get_embedding_scanner() logger.info(f"Updating embedding cache for {actual_file_paths[0]}") + elif model_type == "other": + scanner = await ServiceRegistry.get_other_scanner() + logger.info(f"Updating other-model cache for {actual_file_paths[0]}") adjust_cached_entry = ( getattr(scanner, "adjust_cached_entry", None) @@ -2718,7 +2801,7 @@ class DownloadManager: return {"success": False, "error": str(e)} def _get_supported_extensions_for_type(self, model_type: str) -> Set[str]: - if model_type == "checkpoint": + if model_type in ("checkpoint", "other"): return { ".ckpt", ".pt", diff --git a/py/services/download_routing.py b/py/services/download_routing.py index ec09d24d..a18c2080 100644 --- a/py/services/download_routing.py +++ b/py/services/download_routing.py @@ -10,9 +10,13 @@ two can never disagree. from __future__ import annotations import logging -from typing import Iterable +from typing import Iterable, Optional -from ..utils.constants import DIFFUSION_MODEL_BASE_MODELS +from ..utils.constants import ( + CIVITAI_FILE_TYPE_TO_OTHER_SUB_TYPE, + CIVITAI_TYPE_TO_OTHER_SUB_TYPE, + DIFFUSION_MODEL_BASE_MODELS, +) logger = logging.getLogger(__name__) @@ -51,3 +55,49 @@ def is_diffusion_model_download( return True return False + + +def resolve_other_download_sub_type( + civitai_model_type: str, + file_types: Iterable[str] = (), + selected_file_type: Optional[str] = None, +) -> Optional[str]: + """Resolve the "other"-page sub_type for a download. + + Fixed priority (locked design, docs/plans/other-models-page.md §9.2): + + 1. Explicit user file pick — when the picked file's type maps, it wins + even when model.type maps to something else. + 2. model.type via CIVITAI_TYPE_TO_OTHER_SUB_TYPE. + 3. file.type fallback — only when model.type maps to nothing. Must NOT + override a mapped model.type: checkpoint models routinely bundle + VAE/Text Encoder component files. + 4. Still undecidable -> None (caller must ask the user for a folder). + """ + if selected_file_type: + mapped = CIVITAI_FILE_TYPE_TO_OTHER_SUB_TYPE.get(selected_file_type) + if mapped: + logger.info( + "Explicit file pick type '%s' routes other download to '%s'", + selected_file_type, + mapped, + ) + return mapped + + normalized_model_type = (civitai_model_type or "").strip().lower() + mapped = CIVITAI_TYPE_TO_OTHER_SUB_TYPE.get(normalized_model_type) + if mapped: + return mapped + + for file_type in file_types: + mapped = CIVITAI_FILE_TYPE_TO_OTHER_SUB_TYPE.get(file_type) + if mapped: + logger.info( + "model.type '%s' unmapped; file type '%s' routes other download to '%s'", + civitai_model_type, + file_type, + mapped, + ) + return mapped + + return None diff --git a/py/services/settings_manager.py b/py/services/settings_manager.py index 08d71ac0..cdfb41fb 100644 --- a/py/services/settings_manager.py +++ b/py/services/settings_manager.py @@ -27,7 +27,9 @@ from platformdirs import user_config_dir from ..utils.constants import ( DEFAULT_HASH_CHUNK_SIZE_MB, DEFAULT_PRIORITY_TAG_CONFIG, + OTHER_MODEL_FOLDER_SUBTYPES, SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS, + VALID_OTHER_SUB_TYPES, ) from ..utils.preview_selection import VALID_MATURE_BLUR_LEVELS from ..utils.settings_paths import ( @@ -83,6 +85,7 @@ DEFAULT_SETTINGS: Dict[str, Any] = { "default_checkpoint_root": "", "default_unet_root": "", "default_embedding_root": "", + "default_other_roots": {}, "recipes_path": "", "base_model_path_mappings": {}, "download_path_templates": {}, @@ -309,6 +312,7 @@ class SettingsManager: default_checkpoint_root=merged.get("default_checkpoint_root"), default_unet_root=merged.get("default_unet_root"), default_embedding_root=merged.get("default_embedding_root"), + default_other_roots=merged.get("default_other_roots"), recipes_path=merged.get("recipes_path"), ) } @@ -443,6 +447,7 @@ class SettingsManager: ), default_unet_root=self.settings.get("default_unet_root", ""), default_embedding_root=self.settings.get("default_embedding_root", ""), + default_other_roots=self.settings.get("default_other_roots"), recipes_path=self.settings.get("recipes_path", ""), ) libraries = {library_name: library_payload} @@ -494,6 +499,7 @@ class SettingsManager: default_checkpoint_root=data.get("default_checkpoint_root"), default_unet_root=data.get("default_unet_root"), default_embedding_root=data.get("default_embedding_root"), + default_other_roots=data.get("default_other_roots"), recipes_path=data.get("recipes_path"), metadata=data.get("metadata"), base=data, @@ -541,6 +547,9 @@ class SettingsManager: self.settings["default_embedding_root"] = active_library.get( "default_embedding_root", "" ) + self.settings["default_other_roots"] = self._normalize_default_other_roots( + active_library.get("default_other_roots", {}) + ) self.settings["recipes_path"] = active_library.get("recipes_path", "") if save: @@ -558,6 +567,7 @@ class SettingsManager: default_checkpoint_root: Optional[str] = None, default_unet_root: Optional[str] = None, default_embedding_root: Optional[str] = None, + default_other_roots: Optional[Mapping[str, str]] = None, recipes_path: Optional[str] = None, metadata: Optional[Mapping[str, Any]] = None, base: Optional[Mapping[str, Any]] = None, @@ -597,6 +607,15 @@ class SettingsManager: else: payload.setdefault("default_embedding_root", "") + if default_other_roots is not None: + payload["default_other_roots"] = self._normalize_default_other_roots( + default_other_roots + ) + else: + payload["default_other_roots"] = self._normalize_default_other_roots( + payload.get("default_other_roots", {}) + ) + if recipes_path is not None: payload["recipes_path"] = recipes_path else: @@ -632,6 +651,35 @@ class SettingsManager: normalized[key] = cleaned return normalized + def _normalize_default_other_roots( + self, value: Any, *, strict: bool = False + ) -> Dict[str, str]: + """Normalize a ``default_other_roots`` mapping ({sub_type: root path}). + + Unknown sub_type keys and non-string/empty paths are dropped; with + ``strict=True`` unknown sub_type keys raise instead (used by ``set()`` + so typos in API payloads surface as errors). + """ + if not isinstance(value, Mapping): + if strict and value is not None: + raise ValueError("default_other_roots must be a mapping") + return {} + normalized: Dict[str, str] = {} + for sub_type, path in value.items(): + if sub_type not in VALID_OTHER_SUB_TYPES: + if strict: + raise ValueError( + f"Unknown other-model sub-type '{sub_type}'; " + f"expected one of {sorted(VALID_OTHER_SUB_TYPES)}" + ) + continue + if not isinstance(path, str): + continue + stripped = path.strip() + if stripped: + normalized[sub_type] = stripped + return normalized + def _has_configured_paths(self, folder_paths: Any) -> bool: if not isinstance(folder_paths, Mapping): return False @@ -744,6 +792,7 @@ class SettingsManager: default_checkpoint_root: Optional[str] = None, default_unet_root: Optional[str] = None, default_embedding_root: Optional[str] = None, + default_other_roots: Optional[Mapping[str, str]] = None, recipes_path: Optional[str] = None, ) -> bool: libraries = self.settings.get("libraries", {}) @@ -794,6 +843,14 @@ class SettingsManager: library["default_embedding_root"] = default_embedding_root changed = True + if default_other_roots is not None: + normalized_other_roots = self._normalize_default_other_roots( + default_other_roots + ) + if library.get("default_other_roots") != normalized_other_roots: + library["default_other_roots"] = normalized_other_roots + changed = True + if recipes_path is not None and library.get("recipes_path") != recipes_path: library["recipes_path"] = recipes_path changed = True @@ -894,12 +951,53 @@ 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_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 updated: + self.settings["default_other_roots"] = other_roots self._update_active_library_entry( default_lora_root=self.settings.get("default_lora_root"), default_checkpoint_root=self.settings.get("default_checkpoint_root"), default_unet_root=self.settings.get("default_unet_root"), default_embedding_root=self.settings.get("default_embedding_root"), + default_other_roots=other_roots, ) if self._bootstrap_reason == "missing": self._needs_initial_save = True @@ -1599,6 +1697,8 @@ class SettingsManager: value = self.normalize_download_skip_base_models(value) elif key == "mature_blur_level": value = self.normalize_mature_blur_level(value) + elif key == "default_other_roots": + value = self._normalize_default_other_roots(value, strict=True) elif key == "recipes_path": current_recipes_dir = self._get_effective_recipes_dir() value = self._normalize_recipes_path_value(value) @@ -1626,6 +1726,8 @@ class SettingsManager: self._update_active_library_entry(default_unet_root=str(value)) elif key == "default_embedding_root": self._update_active_library_entry(default_embedding_root=str(value)) + elif key == "default_other_roots": + self._update_active_library_entry(default_other_roots=value) elif key == "recipes_path": self._update_active_library_entry(recipes_path=str(value)) elif key == "model_name_display": @@ -1796,6 +1898,7 @@ class SettingsManager: "lora_scanner", "checkpoint_scanner", "embedding_scanner", + "other_scanner", "recipe_scanner", ): service = ServiceRegistry.get_service_sync(service_name) @@ -1960,6 +2063,7 @@ class SettingsManager: default_checkpoint_root: Optional[str] = None, default_unet_root: Optional[str] = None, default_embedding_root: Optional[str] = None, + default_other_roots: Optional[Mapping[str, str]] = None, recipes_path: Optional[str] = None, metadata: Optional[Mapping[str, Any]] = None, activate: bool = False, @@ -2004,6 +2108,11 @@ class SettingsManager: if default_embedding_root is not None else existing.get("default_embedding_root") ), + default_other_roots=( + default_other_roots + if default_other_roots is not None + else existing.get("default_other_roots") + ), recipes_path=( recipes_path if recipes_path is not None @@ -2036,6 +2145,7 @@ class SettingsManager: default_checkpoint_root: str = "", default_unet_root: str = "", default_embedding_root: str = "", + default_other_roots: Optional[Mapping[str, str]] = None, recipes_path: str = "", metadata: Optional[Mapping[str, Any]] = None, activate: bool = False, @@ -2054,6 +2164,7 @@ class SettingsManager: default_checkpoint_root=default_checkpoint_root, default_unet_root=default_unet_root, default_embedding_root=default_embedding_root, + default_other_roots=default_other_roots, recipes_path=recipes_path, metadata=metadata, activate=activate, @@ -2114,6 +2225,7 @@ class SettingsManager: default_checkpoint_root: Optional[str] = None, default_unet_root: Optional[str] = None, default_embedding_root: Optional[str] = None, + default_other_roots: Optional[Mapping[str, str]] = None, recipes_path: Optional[str] = None, ) -> None: """Update folder paths for the active library.""" @@ -2127,6 +2239,7 @@ class SettingsManager: default_checkpoint_root=default_checkpoint_root, default_unet_root=default_unet_root, default_embedding_root=default_embedding_root, + default_other_roots=default_other_roots, recipes_path=recipes_path, activate=True, ) @@ -2151,6 +2264,7 @@ class SettingsManager: "lora_scanner", "checkpoint_scanner", "embedding_scanner", + "other_scanner", "recipe_scanner", "model_update_service", ): diff --git a/py/utils/constants.py b/py/utils/constants.py index fd8c4812..bc44ea21 100644 --- a/py/utils/constants.py +++ b/py/utils/constants.py @@ -126,6 +126,20 @@ CIVITAI_TYPE_TO_OTHER_SUB_TYPE = { "controlnet": "controlnet", } +# CivitAI ModelFile.type values -> internal sub_type for the "other" model +# page. Used for download routing only, and strictly as an explicit user file +# pick or a fallback when model.type maps to nothing — checkpoint models +# routinely bundle VAE/Text Encoder component files, so file types must never +# override a mapped model.type. +CIVITAI_FILE_TYPE_TO_OTHER_SUB_TYPE = { + "VAE": "vae", + "Upscaler": "upscaler", + "Text Encoder": "text_encoder", + "Vision Encoder": "clip_vision", + "CLIPVision": "clip_vision", + "ControlNet": "controlnet", +} + # Backward compatibility alias VALID_LORA_TYPES = VALID_LORA_SUB_TYPES @@ -134,6 +148,7 @@ CIVITAI_USER_MODEL_TYPES = [ *VALID_LORA_TYPES, "textualinversion", "checkpoint", + *sorted(VALID_OTHER_CIVITAI_TYPES), ] # Default chunk size in megabytes used for hashing large files. diff --git a/py/utils/example_images_download_manager.py b/py/utils/example_images_download_manager.py index 80a754ff..82a0e36f 100644 --- a/py/utils/example_images_download_manager.py +++ b/py/utils/example_images_download_manager.py @@ -420,6 +420,10 @@ class DownloadManager: embedding_scanner = await ServiceRegistry.get_embedding_scanner() scanners.append(("embedding", embedding_scanner)) + if "other" in model_types: + other_scanner = await ServiceRegistry.get_other_scanner() + scanners.append(("other", other_scanner)) + # Load progress file to check processed models (async to avoid blocking) settings_manager = get_settings_manager() active_library = settings_manager.get_active_library_name() @@ -600,6 +604,10 @@ class DownloadManager: embedding_scanner = await ServiceRegistry.get_embedding_scanner() scanners.append(("embedding", embedding_scanner)) + if "other" in model_types: + other_scanner = await ServiceRegistry.get_other_scanner() + scanners.append(("other", other_scanner)) + # Get all models all_models = [] for scanner_type, scanner in scanners: @@ -1098,6 +1106,10 @@ class DownloadManager: embedding_scanner = await ServiceRegistry.get_embedding_scanner() scanners.append(("embedding", embedding_scanner)) + if "other" in model_types: + other_scanner = await ServiceRegistry.get_other_scanner() + scanners.append(("other", other_scanner)) + # Find the specified models models_to_process = [] for scanner_type, scanner in scanners: diff --git a/py/utils/models.py b/py/utils/models.py index f9961cf6..ebbf71c7 100644 --- a/py/utils/models.py +++ b/py/utils/models.py @@ -340,7 +340,8 @@ class OtherModelMetadata(BaseModelMetadata): sha256_value = (file_info.get("hashes") or {}).get("SHA256", "").lower() # Map the CivitAI model type onto our sub_types; unknown types keep the # placeholder until the scanner re-derives sub_type from the location. - civitai_type = str(version_info.get("type", "") or "").lower() + # The type lives at version["model"]["type"], not version["type"]. + civitai_type = str((version_info.get("model") or {}).get("type", "") or "").lower() sub_type = CIVITAI_TYPE_TO_OTHER_SUB_TYPE.get(civitai_type, "vae") # Extract tags and description if available diff --git a/settings.json.example b/settings.json.example index 101a217d..ba09276a 100644 --- a/settings.json.example +++ b/settings.json.example @@ -31,5 +31,6 @@ "C:/path/to/your/clip_vision_folder" ] }, + "default_other_roots": {}, "auto_organize_exclusions": [] } diff --git a/tests/routes/test_download_routing_handlers.py b/tests/routes/test_download_routing_handlers.py index 48bcf767..260a9e11 100644 --- a/tests/routes/test_download_routing_handlers.py +++ b/tests/routes/test_download_routing_handlers.py @@ -91,3 +91,61 @@ async def test_invalid_json_rejected(): FakeRequest(json.JSONDecodeError("bad", "", 0)) ) assert response.status == 400 + + +@pytest.mark.asyncio +async def test_other_model_type_returns_sub_type(): + handler = DownloadRoutingHandler() + response = await handler.get_download_routing( + FakeRequest({"model_type": "TextEncoder", "file_types": ["Model"]}) + ) + payload = json.loads(response.text) + assert response.status == 200 + assert payload == {"success": True, "root_kind": "other", "sub_type": "text_encoder"} + + +@pytest.mark.asyncio +async def test_other_explicit_file_pick_wins(): + handler = DownloadRoutingHandler() + response = await handler.get_download_routing( + FakeRequest( + { + "model_type": "Other", + "file_types": ["Model"], + "selected_file_type": "VAE", + } + ) + ) + payload = json.loads(response.text) + assert payload["root_kind"] == "other" + assert payload["sub_type"] == "vae" + + +@pytest.mark.asyncio +async def test_other_file_type_fallback_when_model_type_unmapped(): + handler = DownloadRoutingHandler() + response = await handler.get_download_routing( + FakeRequest({"model_type": "Other", "file_types": ["Model", "Upscaler"]}) + ) + payload = json.loads(response.text) + assert payload["sub_type"] == "upscaler" + + +@pytest.mark.asyncio +async def test_other_undecidable_sub_type_is_none(): + handler = DownloadRoutingHandler() + response = await handler.get_download_routing( + FakeRequest({"model_type": "Other", "file_types": ["Model"]}) + ) + payload = json.loads(response.text) + assert response.status == 200 + assert payload == {"success": True, "root_kind": "other", "sub_type": None} + + +@pytest.mark.asyncio +async def test_other_invalid_selected_file_type_rejected(): + handler = DownloadRoutingHandler() + response = await handler.get_download_routing( + FakeRequest({"model_type": "VAE", "selected_file_type": 123}) + ) + assert response.status == 400 diff --git a/tests/routes/test_misc_routes.py b/tests/routes/test_misc_routes.py index 6b6b51be..9ada287d 100644 --- a/tests/routes/test_misc_routes.py +++ b/tests/routes/test_misc_routes.py @@ -1133,8 +1133,8 @@ async def test_get_civitai_user_models_marks_library_versions(): }, { "id": 4, - "name": "Unsupported", - "type": "Other", + "name": "VAE Model", + "type": "VAE", "modelVersions": [ { "id": 400, @@ -1142,6 +1142,17 @@ async def test_get_civitai_user_models_marks_library_versions(): } ], }, + { + "id": 5, + "name": "Unsupported", + "type": "Wildcard", + "modelVersions": [ + { + "id": 500, + "name": "v1", + } + ], + }, ] provider = FakeUserModelsProvider(models) @@ -1152,6 +1163,7 @@ async def test_get_civitai_user_models_marks_library_versions(): lora_scanner = FakeExistenceScanner({101}) checkpoint_scanner = FakeExistenceScanner() embedding_scanner = FakeExistenceScanner({202}) + other_scanner = FakeExistenceScanner({400}) async def lora_factory(): return lora_scanner @@ -1162,11 +1174,15 @@ async def test_get_civitai_user_models_marks_library_versions(): async def embedding_factory(): return embedding_scanner + async def other_factory(): + return other_scanner + handler = ModelLibraryHandler( ServiceRegistryAdapter( get_lora_scanner=lora_factory, get_checkpoint_scanner=checkpoint_factory, get_embedding_scanner=embedding_factory, + get_other_scanner=other_factory, get_downloaded_version_history_service=lambda: fake_download_history_service_factory(), ), metadata_provider_factory=provider_factory, @@ -1240,6 +1256,18 @@ async def test_get_civitai_user_models_marks_library_versions(): "inLibrary": False, "hasBeenDownloaded": False, }, + { + "modelId": 4, + "versionId": 400, + "modelName": "VAE Model", + "versionName": "v1", + "type": "VAE", + "tags": [], + "baseModel": None, + "thumbnailUrl": None, + "inLibrary": True, + "hasBeenDownloaded": False, + }, ] assert provider.received_usernames == ["pixel"] @@ -1351,7 +1379,7 @@ async def test_get_civitai_user_models_returns_pagination_fields(): { "id": 2, "name": "Unsupported", - "type": "Other", + "type": "Wildcard", "modelVersions": [{"id": 200, "name": "v1"}], }, ] diff --git a/tests/routes/test_other_routes.py b/tests/routes/test_other_routes.py index e87fee94..c7085858 100644 --- a/tests/routes/test_other_routes.py +++ b/tests/routes/test_other_routes.py @@ -116,3 +116,50 @@ async def test_initialize_services_builds_other_model_service(monkeypatch): assert isinstance(handler.service, OtherModelService) assert handler.service.model_type == "other" assert handler.service.scanner is sentinel_scanner + + +def test_roots_by_subtype_route_registered(): + app = web.Application() + OtherRoutes().setup_routes(app) + + registered = {(route.method, route.resource.canonical) for route in app.router.routes()} + + assert ("GET", "/api/lm/other/roots_by_subtype") in registered + + +async def test_get_roots_by_subtype_aggregates_folder_keys(monkeypatch): + """text_encoders and the legacy clip key both land under text_encoder.""" + from py.config import config + + monkeypatch.setattr( + config, + "other_folder_roots", + { + "vae": ["/models/vae", "/models/vae2"], + "text_encoders": ["/models/text_encoders"], + "clip": ["/models/clip_legacy"], + "upscale_models": ["/models/upscale"], + "unknown_key": ["/models/ignored"], + }, + ) + + response = await OtherRoutes().get_roots_by_subtype(DummyRequest()) + payload = json.loads(response.text) + + assert payload["success"] is True + assert payload["roots_by_subtype"] == { + "vae": ["/models/vae", "/models/vae2"], + "text_encoder": ["/models/text_encoders", "/models/clip_legacy"], + "upscaler": ["/models/upscale"], + } + + +async def test_get_roots_by_subtype_empty_config(monkeypatch): + from py.config import config + + monkeypatch.setattr(config, "other_folder_roots", {}) + + response = await OtherRoutes().get_roots_by_subtype(DummyRequest()) + payload = json.loads(response.text) + + assert payload == {"success": True, "roots_by_subtype": {}} diff --git a/tests/services/test_download_manager_other.py b/tests/services/test_download_manager_other.py new file mode 100644 index 00000000..532b54da --- /dev/null +++ b/tests/services/test_download_manager_other.py @@ -0,0 +1,565 @@ +"""DownloadManager support for the "other" model type (VAE, upscaler, ...). + +Covers the Phase-2 scatter points from docs/plans/other-models-page.md §9.1: +type map acceptance, existence gates consulting the other scanner (never +falling through to the lora scanner), per-sub_type default roots, resume +metadata and the archive extension set. +""" + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from py.services import aria2_transfer_state +from py.services import download_manager +from py.services.download_manager import DownloadManager +from py.services.service_registry import ServiceRegistry +from py.services.settings_manager import SettingsManager, get_settings_manager + + +@pytest.fixture(autouse=True) +def reset_download_manager(): + """Ensure each test operates on a fresh singleton.""" + DownloadManager._instance = None + yield + DownloadManager._instance = None + + +@pytest.fixture(autouse=True) +def isolate_settings(monkeypatch, tmp_path): + """Point settings writes at a temporary directory to avoid touching real files.""" + manager = get_settings_manager() + default_settings = manager._get_default_settings() + default_settings.update( + { + "default_lora_root": str(tmp_path / "loras"), + "default_checkpoint_root": str(tmp_path / "checkpoints"), + "default_embedding_root": str(tmp_path / "embeddings"), + "default_other_roots": { + "vae": str(tmp_path / "vae"), + "upscaler": str(tmp_path / "upscale_models"), + "text_encoder": str(tmp_path / "text_encoders"), + "clip_vision": str(tmp_path / "clip_vision"), + }, + "download_path_templates": { + "lora": "{base_model}/{first_tag}", + "checkpoint": "{base_model}/{first_tag}", + "embedding": "{base_model}/{first_tag}", + "other": "", + }, + "skip_previously_downloaded_model_versions": False, + "download_skip_base_models": [], + } + ) + monkeypatch.setattr(manager, "settings", default_settings) + monkeypatch.setattr(SettingsManager, "_save_settings", lambda self: None) + + +@pytest.fixture(autouse=True) +def isolate_aria2_state(monkeypatch, tmp_path): + state_path = tmp_path / "cache" / "aria2" / "downloads.json" + monkeypatch.setattr( + aria2_transfer_state, + "get_aria2_state_path", + lambda: str(state_path), + ) + + +@pytest.fixture(autouse=True) +def stub_metadata(monkeypatch): + class _StubMetadata: + def __init__(self, save_path: str): + self.file_path = save_path + self.sha256 = "sha256" + self.file_name = Path(save_path).stem + + def _make_class(name): + @staticmethod + def from_civitai_info(_version_info, _file_info, save_path): + metadata = _StubMetadata(save_path) + metadata.metadata_class = name + return metadata + + return type(name, (), {"from_civitai_info": from_civitai_info}) + + monkeypatch.setattr(download_manager, "LoraMetadata", _make_class("LoraMetadata")) + monkeypatch.setattr( + download_manager, "CheckpointMetadata", _make_class("CheckpointMetadata") + ) + monkeypatch.setattr( + download_manager, "EmbeddingMetadata", _make_class("EmbeddingMetadata") + ) + monkeypatch.setattr( + download_manager, "OtherModelMetadata", _make_class("OtherModelMetadata") + ) + + +class DummyScanner: + def __init__(self, exists: bool = False, raw_data=None): + self.exists = exists + self.calls = [] + self._cache = SimpleNamespace(raw_data=list(raw_data or [])) + + async def check_model_version_exists(self, version_id): + self.calls.append(version_id) + return self.exists + + async def get_cached_data(self): + return self._cache + + +@pytest.fixture +def scanners(monkeypatch): + lora_scanner = DummyScanner() + checkpoint_scanner = DummyScanner() + embedding_scanner = DummyScanner() + other_scanner = DummyScanner() + + monkeypatch.setattr( + ServiceRegistry, "get_lora_scanner", AsyncMock(return_value=lora_scanner) + ) + monkeypatch.setattr( + ServiceRegistry, + "get_checkpoint_scanner", + AsyncMock(return_value=checkpoint_scanner), + ) + monkeypatch.setattr( + ServiceRegistry, + "get_embedding_scanner", + AsyncMock(return_value=embedding_scanner), + ) + monkeypatch.setattr( + ServiceRegistry, + "get_other_scanner", + AsyncMock(return_value=other_scanner), + ) + + return SimpleNamespace( + lora=lora_scanner, + checkpoint=checkpoint_scanner, + embedding=embedding_scanner, + other=other_scanner, + ) + + +def _other_payload(civitai_type: str, *, files=None) -> dict: + return { + "id": 42, + "model": {"type": civitai_type, "tags": ["utility"]}, + "baseModel": "SDXL 1.0", + "creator": {"username": "Author"}, + "files": files + or [ + { + "type": "Model", + "primary": True, + "downloadUrl": "https://example.invalid/file.safetensors", + "name": "file.safetensors", + } + ], + } + + +@pytest.fixture +def metadata_provider(monkeypatch): + class DummyProvider: + def __init__(self): + self.calls = [] + self.payload = _other_payload("VAE") + + async def get_model_version(self, model_id, model_version_id): + self.calls.append((model_id, model_version_id)) + return self.payload + + provider = DummyProvider() + monkeypatch.setattr( + download_manager, + "get_default_metadata_provider", + AsyncMock(return_value=provider), + ) + return provider + + +def _capture_execute(monkeypatch, captured): + async def fake_execute_download(self, **kwargs): + captured.update(kwargs) + return {"success": True} + + monkeypatch.setattr( + DownloadManager, "_execute_download", fake_execute_download, raising=False + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "civitai_type", + ["VAE", "Upscaler", "TextEncoder", "CLIP", "CLIPVision", "Controlnet", "Other"], +) +async def test_download_accepts_other_model_types( + monkeypatch, scanners, metadata_provider, tmp_path, civitai_type +): + """All VALID_OTHER_CIVITAI_TYPES route to model_type 'other'.""" + metadata_provider.payload = _other_payload(civitai_type) + + captured = {} + _capture_execute(monkeypatch, captured) + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=99, save_dir=str(tmp_path) + ) + + assert result["success"] is True + assert captured["model_type"] == "other" + assert captured["metadata"].metadata_class == "OtherModelMetadata" + + +@pytest.mark.asyncio +async def test_download_rejects_unknown_model_type( + monkeypatch, scanners, metadata_provider, tmp_path +): + metadata_provider.payload = _other_payload("Workflow") + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=99, save_dir=str(tmp_path) + ) + + assert result["success"] is False + assert result["error"].startswith("Model type") + + +@pytest.mark.asyncio +async def test_early_gate_checks_other_scanner( + monkeypatch, scanners, metadata_provider, tmp_path +): + scanners.other.exists = True + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=101, save_dir=str(tmp_path) + ) + + assert result["success"] is False + assert result["error"] == "Model version already exists in other library" + assert scanners.other.calls == [101] + + +@pytest.mark.asyncio +async def test_scanner_dispatch_has_no_lora_fall_through(scanners): + """The Phase-2 trap: 'other' must reach the other scanner explicitly, and + unknown types must raise instead of silently deduping against loras.""" + manager = DownloadManager() + + scanner = await manager._get_scanner_for_model_type("other") + assert scanner is scanners.other + + scanner = await manager._get_scanner_for_model_type("lora") + assert scanner is scanners.lora + + with pytest.raises(ValueError): + await manager._get_scanner_for_model_type("bogus") + + +@pytest.mark.asyncio +async def test_explicit_file_gate_uses_other_scanner_not_lora( + monkeypatch, scanners, metadata_provider, tmp_path +): + """A matching local entry in the LORA library must not block an 'other' + download; only the other scanner's library is consulted.""" + local_entry = { + "file_name": "file", + "sha256": "deadbeef", + "civitai": {"id": 42}, + } + scanners.lora._cache = SimpleNamespace(raw_data=[dict(local_entry)]) + scanners.other._cache = SimpleNamespace(raw_data=[]) + + metadata_provider.payload = _other_payload( + "VAE", + files=[ + { + "id": 7, + "type": "Model", + "primary": True, + "name": "file.safetensors", + "hashes": {"SHA256": "deadbeef"}, + "downloadUrl": "https://example.invalid/file.safetensors", + } + ], + ) + + captured = {} + _capture_execute(monkeypatch, captured) + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=42, + save_dir=str(tmp_path), + file_params={"id": 7, "type": "Model"}, + ) + + assert result["success"] is True + assert captured["model_type"] == "other" + + +@pytest.mark.asyncio +async def test_explicit_file_gate_blocks_when_other_scanner_has_file( + monkeypatch, scanners, metadata_provider, tmp_path +): + local_entry = { + "file_name": "file", + "sha256": "deadbeef", + "civitai": {"id": 42}, + } + scanners.other._cache = SimpleNamespace(raw_data=[dict(local_entry)]) + + metadata_provider.payload = _other_payload( + "VAE", + files=[ + { + "id": 7, + "type": "Model", + "primary": True, + "name": "file.safetensors", + "hashes": {"SHA256": "deadbeef"}, + "downloadUrl": "https://example.invalid/file.safetensors", + } + ], + ) + + execute_mock = AsyncMock(return_value={"success": True}) + monkeypatch.setattr(DownloadManager, "_execute_download", execute_mock) + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=42, + save_dir=str(tmp_path), + file_params={"id": 7, "type": "Model"}, + ) + + assert result["success"] is False + assert "already exists in other library" in result["error"] + assert execute_mock.await_count == 0 + + +@pytest.mark.asyncio +async def test_version_level_fallback_gate_checks_other_scanner( + monkeypatch, scanners, metadata_provider, tmp_path +): + """file_params that resolve to nothing fall back to the version-level + gate, which must consult the other scanner.""" + scanners.other.exists = True + + execute_mock = AsyncMock(return_value={"success": True}) + monkeypatch.setattr(DownloadManager, "_execute_download", execute_mock) + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=101, + save_dir=str(tmp_path), + file_params={"id": 999999, "type": "Model"}, + ) + + assert result["success"] is False + assert result["error"] == "Model version already exists in other library" + assert scanners.other.calls == [101] + assert execute_mock.await_count == 0 + + +@pytest.mark.asyncio +async def test_default_paths_use_per_sub_type_root( + monkeypatch, scanners, metadata_provider, tmp_path +): + """model.type VAE -> default_other_roots['vae'].""" + captured = {} + _capture_execute(monkeypatch, captured) + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=99, use_default_paths=True + ) + + assert result["success"] is True + assert str(tmp_path / "vae") in str(captured["save_dir"]) + + +@pytest.mark.asyncio +async def test_default_paths_file_type_fallback_for_unmapped_model_type( + monkeypatch, scanners, metadata_provider, tmp_path +): + """model.type 'Other' maps to nothing; a 'Upscaler' file type decides.""" + metadata_provider.payload = _other_payload( + "Other", + files=[ + { + "type": "Upscaler", + "primary": True, + "downloadUrl": "https://example.invalid/upscaler.safetensors", + "name": "upscaler.safetensors", + } + ], + ) + + captured = {} + _capture_execute(monkeypatch, captured) + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=99, use_default_paths=True + ) + + assert result["success"] is True + assert str(tmp_path / "upscale_models") in str(captured["save_dir"]) + + +@pytest.mark.asyncio +async def test_default_paths_explicit_file_pick_wins( + monkeypatch, scanners, metadata_provider, tmp_path +): + """An explicit pick of a bundled VAE component file routes to the vae + root even though model.type maps to upscaler.""" + metadata_provider.payload = _other_payload( + "Upscaler", + files=[ + { + "id": 1, + "type": "Model", + "primary": True, + "downloadUrl": "https://example.invalid/model.safetensors", + "name": "model.safetensors", + }, + { + "id": 2, + "type": "VAE", + "downloadUrl": "https://example.invalid/bundled-vae.safetensors", + "name": "bundled-vae.safetensors", + }, + ], + ) + + captured = {} + _capture_execute(monkeypatch, captured) + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=99, + use_default_paths=True, + file_params={"id": 2, "type": "VAE"}, + ) + + assert result["success"] is True + assert str(tmp_path / "vae") in str(captured["save_dir"]) + assert captured["download_urls"] == [ + "https://example.invalid/bundled-vae.safetensors" + ] + + +@pytest.mark.asyncio +async def test_default_paths_errors_when_sub_type_root_unconfigured( + monkeypatch, scanners, metadata_provider, tmp_path +): + """controlnet has no configured default root in the fixture settings.""" + metadata_provider.payload = _other_payload("Controlnet") + + execute_mock = AsyncMock(return_value={"success": True}) + monkeypatch.setattr(DownloadManager, "_execute_download", execute_mock) + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=99, use_default_paths=True + ) + + assert result["success"] is False + assert "controlnet" in result["error"] + assert execute_mock.await_count == 0 + + +@pytest.mark.asyncio +async def test_default_paths_errors_when_sub_type_undecidable( + monkeypatch, scanners, metadata_provider, tmp_path +): + """model.type 'Other' with only plain 'Model' files: never silently + default to the vae folder — error and ask for an explicit folder.""" + metadata_provider.payload = _other_payload("Other") + + execute_mock = AsyncMock(return_value={"success": True}) + monkeypatch.setattr(DownloadManager, "_execute_download", execute_mock) + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=99, use_default_paths=True + ) + + assert result["success"] is False + assert "sub-type" in result["error"] + assert execute_mock.await_count == 0 + + +@pytest.mark.asyncio +async def test_civarchive_source_same_payload_shape( + monkeypatch, scanners, metadata_provider, tmp_path +): + """CivArchive downloads walk the same path with the same payload shape.""" + metadata_provider.payload = _other_payload("TextEncoder") + + captured = {} + _capture_execute(monkeypatch, captured) + + manager = DownloadManager() + result = await manager.download_from_civitai( + model_version_id=99, save_dir=str(tmp_path), source="civarchive" + ) + + assert result["success"] is True + assert captured["model_type"] == "other" + + +def test_build_metadata_for_resume_uses_other_metadata(): + manager = DownloadManager() + metadata = manager._build_metadata_for_resume( + model_type="other", + version_info={"model": {"type": "VAE"}}, + file_info={"name": "file.safetensors"}, + save_path="/tmp/file.safetensors", + ) + assert metadata.metadata_class == "OtherModelMetadata" + + +def test_other_extension_set_matches_checkpoint(): + manager = DownloadManager() + extensions = manager._get_supported_extensions_for_type("other") + assert extensions == manager._get_supported_extensions_for_type("checkpoint") + assert ".gguf" in extensions + assert ".safetensors" in extensions + + +@pytest.mark.asyncio +async def test_sync_downloaded_version_uses_other_scanner(monkeypatch, scanners): + """Update tracking for a downloaded other-model version consults the + other scanner for local versions.""" + + class FakeUpdateService: + def __init__(self): + self.calls = [] + + async def update_in_library_versions( + self, model_type, model_id, version_ids, version_info=None + ): + self.calls.append((model_type, model_id, version_ids)) + + update_service = FakeUpdateService() + monkeypatch.setattr( + ServiceRegistry, + "get_model_update_service", + AsyncMock(return_value=update_service), + ) + + manager = DownloadManager() + await manager._sync_downloaded_version( + "other", 7, {"id": 42, "model": {"id": 7}} + ) + + assert update_service.calls == [("other", 7, [42])] diff --git a/tests/services/test_download_routing.py b/tests/services/test_download_routing.py index 2d9dfb22..07d78bd1 100644 --- a/tests/services/test_download_routing.py +++ b/tests/services/test_download_routing.py @@ -38,3 +38,105 @@ def test_non_checkpoint_types_never_route_to_unet(): def test_empty_inputs_stay_on_checkpoint_roots(): assert not is_diffusion_model_download("checkpoint") assert not is_diffusion_model_download("checkpoint", file_types=[], base_model="") + + +from py.services.download_routing import resolve_other_download_sub_type + + +class TestResolveOtherDownloadSubType: + """Fixed priority: explicit file pick > model.type > file.type fallback.""" + + def test_explicit_file_pick_wins_over_model_type(self): + """User explicitly picked a VAE component file of a Checkpoint model — + the picked file type wins.""" + assert ( + resolve_other_download_sub_type( + "Checkpoint", file_types=["Model", "VAE"], selected_file_type="VAE" + ) + == "vae" + ) + + @pytest.mark.parametrize( + "selected,expected", + [ + ("VAE", "vae"), + ("Upscaler", "upscaler"), + ("Text Encoder", "text_encoder"), + ("Vision Encoder", "clip_vision"), + ("CLIPVision", "clip_vision"), + ("ControlNet", "controlnet"), + ], + ) + def test_explicit_file_pick_maps_all_known_types(self, selected, expected): + assert ( + resolve_other_download_sub_type("Other", selected_file_type=selected) + == expected + ) + + @pytest.mark.parametrize( + "model_type,expected", + [ + ("VAE", "vae"), + ("Upscaler", "upscaler"), + ("TextEncoder", "text_encoder"), + ("CLIP", "text_encoder"), + ("CLIPVision", "clip_vision"), + ("Controlnet", "controlnet"), + ], + ) + def test_model_type_mapping(self, model_type, expected): + assert resolve_other_download_sub_type(model_type) == expected + + def test_model_type_beats_unmappable_file_pick(self): + """An explicit pick whose file type does not map (e.g. plain 'Model') + falls through to model.type.""" + assert ( + resolve_other_download_sub_type( + "TextEncoder", selected_file_type="Model" + ) + == "text_encoder" + ) + + def test_bundled_component_files_never_override_model_type(self): + """Anti-misrouting: a TextEncoder model bundling a VAE component file + must stay text_encoder — file types are a fallback, not an override.""" + assert ( + resolve_other_download_sub_type( + "TextEncoder", file_types=["Model", "VAE"] + ) + == "text_encoder" + ) + assert ( + resolve_other_download_sub_type( + "Controlnet", file_types=["Model", "Text Encoder"] + ) + == "controlnet" + ) + + def test_file_type_fallback_when_model_type_unmapped(self): + """model.type 'Other' (or retired values) maps to nothing, so the + first mappable file type decides.""" + assert ( + resolve_other_download_sub_type("Other", file_types=["Model", "Upscaler"]) + == "upscaler" + ) + + def test_file_type_fallback_for_civarchive_payload(self): + """CivArchive-shaped payload: same fields, same decision path.""" + assert ( + resolve_other_download_sub_type( + "Other", + file_types=["Config", "Text Encoder"], + ) + == "text_encoder" + ) + + @pytest.mark.parametrize("model_type", ["Other", "", "SomethingNew"]) + def test_undecidable_returns_none(self, model_type): + assert ( + resolve_other_download_sub_type(model_type, file_types=["Model"]) is None + ) + assert resolve_other_download_sub_type(model_type) is None + + def test_model_type_matching_is_case_insensitive(self): + assert resolve_other_download_sub_type("vAe") == "vae" diff --git a/tests/services/test_example_images_other_dispatch.py b/tests/services/test_example_images_other_dispatch.py new file mode 100644 index 00000000..15744eae --- /dev/null +++ b/tests/services/test_example_images_other_dispatch.py @@ -0,0 +1,168 @@ +"""Example-images download dispatch accepts the "other" model type. + +Covers the three scanner-dispatch sites from docs/plans/other-models-page.md +§9.1: check_pending_models, _download_all_example_images and +_download_specific_models_example_images_sync. +""" +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from py.services.settings_manager import get_settings_manager +from py.utils import example_images_download_manager as download_module + + +class StubScanner: + """Scanner double returning predetermined cache contents.""" + + def __init__(self, models: list[dict[str, Any]]) -> None: + self._cache = SimpleNamespace(raw_data=models) + + async def get_cached_data(self): + return self._cache + + +class RecordingWebSocketManager: + def __init__(self) -> None: + self.payloads: list[dict[str, Any]] = [] + + async def broadcast(self, payload: dict[str, Any]) -> None: + self.payloads.append(payload) + + +def _patch_all_scanners(monkeypatch: pytest.MonkeyPatch, **scanners) -> None: + for name, getter in ( + ("lora", "get_lora_scanner"), + ("checkpoint", "get_checkpoint_scanner"), + ("embedding", "get_embedding_scanner"), + ("other", "get_other_scanner"), + ): + scanner = scanners.get(name) or StubScanner([]) + + async def _get_scanner(cls, _scanner=scanner): + return _scanner + + monkeypatch.setattr( + download_module.ServiceRegistry, + getter, + classmethod(_get_scanner), + ) + + +@pytest.mark.asyncio +async def test_check_pending_models_includes_other_scanner( + monkeypatch: pytest.MonkeyPatch, + tmp_path, + settings_manager, +): + ws_manager = RecordingWebSocketManager() + manager = download_module.DownloadManager(ws_manager=ws_manager) + + monkeypatch.setitem(settings_manager.settings, "example_images_path", str(tmp_path)) + + other_models = [{"sha256": "d" * 64, "model_name": "VAE Model"}] + _patch_all_scanners(monkeypatch, other=StubScanner(other_models)) + + result = await manager.check_pending_models(["other"]) + + assert result["success"] is True + assert result["total_models"] == 1 + assert result["pending_count"] == 1 + + +@pytest.mark.asyncio +async def test_download_all_example_images_processes_other_models( + monkeypatch: pytest.MonkeyPatch, + tmp_path, + settings_manager, +): + ws_manager = RecordingWebSocketManager() + manager = download_module.DownloadManager(ws_manager=ws_manager) + + monkeypatch.setitem(settings_manager.settings, "example_images_path", str(tmp_path)) + + other_models = [{"sha256": "e" * 64, "model_name": "Upscaler Model"}] + _patch_all_scanners(monkeypatch, other=StubScanner(other_models)) + + async def fake_get_downloader(): + return object() + + processed: list[tuple[str, dict[str, Any]]] = [] + + async def fake_process_model(self, scanner_type, model, scanner, *_args, **_kwargs): + processed.append((scanner_type, model)) + return False + + monkeypatch.setattr(download_module, "get_downloader", fake_get_downloader) + monkeypatch.setattr( + download_module.DownloadManager, "_process_model", fake_process_model + ) + + # Simulate the running state that start_download establishes. + manager._progress["status"] = "running" + + await manager._download_all_example_images( + str(tmp_path), + optimize=False, + model_types=["other"], + delay=0, + library_name="default", + ) + + assert processed == [("other", other_models[0])] + + +@pytest.mark.asyncio +async def test_download_specific_models_example_images_processes_other_models( + monkeypatch: pytest.MonkeyPatch, + tmp_path, + settings_manager, +): + ws_manager = RecordingWebSocketManager() + manager = download_module.DownloadManager(ws_manager=ws_manager) + + monkeypatch.setitem(settings_manager.settings, "example_images_path", str(tmp_path)) + + model_hash = "f" * 64 + other_models = [{"sha256": model_hash, "model_name": "Text Encoder Model"}] + _patch_all_scanners(monkeypatch, other=StubScanner(other_models)) + + async def fake_get_downloader(): + return object() + + processed: list[tuple[str, dict[str, Any]]] = [] + + async def fake_process_specific_model( + self, scanner_type, model, scanner, *_args, **_kwargs + ): + processed.append((scanner_type, model)) + return True + + monkeypatch.setattr(download_module, "get_downloader", fake_get_downloader) + monkeypatch.setattr( + download_module.DownloadManager, + "_process_specific_model", + fake_process_specific_model, + ) + + # Simulate the running state that start_force_download establishes. + manager._progress["status"] = "running" + + await manager._download_specific_models_example_images_sync( + [model_hash], + str(tmp_path), + optimize=False, + model_types=["other"], + delay=0, + library_name="default", + ) + + assert processed == [("other", other_models[0])] + + +@pytest.fixture +def settings_manager(): + return get_settings_manager() diff --git a/tests/services/test_other_scanner.py b/tests/services/test_other_scanner.py index 63c4dc3a..14882acd 100644 --- a/tests/services/test_other_scanner.py +++ b/tests/services/test_other_scanner.py @@ -303,9 +303,13 @@ class TestOtherModelMetadataFromCivitai: def _build(self, civitai_type: str) -> OtherModelMetadata: return OtherModelMetadata.from_civitai_info( { - "type": civitai_type, "baseModel": "SDXL", - "model": {"name": "Model", "tags": ["tag"], "description": "desc"}, + "model": { + "name": "Model", + "tags": ["tag"], + "description": "desc", + "type": civitai_type, + }, }, {"name": "model.safetensors", "sizeKB": 1, "hashes": {"SHA256": "AB"}}, "/tmp/model.safetensors", @@ -329,6 +333,20 @@ class TestOtherModelMetadataFromCivitai: assert metadata.sha256 == "ab" assert metadata.tags == ["tag"] + def test_top_level_type_key_is_ignored(self): + """Regression: the CivitAI type lives at version["model"]["type"]; a + top-level version["type"] key must not drive the mapping (#Phase-1 bug).""" + metadata = OtherModelMetadata.from_civitai_info( + { + "type": "Upscaler", + "baseModel": "SDXL", + "model": {"name": "Model", "type": "VAE"}, + }, + {"name": "model.safetensors", "sizeKB": 1, "hashes": {"SHA256": "AB"}}, + "/tmp/model.safetensors", + ) + assert metadata.sub_type == "vae" + def test_page_type_maps_to_other(): """The WS progress page type for the other scanner is 'other'.""" diff --git a/tests/services/test_settings_manager.py b/tests/services/test_settings_manager.py index 980de6e8..9d4dc2be 100644 --- a/tests/services/test_settings_manager.py +++ b/tests/services/test_settings_manager.py @@ -1208,3 +1208,137 @@ def test_skip_previously_downloaded_model_versions_coerces_string_input(manager) assert manager.get_skip_previously_downloaded_model_versions() is True assert manager.settings["skip_previously_downloaded_model_versions"] is True + + +def test_default_other_roots_stay_empty_without_other_folders(manager): + assert manager._get_default_settings()["default_other_roots"] == {} + + manager.settings["default_other_roots"] = {} + manager.settings["folder_paths"] = {} + manager.settings["extra_folder_paths"] = {} + + manager._auto_set_default_roots() + + assert manager.get("default_other_roots") == {} + + +def test_auto_set_default_other_roots(manager): + manager.settings["default_other_roots"] = {} + manager.settings["folder_paths"] = { + "vae": ["/vae"], + "upscale_models": ["/upscalers"], + "clip_vision": ["/clip_vision"], + } + + manager._auto_set_default_roots() + + roots = manager.get("default_other_roots") + assert roots["vae"] == "/vae" + assert roots["upscaler"] == "/upscalers" + assert roots["clip_vision"] == "/clip_vision" + # text_encoder has no configured folders -> no entry + assert "text_encoder" not in roots + assert "controlnet" not in roots + + +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["default_other_roots"] = {} + manager.settings["folder_paths"] = { + "clip": ["/legacy-clip"], + "text_encoders": ["/text-encoders"], + } + + manager._auto_set_default_roots() + + roots = manager.get("default_other_roots") + assert roots["text_encoder"] in {"/legacy-clip", "/text-encoders"} + + # A value pointing at either key's root is considered valid + manager.settings["default_other_roots"] = {"text_encoder": "/legacy-clip"} + manager._auto_set_default_roots() + assert manager.get("default_other_roots")["text_encoder"] == "/legacy-clip" + + +def test_auto_set_default_other_roots_repairs_stale(manager): + manager.settings["default_other_roots"] = {"vae": "/stale-vae"} + manager.settings["folder_paths"] = {"vae": ["/vae"]} + + manager._auto_set_default_roots() + + assert manager.get("default_other_roots")["vae"] == "/vae" + + +def test_auto_set_default_other_roots_uses_extra_folder_paths(manager): + manager.settings["default_other_roots"] = {} + manager.settings["folder_paths"] = {"vae": []} + manager.settings["extra_folder_paths"] = {"vae": ["/extra-vae"]} + + manager._auto_set_default_roots() + + assert manager.get("default_other_roots")["vae"] == "/extra-vae" + + +def test_set_default_other_roots_syncs_active_library(manager): + manager.set("default_other_roots", {"vae": "/vae"}) + + libraries = manager.get_libraries() + active = manager.get_active_library_name() + assert libraries[active]["default_other_roots"] == {"vae": "/vae"} + assert manager.get("default_other_roots") == {"vae": "/vae"} + + +def test_set_default_other_roots_rejects_illegal_sub_type(manager): + with pytest.raises(ValueError, match="Unknown other-model sub-type"): + manager.set("default_other_roots", {"vae": "/vae", "lora": "/loras"}) + + +def test_set_default_other_roots_normalizes_values(manager): + manager.set("default_other_roots", {"vae": " /vae ", "upscaler": ""}) + assert manager.get("default_other_roots") == {"vae": "/vae"} + + +def test_upsert_library_passthrough_default_other_roots(manager, tmp_path): + manager.upsert_library( + "studio", + folder_paths={"loras": ["/studio/loras"], "vae": ["/studio/vae"]}, + default_other_roots={"vae": "/studio/vae"}, + activate=True, + ) + + libraries = manager.get_libraries() + assert libraries["studio"]["default_other_roots"] == {"vae": "/studio/vae"} + assert manager.get("default_other_roots") == {"vae": "/studio/vae"} + + # Omitting the argument preserves the stored value + manager.upsert_library("studio", folder_paths={"loras": ["/studio/loras"]}) + libraries = manager.get_libraries() + assert libraries["studio"]["default_other_roots"] == {"vae": "/studio/vae"} + + +def test_library_switch_restores_default_other_roots(manager): + manager.set("default_other_roots", {"vae": "/default-vae"}) + manager.create_library( + "studio", + folder_paths={"loras": ["/studio/loras"]}, + default_other_roots={"vae": "/studio-vae"}, + ) + + manager.activate_library("studio") + assert manager.get("default_other_roots") == {"vae": "/studio-vae"} + + manager.activate_library("default") + assert manager.get("default_other_roots") == {"vae": "/default-vae"} + + +def test_migrate_sanitizes_legacy_libraries_includes_other_roots(tmp_path, monkeypatch): + initial = { + "libraries": {"legacy": "not-a-dict"}, + "active_library": "legacy", + "folder_paths": {"loras": ["/old"]}, + } + + manager = _create_manager_with_settings(tmp_path, monkeypatch, initial) + + payload = manager.get_libraries()["legacy"] + assert payload["default_other_roots"] == {}