mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 11:11:26 -03:00
feat(backend): CivitAI download support for other model types with subtype routing
This commit is contained in:
@@ -7,7 +7,11 @@ import logging
|
|||||||
|
|
||||||
from aiohttp import web
|
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__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -31,6 +35,7 @@ class DownloadRoutingHandler:
|
|||||||
model_type = payload.get("model_type", "")
|
model_type = payload.get("model_type", "")
|
||||||
base_model = payload.get("base_model") or ""
|
base_model = payload.get("base_model") or ""
|
||||||
file_types = payload.get("file_types") 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:
|
if not isinstance(model_type, str) or not model_type:
|
||||||
return web.json_response(
|
return web.json_response(
|
||||||
@@ -44,6 +49,25 @@ class DownloadRoutingHandler:
|
|||||||
},
|
},
|
||||||
status=400,
|
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(
|
is_diffusion = is_diffusion_model_download(
|
||||||
model_type,
|
model_type,
|
||||||
|
|||||||
@@ -53,6 +53,7 @@ from ...utils.constants import (
|
|||||||
PREVIEW_EXTENSIONS,
|
PREVIEW_EXTENSIONS,
|
||||||
SUPPORTED_MEDIA_EXTENSIONS,
|
SUPPORTED_MEDIA_EXTENSIONS,
|
||||||
VALID_LORA_TYPES,
|
VALID_LORA_TYPES,
|
||||||
|
VALID_OTHER_CIVITAI_TYPES,
|
||||||
)
|
)
|
||||||
from .hf_handlers import HfHandler
|
from .hf_handlers import HfHandler
|
||||||
from .agent_handlers import AgentHandler
|
from .agent_handlers import AgentHandler
|
||||||
@@ -2068,6 +2069,7 @@ class ServiceRegistryAdapter:
|
|||||||
get_embedding_scanner: Callable[[], Awaitable[Any]]
|
get_embedding_scanner: Callable[[], Awaitable[Any]]
|
||||||
get_downloaded_version_history_service: Callable[[], Awaitable[Any]]
|
get_downloaded_version_history_service: Callable[[], Awaitable[Any]]
|
||||||
get_backup_service: Callable[[], Awaitable[Any]] = _noop_backup_service
|
get_backup_service: Callable[[], Awaitable[Any]] = _noop_backup_service
|
||||||
|
get_other_scanner: Callable[[], Awaitable[Any]] = ServiceRegistry.get_other_scanner
|
||||||
|
|
||||||
|
|
||||||
class ModelLibraryHandler:
|
class ModelLibraryHandler:
|
||||||
@@ -2789,12 +2791,30 @@ class ModelLibraryHandler:
|
|||||||
model_type.lower() for model_type in CIVITAI_USER_MODEL_TYPES
|
model_type.lower() for model_type in CIVITAI_USER_MODEL_TYPES
|
||||||
}
|
}
|
||||||
lora_type_aliases = {model_type.lower() for model_type in VALID_LORA_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] = {
|
type_scanner_map: Dict[str, Any] = {
|
||||||
**{alias: lora_scanner for alias in lora_type_aliases},
|
**{alias: lora_scanner for alias in lora_type_aliases},
|
||||||
"checkpoint": checkpoint_scanner,
|
"checkpoint": checkpoint_scanner,
|
||||||
"textualinversion": embedding_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]] = []
|
versions: list[dict[str, Any]] = []
|
||||||
history_service = await self._get_download_history_service()
|
history_service = await self._get_download_history_service()
|
||||||
@@ -2818,12 +2838,17 @@ class ModelLibraryHandler:
|
|||||||
"embedding",
|
"embedding",
|
||||||
model_ids,
|
model_ids,
|
||||||
)
|
)
|
||||||
|
other_downloaded = await history_service.get_downloaded_version_ids_bulk(
|
||||||
|
"other",
|
||||||
|
model_ids,
|
||||||
|
)
|
||||||
downloaded_version_map: Dict[str, Dict[int, set[int]]] = {
|
downloaded_version_map: Dict[str, Dict[int, set[int]]] = {
|
||||||
"lora": lora_downloaded,
|
"lora": lora_downloaded,
|
||||||
"locon": lora_downloaded,
|
"locon": lora_downloaded,
|
||||||
"dora": lora_downloaded,
|
"dora": lora_downloaded,
|
||||||
"checkpoint": checkpoint_downloaded,
|
"checkpoint": checkpoint_downloaded,
|
||||||
"textualinversion": embedding_downloaded,
|
"textualinversion": embedding_downloaded,
|
||||||
|
**{alias: other_downloaded for alias in VALID_OTHER_CIVITAI_TYPES},
|
||||||
}
|
}
|
||||||
for model in models:
|
for model in models:
|
||||||
if not isinstance(model, dict):
|
if not isinstance(model, dict):
|
||||||
@@ -3982,6 +4007,7 @@ def build_service_registry_adapter() -> ServiceRegistryAdapter:
|
|||||||
get_lora_scanner=ServiceRegistry.get_lora_scanner,
|
get_lora_scanner=ServiceRegistry.get_lora_scanner,
|
||||||
get_checkpoint_scanner=ServiceRegistry.get_checkpoint_scanner,
|
get_checkpoint_scanner=ServiceRegistry.get_checkpoint_scanner,
|
||||||
get_embedding_scanner=ServiceRegistry.get_embedding_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_downloaded_version_history_service=ServiceRegistry.get_downloaded_version_history_service,
|
||||||
get_backup_service=ServiceRegistry.get_backup_service,
|
get_backup_service=ServiceRegistry.get_backup_service,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -1,12 +1,13 @@
|
|||||||
import logging
|
import logging
|
||||||
from typing import Any, Dict
|
from typing import Any, Dict, List
|
||||||
from aiohttp import web
|
from aiohttp import web
|
||||||
|
|
||||||
from .base_model_routes import BaseModelRoutes
|
from .base_model_routes import BaseModelRoutes
|
||||||
from .model_route_registrar import ModelRouteRegistrar
|
from .model_route_registrar import ModelRouteRegistrar
|
||||||
|
from ..config import config
|
||||||
from ..services.other_model_service import OtherModelService
|
from ..services.other_model_service import OtherModelService
|
||||||
from ..services.service_registry import ServiceRegistry
|
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__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -41,6 +42,9 @@ class OtherRoutes(BaseModelRoutes):
|
|||||||
"""Setup Other-model-specific routes"""
|
"""Setup Other-model-specific routes"""
|
||||||
# Other-model info by name
|
# Other-model info by name
|
||||||
registrar.add_prefixed_route('GET', '/api/lm/{prefix}/info/{name}', prefix, self.get_other_model_info)
|
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:
|
def _validate_civitai_model_type(self, model_type: str) -> bool:
|
||||||
"""Validate CivitAI model type for other models.
|
"""Validate CivitAI model type for other models.
|
||||||
@@ -58,6 +62,32 @@ class OtherRoutes(BaseModelRoutes):
|
|||||||
"""Parse other-model-specific parameters (none in Phase 1)."""
|
"""Parse other-model-specific parameters (none in Phase 1)."""
|
||||||
return {}
|
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:
|
async def get_other_model_info(self, request: web.Request) -> web.Response:
|
||||||
"""Get detailed information for a specific other model by name"""
|
"""Get detailed information for a specific other model by name"""
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -17,12 +17,18 @@ from dataclasses import dataclass, field
|
|||||||
import uuid
|
import uuid
|
||||||
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, cast
|
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, cast
|
||||||
from urllib.parse import urlparse
|
from urllib.parse import urlparse
|
||||||
from ..utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata
|
from ..utils.models import (
|
||||||
|
LoraMetadata,
|
||||||
|
CheckpointMetadata,
|
||||||
|
EmbeddingMetadata,
|
||||||
|
OtherModelMetadata,
|
||||||
|
)
|
||||||
from ..utils.constants import (
|
from ..utils.constants import (
|
||||||
CARD_PREVIEW_WIDTH,
|
CARD_PREVIEW_WIDTH,
|
||||||
MODEL_WEIGHT_FILE_TYPES,
|
MODEL_WEIGHT_FILE_TYPES,
|
||||||
SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS,
|
SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS,
|
||||||
VALID_LORA_TYPES,
|
VALID_LORA_TYPES,
|
||||||
|
VALID_OTHER_CIVITAI_TYPES,
|
||||||
)
|
)
|
||||||
from ..utils.civitai_utils import normalize_civitai_download_url, rewrite_preview_url
|
from ..utils.civitai_utils import normalize_civitai_download_url, rewrite_preview_url
|
||||||
from ..utils.file_utils import calculate_sha256, calculate_autov3
|
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.exif_utils import ExifUtils
|
||||||
from ..utils.metadata_manager import MetadataManager
|
from ..utils.metadata_manager import MetadataManager
|
||||||
from .service_registry import ServiceRegistry
|
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 .settings_manager import get_settings_manager
|
||||||
from .metadata_service import get_default_metadata_provider, get_metadata_provider
|
from .metadata_service import get_default_metadata_provider, get_metadata_provider
|
||||||
from .downloader import get_downloader, DownloadProgress, DownloadStreamControl
|
from .downloader import get_downloader, DownloadProgress, DownloadStreamControl
|
||||||
@@ -228,12 +234,21 @@ class DownloadManager:
|
|||||||
return False
|
return False
|
||||||
|
|
||||||
async def _get_scanner_for_model_type(self, model_type: str):
|
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":
|
if model_type == "checkpoint":
|
||||||
return await self._get_checkpoint_scanner()
|
return await self._get_checkpoint_scanner()
|
||||||
if model_type == "embedding":
|
if model_type == "embedding":
|
||||||
return await ServiceRegistry.get_embedding_scanner()
|
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
|
@staticmethod
|
||||||
def _resolve_target_file(
|
def _resolve_target_file(
|
||||||
@@ -978,6 +993,8 @@ class DownloadManager:
|
|||||||
return CheckpointMetadata.from_civitai_info(version_info, file_info, save_path)
|
return CheckpointMetadata.from_civitai_info(version_info, file_info, save_path)
|
||||||
if model_type == "embedding":
|
if model_type == "embedding":
|
||||||
return EmbeddingMetadata.from_civitai_info(version_info, file_info, save_path)
|
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)
|
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]:
|
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()
|
lora_scanner = await self._get_lora_scanner()
|
||||||
checkpoint_scanner = await self._get_checkpoint_scanner()
|
checkpoint_scanner = await self._get_checkpoint_scanner()
|
||||||
embedding_scanner = await ServiceRegistry.get_embedding_scanner()
|
embedding_scanner = await ServiceRegistry.get_embedding_scanner()
|
||||||
|
other_scanner = await ServiceRegistry.get_other_scanner()
|
||||||
|
|
||||||
# Check lora scanner first
|
# Check lora scanner first
|
||||||
if await lora_scanner.check_model_version_exists(model_version_id):
|
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",
|
"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'
|
# Use CivArchive provider directly when source is 'civarchive'
|
||||||
# This prioritizes CivArchive metadata (with mirror availability info) over Civitai
|
# This prioritizes CivArchive metadata (with mirror availability info) over Civitai
|
||||||
if source == "civarchive":
|
if source == "civarchive":
|
||||||
@@ -1500,6 +1525,8 @@ class DownloadManager:
|
|||||||
model_type = "lora"
|
model_type = "lora"
|
||||||
elif model_type_from_info == "textualinversion":
|
elif model_type_from_info == "textualinversion":
|
||||||
model_type = "embedding"
|
model_type = "embedding"
|
||||||
|
elif model_type_from_info in VALID_OTHER_CIVITAI_TYPES:
|
||||||
|
model_type = "other"
|
||||||
else:
|
else:
|
||||||
return {
|
return {
|
||||||
"success": False,
|
"success": False,
|
||||||
@@ -1686,6 +1713,13 @@ class DownloadManager:
|
|||||||
"success": False,
|
"success": False,
|
||||||
"error": "Model version already exists in embedding library",
|
"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
|
# Handle use_default_paths
|
||||||
if use_default_paths:
|
if use_default_paths:
|
||||||
@@ -1725,6 +1759,45 @@ class DownloadManager:
|
|||||||
"error": "Default embedding root path not set in settings",
|
"error": "Default embedding root path not set in settings",
|
||||||
}
|
}
|
||||||
save_dir = default_path
|
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
|
# Calculate relative path using template
|
||||||
relative_path = self._calculate_relative_path(version_info, model_type)
|
relative_path = self._calculate_relative_path(version_info, model_type)
|
||||||
@@ -1921,6 +1994,11 @@ class DownloadManager:
|
|||||||
version_info, file_info, save_path
|
version_info, file_info, save_path
|
||||||
)
|
)
|
||||||
logger.info(f"Creating EmbeddingMetadata for {file_name}")
|
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:
|
else:
|
||||||
return {
|
return {
|
||||||
"success": False,
|
"success": False,
|
||||||
@@ -2133,6 +2211,8 @@ class DownloadManager:
|
|||||||
scanner = await self._get_checkpoint_scanner()
|
scanner = await self._get_checkpoint_scanner()
|
||||||
elif model_type == "embedding":
|
elif model_type == "embedding":
|
||||||
scanner = await ServiceRegistry.get_embedding_scanner()
|
scanner = await ServiceRegistry.get_embedding_scanner()
|
||||||
|
elif model_type == "other":
|
||||||
|
scanner = await ServiceRegistry.get_other_scanner()
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
logger.debug("Failed to acquire scanner for %s models: %s", model_type, exc)
|
logger.debug("Failed to acquire scanner for %s models: %s", model_type, exc)
|
||||||
|
|
||||||
@@ -2629,6 +2709,9 @@ class DownloadManager:
|
|||||||
elif model_type == "embedding":
|
elif model_type == "embedding":
|
||||||
scanner = await ServiceRegistry.get_embedding_scanner()
|
scanner = await ServiceRegistry.get_embedding_scanner()
|
||||||
logger.info(f"Updating embedding cache for {actual_file_paths[0]}")
|
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 = (
|
adjust_cached_entry = (
|
||||||
getattr(scanner, "adjust_cached_entry", None)
|
getattr(scanner, "adjust_cached_entry", None)
|
||||||
@@ -2718,7 +2801,7 @@ class DownloadManager:
|
|||||||
return {"success": False, "error": str(e)}
|
return {"success": False, "error": str(e)}
|
||||||
|
|
||||||
def _get_supported_extensions_for_type(self, model_type: str) -> Set[str]:
|
def _get_supported_extensions_for_type(self, model_type: str) -> Set[str]:
|
||||||
if model_type == "checkpoint":
|
if model_type in ("checkpoint", "other"):
|
||||||
return {
|
return {
|
||||||
".ckpt",
|
".ckpt",
|
||||||
".pt",
|
".pt",
|
||||||
|
|||||||
@@ -10,9 +10,13 @@ two can never disagree.
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import logging
|
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__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -51,3 +55,49 @@ def is_diffusion_model_download(
|
|||||||
return True
|
return True
|
||||||
|
|
||||||
return False
|
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
|
||||||
|
|||||||
@@ -27,7 +27,9 @@ from platformdirs import user_config_dir
|
|||||||
from ..utils.constants import (
|
from ..utils.constants import (
|
||||||
DEFAULT_HASH_CHUNK_SIZE_MB,
|
DEFAULT_HASH_CHUNK_SIZE_MB,
|
||||||
DEFAULT_PRIORITY_TAG_CONFIG,
|
DEFAULT_PRIORITY_TAG_CONFIG,
|
||||||
|
OTHER_MODEL_FOLDER_SUBTYPES,
|
||||||
SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS,
|
SUPPORTED_DOWNLOAD_SKIP_BASE_MODELS,
|
||||||
|
VALID_OTHER_SUB_TYPES,
|
||||||
)
|
)
|
||||||
from ..utils.preview_selection import VALID_MATURE_BLUR_LEVELS
|
from ..utils.preview_selection import VALID_MATURE_BLUR_LEVELS
|
||||||
from ..utils.settings_paths import (
|
from ..utils.settings_paths import (
|
||||||
@@ -83,6 +85,7 @@ DEFAULT_SETTINGS: Dict[str, Any] = {
|
|||||||
"default_checkpoint_root": "",
|
"default_checkpoint_root": "",
|
||||||
"default_unet_root": "",
|
"default_unet_root": "",
|
||||||
"default_embedding_root": "",
|
"default_embedding_root": "",
|
||||||
|
"default_other_roots": {},
|
||||||
"recipes_path": "",
|
"recipes_path": "",
|
||||||
"base_model_path_mappings": {},
|
"base_model_path_mappings": {},
|
||||||
"download_path_templates": {},
|
"download_path_templates": {},
|
||||||
@@ -309,6 +312,7 @@ class SettingsManager:
|
|||||||
default_checkpoint_root=merged.get("default_checkpoint_root"),
|
default_checkpoint_root=merged.get("default_checkpoint_root"),
|
||||||
default_unet_root=merged.get("default_unet_root"),
|
default_unet_root=merged.get("default_unet_root"),
|
||||||
default_embedding_root=merged.get("default_embedding_root"),
|
default_embedding_root=merged.get("default_embedding_root"),
|
||||||
|
default_other_roots=merged.get("default_other_roots"),
|
||||||
recipes_path=merged.get("recipes_path"),
|
recipes_path=merged.get("recipes_path"),
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
@@ -443,6 +447,7 @@ class SettingsManager:
|
|||||||
),
|
),
|
||||||
default_unet_root=self.settings.get("default_unet_root", ""),
|
default_unet_root=self.settings.get("default_unet_root", ""),
|
||||||
default_embedding_root=self.settings.get("default_embedding_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", ""),
|
recipes_path=self.settings.get("recipes_path", ""),
|
||||||
)
|
)
|
||||||
libraries = {library_name: library_payload}
|
libraries = {library_name: library_payload}
|
||||||
@@ -494,6 +499,7 @@ class SettingsManager:
|
|||||||
default_checkpoint_root=data.get("default_checkpoint_root"),
|
default_checkpoint_root=data.get("default_checkpoint_root"),
|
||||||
default_unet_root=data.get("default_unet_root"),
|
default_unet_root=data.get("default_unet_root"),
|
||||||
default_embedding_root=data.get("default_embedding_root"),
|
default_embedding_root=data.get("default_embedding_root"),
|
||||||
|
default_other_roots=data.get("default_other_roots"),
|
||||||
recipes_path=data.get("recipes_path"),
|
recipes_path=data.get("recipes_path"),
|
||||||
metadata=data.get("metadata"),
|
metadata=data.get("metadata"),
|
||||||
base=data,
|
base=data,
|
||||||
@@ -541,6 +547,9 @@ class SettingsManager:
|
|||||||
self.settings["default_embedding_root"] = active_library.get(
|
self.settings["default_embedding_root"] = active_library.get(
|
||||||
"default_embedding_root", ""
|
"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", "")
|
self.settings["recipes_path"] = active_library.get("recipes_path", "")
|
||||||
|
|
||||||
if save:
|
if save:
|
||||||
@@ -558,6 +567,7 @@ class SettingsManager:
|
|||||||
default_checkpoint_root: Optional[str] = None,
|
default_checkpoint_root: Optional[str] = None,
|
||||||
default_unet_root: Optional[str] = None,
|
default_unet_root: Optional[str] = None,
|
||||||
default_embedding_root: Optional[str] = None,
|
default_embedding_root: Optional[str] = None,
|
||||||
|
default_other_roots: Optional[Mapping[str, str]] = None,
|
||||||
recipes_path: Optional[str] = None,
|
recipes_path: Optional[str] = None,
|
||||||
metadata: Optional[Mapping[str, Any]] = None,
|
metadata: Optional[Mapping[str, Any]] = None,
|
||||||
base: Optional[Mapping[str, Any]] = None,
|
base: Optional[Mapping[str, Any]] = None,
|
||||||
@@ -597,6 +607,15 @@ class SettingsManager:
|
|||||||
else:
|
else:
|
||||||
payload.setdefault("default_embedding_root", "")
|
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:
|
if recipes_path is not None:
|
||||||
payload["recipes_path"] = recipes_path
|
payload["recipes_path"] = recipes_path
|
||||||
else:
|
else:
|
||||||
@@ -632,6 +651,35 @@ class SettingsManager:
|
|||||||
normalized[key] = cleaned
|
normalized[key] = cleaned
|
||||||
return normalized
|
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:
|
def _has_configured_paths(self, folder_paths: Any) -> bool:
|
||||||
if not isinstance(folder_paths, Mapping):
|
if not isinstance(folder_paths, Mapping):
|
||||||
return False
|
return False
|
||||||
@@ -744,6 +792,7 @@ class SettingsManager:
|
|||||||
default_checkpoint_root: Optional[str] = None,
|
default_checkpoint_root: Optional[str] = None,
|
||||||
default_unet_root: Optional[str] = None,
|
default_unet_root: Optional[str] = None,
|
||||||
default_embedding_root: Optional[str] = None,
|
default_embedding_root: Optional[str] = None,
|
||||||
|
default_other_roots: Optional[Mapping[str, str]] = None,
|
||||||
recipes_path: Optional[str] = None,
|
recipes_path: Optional[str] = None,
|
||||||
) -> bool:
|
) -> bool:
|
||||||
libraries = self.settings.get("libraries", {})
|
libraries = self.settings.get("libraries", {})
|
||||||
@@ -794,6 +843,14 @@ class SettingsManager:
|
|||||||
library["default_embedding_root"] = default_embedding_root
|
library["default_embedding_root"] = default_embedding_root
|
||||||
changed = True
|
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:
|
if recipes_path is not None and library.get("recipes_path") != recipes_path:
|
||||||
library["recipes_path"] = recipes_path
|
library["recipes_path"] = recipes_path
|
||||||
changed = True
|
changed = True
|
||||||
@@ -894,12 +951,53 @@ class SettingsManager:
|
|||||||
updated = _check_and_auto_set("unet", "default_unet_root") or updated
|
updated = _check_and_auto_set("unet", "default_unet_root") or updated
|
||||||
updated = _check_and_auto_set("embeddings", "default_embedding_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:
|
if updated:
|
||||||
|
self.settings["default_other_roots"] = other_roots
|
||||||
self._update_active_library_entry(
|
self._update_active_library_entry(
|
||||||
default_lora_root=self.settings.get("default_lora_root"),
|
default_lora_root=self.settings.get("default_lora_root"),
|
||||||
default_checkpoint_root=self.settings.get("default_checkpoint_root"),
|
default_checkpoint_root=self.settings.get("default_checkpoint_root"),
|
||||||
default_unet_root=self.settings.get("default_unet_root"),
|
default_unet_root=self.settings.get("default_unet_root"),
|
||||||
default_embedding_root=self.settings.get("default_embedding_root"),
|
default_embedding_root=self.settings.get("default_embedding_root"),
|
||||||
|
default_other_roots=other_roots,
|
||||||
)
|
)
|
||||||
if self._bootstrap_reason == "missing":
|
if self._bootstrap_reason == "missing":
|
||||||
self._needs_initial_save = True
|
self._needs_initial_save = True
|
||||||
@@ -1599,6 +1697,8 @@ class SettingsManager:
|
|||||||
value = self.normalize_download_skip_base_models(value)
|
value = self.normalize_download_skip_base_models(value)
|
||||||
elif key == "mature_blur_level":
|
elif key == "mature_blur_level":
|
||||||
value = self.normalize_mature_blur_level(value)
|
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":
|
elif key == "recipes_path":
|
||||||
current_recipes_dir = self._get_effective_recipes_dir()
|
current_recipes_dir = self._get_effective_recipes_dir()
|
||||||
value = self._normalize_recipes_path_value(value)
|
value = self._normalize_recipes_path_value(value)
|
||||||
@@ -1626,6 +1726,8 @@ class SettingsManager:
|
|||||||
self._update_active_library_entry(default_unet_root=str(value))
|
self._update_active_library_entry(default_unet_root=str(value))
|
||||||
elif key == "default_embedding_root":
|
elif key == "default_embedding_root":
|
||||||
self._update_active_library_entry(default_embedding_root=str(value))
|
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":
|
elif key == "recipes_path":
|
||||||
self._update_active_library_entry(recipes_path=str(value))
|
self._update_active_library_entry(recipes_path=str(value))
|
||||||
elif key == "model_name_display":
|
elif key == "model_name_display":
|
||||||
@@ -1796,6 +1898,7 @@ class SettingsManager:
|
|||||||
"lora_scanner",
|
"lora_scanner",
|
||||||
"checkpoint_scanner",
|
"checkpoint_scanner",
|
||||||
"embedding_scanner",
|
"embedding_scanner",
|
||||||
|
"other_scanner",
|
||||||
"recipe_scanner",
|
"recipe_scanner",
|
||||||
):
|
):
|
||||||
service = ServiceRegistry.get_service_sync(service_name)
|
service = ServiceRegistry.get_service_sync(service_name)
|
||||||
@@ -1960,6 +2063,7 @@ class SettingsManager:
|
|||||||
default_checkpoint_root: Optional[str] = None,
|
default_checkpoint_root: Optional[str] = None,
|
||||||
default_unet_root: Optional[str] = None,
|
default_unet_root: Optional[str] = None,
|
||||||
default_embedding_root: Optional[str] = None,
|
default_embedding_root: Optional[str] = None,
|
||||||
|
default_other_roots: Optional[Mapping[str, str]] = None,
|
||||||
recipes_path: Optional[str] = None,
|
recipes_path: Optional[str] = None,
|
||||||
metadata: Optional[Mapping[str, Any]] = None,
|
metadata: Optional[Mapping[str, Any]] = None,
|
||||||
activate: bool = False,
|
activate: bool = False,
|
||||||
@@ -2004,6 +2108,11 @@ class SettingsManager:
|
|||||||
if default_embedding_root is not None
|
if default_embedding_root is not None
|
||||||
else existing.get("default_embedding_root")
|
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=(
|
||||||
recipes_path
|
recipes_path
|
||||||
if recipes_path is not None
|
if recipes_path is not None
|
||||||
@@ -2036,6 +2145,7 @@ class SettingsManager:
|
|||||||
default_checkpoint_root: str = "",
|
default_checkpoint_root: str = "",
|
||||||
default_unet_root: str = "",
|
default_unet_root: str = "",
|
||||||
default_embedding_root: str = "",
|
default_embedding_root: str = "",
|
||||||
|
default_other_roots: Optional[Mapping[str, str]] = None,
|
||||||
recipes_path: str = "",
|
recipes_path: str = "",
|
||||||
metadata: Optional[Mapping[str, Any]] = None,
|
metadata: Optional[Mapping[str, Any]] = None,
|
||||||
activate: bool = False,
|
activate: bool = False,
|
||||||
@@ -2054,6 +2164,7 @@ class SettingsManager:
|
|||||||
default_checkpoint_root=default_checkpoint_root,
|
default_checkpoint_root=default_checkpoint_root,
|
||||||
default_unet_root=default_unet_root,
|
default_unet_root=default_unet_root,
|
||||||
default_embedding_root=default_embedding_root,
|
default_embedding_root=default_embedding_root,
|
||||||
|
default_other_roots=default_other_roots,
|
||||||
recipes_path=recipes_path,
|
recipes_path=recipes_path,
|
||||||
metadata=metadata,
|
metadata=metadata,
|
||||||
activate=activate,
|
activate=activate,
|
||||||
@@ -2114,6 +2225,7 @@ class SettingsManager:
|
|||||||
default_checkpoint_root: Optional[str] = None,
|
default_checkpoint_root: Optional[str] = None,
|
||||||
default_unet_root: Optional[str] = None,
|
default_unet_root: Optional[str] = None,
|
||||||
default_embedding_root: Optional[str] = None,
|
default_embedding_root: Optional[str] = None,
|
||||||
|
default_other_roots: Optional[Mapping[str, str]] = None,
|
||||||
recipes_path: Optional[str] = None,
|
recipes_path: Optional[str] = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Update folder paths for the active library."""
|
"""Update folder paths for the active library."""
|
||||||
@@ -2127,6 +2239,7 @@ class SettingsManager:
|
|||||||
default_checkpoint_root=default_checkpoint_root,
|
default_checkpoint_root=default_checkpoint_root,
|
||||||
default_unet_root=default_unet_root,
|
default_unet_root=default_unet_root,
|
||||||
default_embedding_root=default_embedding_root,
|
default_embedding_root=default_embedding_root,
|
||||||
|
default_other_roots=default_other_roots,
|
||||||
recipes_path=recipes_path,
|
recipes_path=recipes_path,
|
||||||
activate=True,
|
activate=True,
|
||||||
)
|
)
|
||||||
@@ -2151,6 +2264,7 @@ class SettingsManager:
|
|||||||
"lora_scanner",
|
"lora_scanner",
|
||||||
"checkpoint_scanner",
|
"checkpoint_scanner",
|
||||||
"embedding_scanner",
|
"embedding_scanner",
|
||||||
|
"other_scanner",
|
||||||
"recipe_scanner",
|
"recipe_scanner",
|
||||||
"model_update_service",
|
"model_update_service",
|
||||||
):
|
):
|
||||||
|
|||||||
@@ -126,6 +126,20 @@ CIVITAI_TYPE_TO_OTHER_SUB_TYPE = {
|
|||||||
"controlnet": "controlnet",
|
"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
|
# Backward compatibility alias
|
||||||
VALID_LORA_TYPES = VALID_LORA_SUB_TYPES
|
VALID_LORA_TYPES = VALID_LORA_SUB_TYPES
|
||||||
|
|
||||||
@@ -134,6 +148,7 @@ CIVITAI_USER_MODEL_TYPES = [
|
|||||||
*VALID_LORA_TYPES,
|
*VALID_LORA_TYPES,
|
||||||
"textualinversion",
|
"textualinversion",
|
||||||
"checkpoint",
|
"checkpoint",
|
||||||
|
*sorted(VALID_OTHER_CIVITAI_TYPES),
|
||||||
]
|
]
|
||||||
|
|
||||||
# Default chunk size in megabytes used for hashing large files.
|
# Default chunk size in megabytes used for hashing large files.
|
||||||
|
|||||||
@@ -420,6 +420,10 @@ class DownloadManager:
|
|||||||
embedding_scanner = await ServiceRegistry.get_embedding_scanner()
|
embedding_scanner = await ServiceRegistry.get_embedding_scanner()
|
||||||
scanners.append(("embedding", 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)
|
# Load progress file to check processed models (async to avoid blocking)
|
||||||
settings_manager = get_settings_manager()
|
settings_manager = get_settings_manager()
|
||||||
active_library = settings_manager.get_active_library_name()
|
active_library = settings_manager.get_active_library_name()
|
||||||
@@ -600,6 +604,10 @@ class DownloadManager:
|
|||||||
embedding_scanner = await ServiceRegistry.get_embedding_scanner()
|
embedding_scanner = await ServiceRegistry.get_embedding_scanner()
|
||||||
scanners.append(("embedding", 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
|
# Get all models
|
||||||
all_models = []
|
all_models = []
|
||||||
for scanner_type, scanner in scanners:
|
for scanner_type, scanner in scanners:
|
||||||
@@ -1098,6 +1106,10 @@ class DownloadManager:
|
|||||||
embedding_scanner = await ServiceRegistry.get_embedding_scanner()
|
embedding_scanner = await ServiceRegistry.get_embedding_scanner()
|
||||||
scanners.append(("embedding", 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
|
# Find the specified models
|
||||||
models_to_process = []
|
models_to_process = []
|
||||||
for scanner_type, scanner in scanners:
|
for scanner_type, scanner in scanners:
|
||||||
|
|||||||
+2
-1
@@ -340,7 +340,8 @@ class OtherModelMetadata(BaseModelMetadata):
|
|||||||
sha256_value = (file_info.get("hashes") or {}).get("SHA256", "").lower()
|
sha256_value = (file_info.get("hashes") or {}).get("SHA256", "").lower()
|
||||||
# Map the CivitAI model type onto our sub_types; unknown types keep the
|
# Map the CivitAI model type onto our sub_types; unknown types keep the
|
||||||
# placeholder until the scanner re-derives sub_type from the location.
|
# 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")
|
sub_type = CIVITAI_TYPE_TO_OTHER_SUB_TYPE.get(civitai_type, "vae")
|
||||||
|
|
||||||
# Extract tags and description if available
|
# Extract tags and description if available
|
||||||
|
|||||||
@@ -31,5 +31,6 @@
|
|||||||
"C:/path/to/your/clip_vision_folder"
|
"C:/path/to/your/clip_vision_folder"
|
||||||
]
|
]
|
||||||
},
|
},
|
||||||
|
"default_other_roots": {},
|
||||||
"auto_organize_exclusions": []
|
"auto_organize_exclusions": []
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -91,3 +91,61 @@ async def test_invalid_json_rejected():
|
|||||||
FakeRequest(json.JSONDecodeError("bad", "", 0))
|
FakeRequest(json.JSONDecodeError("bad", "", 0))
|
||||||
)
|
)
|
||||||
assert response.status == 400
|
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
|
||||||
|
|||||||
@@ -1133,8 +1133,8 @@ async def test_get_civitai_user_models_marks_library_versions():
|
|||||||
},
|
},
|
||||||
{
|
{
|
||||||
"id": 4,
|
"id": 4,
|
||||||
"name": "Unsupported",
|
"name": "VAE Model",
|
||||||
"type": "Other",
|
"type": "VAE",
|
||||||
"modelVersions": [
|
"modelVersions": [
|
||||||
{
|
{
|
||||||
"id": 400,
|
"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)
|
provider = FakeUserModelsProvider(models)
|
||||||
@@ -1152,6 +1163,7 @@ async def test_get_civitai_user_models_marks_library_versions():
|
|||||||
lora_scanner = FakeExistenceScanner({101})
|
lora_scanner = FakeExistenceScanner({101})
|
||||||
checkpoint_scanner = FakeExistenceScanner()
|
checkpoint_scanner = FakeExistenceScanner()
|
||||||
embedding_scanner = FakeExistenceScanner({202})
|
embedding_scanner = FakeExistenceScanner({202})
|
||||||
|
other_scanner = FakeExistenceScanner({400})
|
||||||
|
|
||||||
async def lora_factory():
|
async def lora_factory():
|
||||||
return lora_scanner
|
return lora_scanner
|
||||||
@@ -1162,11 +1174,15 @@ async def test_get_civitai_user_models_marks_library_versions():
|
|||||||
async def embedding_factory():
|
async def embedding_factory():
|
||||||
return embedding_scanner
|
return embedding_scanner
|
||||||
|
|
||||||
|
async def other_factory():
|
||||||
|
return other_scanner
|
||||||
|
|
||||||
handler = ModelLibraryHandler(
|
handler = ModelLibraryHandler(
|
||||||
ServiceRegistryAdapter(
|
ServiceRegistryAdapter(
|
||||||
get_lora_scanner=lora_factory,
|
get_lora_scanner=lora_factory,
|
||||||
get_checkpoint_scanner=checkpoint_factory,
|
get_checkpoint_scanner=checkpoint_factory,
|
||||||
get_embedding_scanner=embedding_factory,
|
get_embedding_scanner=embedding_factory,
|
||||||
|
get_other_scanner=other_factory,
|
||||||
get_downloaded_version_history_service=lambda: fake_download_history_service_factory(),
|
get_downloaded_version_history_service=lambda: fake_download_history_service_factory(),
|
||||||
),
|
),
|
||||||
metadata_provider_factory=provider_factory,
|
metadata_provider_factory=provider_factory,
|
||||||
@@ -1240,6 +1256,18 @@ async def test_get_civitai_user_models_marks_library_versions():
|
|||||||
"inLibrary": False,
|
"inLibrary": False,
|
||||||
"hasBeenDownloaded": 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"]
|
assert provider.received_usernames == ["pixel"]
|
||||||
@@ -1351,7 +1379,7 @@ async def test_get_civitai_user_models_returns_pagination_fields():
|
|||||||
{
|
{
|
||||||
"id": 2,
|
"id": 2,
|
||||||
"name": "Unsupported",
|
"name": "Unsupported",
|
||||||
"type": "Other",
|
"type": "Wildcard",
|
||||||
"modelVersions": [{"id": 200, "name": "v1"}],
|
"modelVersions": [{"id": 200, "name": "v1"}],
|
||||||
},
|
},
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -116,3 +116,50 @@ async def test_initialize_services_builds_other_model_service(monkeypatch):
|
|||||||
assert isinstance(handler.service, OtherModelService)
|
assert isinstance(handler.service, OtherModelService)
|
||||||
assert handler.service.model_type == "other"
|
assert handler.service.model_type == "other"
|
||||||
assert handler.service.scanner is sentinel_scanner
|
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": {}}
|
||||||
|
|||||||
@@ -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])]
|
||||||
@@ -38,3 +38,105 @@ def test_non_checkpoint_types_never_route_to_unet():
|
|||||||
def test_empty_inputs_stay_on_checkpoint_roots():
|
def test_empty_inputs_stay_on_checkpoint_roots():
|
||||||
assert not is_diffusion_model_download("checkpoint")
|
assert not is_diffusion_model_download("checkpoint")
|
||||||
assert not is_diffusion_model_download("checkpoint", file_types=[], base_model="")
|
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"
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -303,9 +303,13 @@ class TestOtherModelMetadataFromCivitai:
|
|||||||
def _build(self, civitai_type: str) -> OtherModelMetadata:
|
def _build(self, civitai_type: str) -> OtherModelMetadata:
|
||||||
return OtherModelMetadata.from_civitai_info(
|
return OtherModelMetadata.from_civitai_info(
|
||||||
{
|
{
|
||||||
"type": civitai_type,
|
|
||||||
"baseModel": "SDXL",
|
"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"}},
|
{"name": "model.safetensors", "sizeKB": 1, "hashes": {"SHA256": "AB"}},
|
||||||
"/tmp/model.safetensors",
|
"/tmp/model.safetensors",
|
||||||
@@ -329,6 +333,20 @@ class TestOtherModelMetadataFromCivitai:
|
|||||||
assert metadata.sha256 == "ab"
|
assert metadata.sha256 == "ab"
|
||||||
assert metadata.tags == ["tag"]
|
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():
|
def test_page_type_maps_to_other():
|
||||||
"""The WS progress page type for the other scanner is 'other'."""
|
"""The WS progress page type for the other scanner is 'other'."""
|
||||||
|
|||||||
@@ -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.get_skip_previously_downloaded_model_versions() is True
|
||||||
assert manager.settings["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"] == {}
|
||||||
|
|||||||
Reference in New Issue
Block a user