alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-08-11 08:20:14 -03:00
Code Issues Packages Projects Releases Wiki Activity

fix(types): resolve pre-existing basedpyright errors in py/ and standalone.py

Browse Source 
Fix ~950 basedpyright errors across the backend:
- Convert ineffective # type: ignore comments to # pyright: ignore[rule]
- Add missing generic type arguments (Dict[str, Any], list[Any], ...)
- Annotate dynamic dict literals and runtime-initialized attributes
- Widen CivitAI provider tuple signatures in recipe parsers
- Remove dead LoraRoutes handlers calling nonexistent LoraService methods
- Suppress unavoidable ServiceRegistry import cycles (basedpyright counts
  function-local imports as cycle edges)
This commit is contained in:
Will Miao
2026-08-08 20:12:52 +08:00
parent 6fcdeb799d
commit 8e724538bd
103 changed files with 1184 additions and 1015 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
py/services/agent/agent_service.py
View File

@@ -117,7 +117,7 @@ def _render_prompt(template: str, variables: Dict[str, Any]) -> str:
Uses simple regex substitution — no Jinja2 dependency needed.
"""
def replace(match: re.Match) -> str:
def replace(match: re.Match[str]) -> str:
key = match.group(1).strip()
value = variables.get(key, "")
if isinstance(value, (dict, list)):

2
py/services/agent/post_processor.py
View File

@@ -295,7 +295,7 @@ class PostProcessor:
normalises every tag to lowercase for case-insensitive dedup.
"""
merged: List[str] = []
seen: set = set()
seen: set[str] = set()
for tag in list(existing) + list(new):
t = tag.strip().lower()
if t and t not in seen:

2
py/services/agent/skill_registry.py
View File

@@ -49,7 +49,7 @@ _FRONTMATTER_RE = re.compile(
)
def _parse_skill_file(path: Path) -> tuple[dict, str]:
def _parse_skill_file(path: Path) -> tuple[dict[str, Any], str]:
"""Read a prompt definition file (``prompt.md`` or legacy ``SKILL.md``) and
return (frontmatter_dict, body_text).

36
py/services/agent/skills/enrich_hf_metadata/readme_processor.py
View File

@@ -9,7 +9,7 @@ from __future__ import annotations
import html as html_module
import re
from typing import List, Tuple
from typing import Any, List, Tuple
_REPO_URL_PATTERN = re.compile(r"https?://huggingface\.co/([^/]+/[^/]+)")
@@ -18,10 +18,10 @@ _REPO_URL_PATTERN = re.compile(r"https?://huggingface\.co/([^/]+/[^/]+)")
def extract_simple_markdown_images(
markdown_text: str,
repo: str,
existing_urls: set | None = None,
existing_urls: set[str] | None = None,
default_width: int = 512,
default_height: int = 512,
) -> list[dict]:
) -> list[dict[str, Any]]:
"""Extract standalone markdown images from the README body.
Matches ``![alt](url)`` on lines that are NOT part of a markdown table
@@ -36,8 +36,8 @@ def extract_simple_markdown_images(
return []
base_url = f"https://huggingface.co/{repo}/resolve/main"
images: list[dict] = []
seen_urls: set = set(existing_urls) if existing_urls else set()
images: list[dict[str, Any]] = []
seen_urls: set[str] = set(existing_urls) if existing_urls else set()
# Collect lines that are NOT inside fenced code blocks
lines = markdown_text.split("\n")
@@ -86,10 +86,10 @@ def extract_simple_markdown_images(
def extract_html_img_tags(
markdown_text: str,
repo: str,
existing_urls: set | None = None,
existing_urls: set[str] | None = None,
default_width: int = 512,
default_height: int = 512,
) -> list[dict]:
) -> list[dict[str, Any]]:
"""Extract image URLs from HTML ``<img src=\"...\">`` tags in the README.
Many HF collection repos (e.g. ``deadman44/Z-Image_LoRA``) use raw HTML
@@ -103,8 +103,8 @@ def extract_html_img_tags(
return []
base_url = f"https://huggingface.co/{repo}/resolve/main"
images: list[dict] = []
seen_urls: set = set(existing_urls) if existing_urls else set()
images: list[dict[str, Any]] = []
seen_urls: set[str] = set(existing_urls) if existing_urls else set()
for m in re.finditer(
r'<img\s[^>]*src=\"([^\"]+)\"',
@@ -175,7 +175,7 @@ def extract_gallery_images(
repo: str,
default_width: int = 512,
default_height: int = 512,
) -> List[dict]:
) -> List[dict[str, Any]]:
"""Extract widget/gallery images from the YAML frontmatter of a HF README.
Args:
@@ -196,7 +196,7 @@ def extract_gallery_images(
if not frontmatter:
return []
images: List[dict] = []
images: List[dict[str, Any]] = []
base_url = f"https://huggingface.co/{repo}/resolve/main"
w = default_width or 512
h = default_height or 512
@@ -258,7 +258,7 @@ def extract_gallery_images(
text = raw_text
if url:
image: dict = {
image: dict[str, Any] = {
"url": url,
"type": "image",
"nsfwLevel": 0,
@@ -276,10 +276,10 @@ def extract_gallery_images(
def extract_gallery_table_images(
markdown_text: str,
repo: str,
existing_urls: set | None = None,
existing_urls: set[str] | None = None,
default_width: int = 512,
default_height: int = 512,
) -> list[dict]:
) -> list[dict[str, Any]]:
"""Extract images from ``| Preview | Prompt |`` markdown gallery tables.
Many HF READMEs include a sample-gallery table in the body (outside
@@ -295,8 +295,8 @@ def extract_gallery_table_images(
return []
base_url = f"https://huggingface.co/{repo}/resolve/main"
images: list[dict] = []
seen_urls: set = set(existing_urls) if existing_urls else set()
images: list[dict[str, Any]] = []
seen_urls: set[str] = set(existing_urls) if existing_urls else set()
lines = markdown_text.split("\n")
n = len(lines)
i = 0
@@ -514,7 +514,7 @@ def _strip_standalone_images(text: str) -> str:
URL was stripped entirely, making it impossible for the LLM to return
a ``preview_url`` for repos that use HTML ``<img>`` tags exclusively.
"""
def _img_to_md(match: re.Match) -> str:
def _img_to_md(match: re.Match[str]) -> str:
"""Convert an ``<img>`` tag to markdown image syntax ``![alt](src)``."""
tag = match.group(0)
src_m = re.search(r'src="([^"]+)"', tag) or re.search(r"src='([^']+)'", tag)
@@ -942,7 +942,7 @@ def _strip_badge_images(text: str) -> str:
"twitter", "colab", "gradio", "space",
)
def _should_remove(m: re.Match) -> str:
def _should_remove(m: re.Match[str]) -> str:
alt = (m.group(1) or "").lower()
for kw in badge_keywords:
if kw in alt:

10
py/services/aria2_downloader.py
View File

@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
from __future__ import annotations
import asyncio
@@ -23,7 +27,7 @@ logger = logging.getLogger(__name__)
def _try_certifi_ca_path() -> str | None:
"""Return the certifi CA bundle path if available, else None."""
try:
import certifi # type: ignore[import-untyped]
import certifi # pyright: ignore[reportMissingTypeStubs]
path = certifi.where()
if os.path.isfile(path):
@@ -84,7 +88,7 @@ class Aria2Downloader:
self._transfers: Dict[str, Aria2Transfer] = {}
self._poll_interval = 0.5
self._state_store = Aria2TransferStateStore()
self._stderr_reader_task: Optional[asyncio.Task] = None
self._stderr_reader_task: Optional[asyncio.Task[Any]] = None
@property
def is_running(self) -> bool:
@@ -190,7 +194,7 @@ class Aria2Downloader:
download_id,
)
options: Dict[str, str] = {
options: Dict[str, Any] = {
"dir": save_dir,
"out": out_name,
"continue": "true",

6
py/services/auto_tag_service.py
View File

@@ -8,7 +8,7 @@ from filename, base_model, and CivitAI version name — no manual tagging requir
from __future__ import annotations
import re
from typing import Dict, List, Set
from typing import Any, Dict, List, Set
# ── Tag category definitions ──────────────────────────────────────────
# Each category maps a display label to a regex pattern.
@@ -52,7 +52,7 @@ AUTO_TAG_GROUPS = {
DEFAULT_ENABLED_GROUPS = {"mode", "video"}
def _collect_sources(model_data: Dict) -> List[str]:
def _collect_sources(model_data: Dict[str, Any]) -> List[str]:
"""Collect all text sources from model data for tag matching."""
sources: List[str] = []
@@ -73,7 +73,7 @@ def _collect_sources(model_data: Dict) -> List[str]:
return sources
def extract_auto_tags(model_data: Dict) -> List[str]:
def extract_auto_tags(model_data: Dict[str, Any]) -> List[str]:
"""Extract auto-detected tags from model metadata.
Uses a two-layer approach:

6
py/services/autov3_backfill_service.py
View File

@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
"""Backfill the AutoV3 checked state for models loaded from a persisted snapshot.
The SQLite persistent cache predates the AutoV3 feature, so entries hydrated
@@ -66,7 +70,7 @@ class Autov3BackfillService:
# initialize concurrently (lora_manager.py), so a global guard would
# silently skip every type but the first to start. Each model type
# runs its own backfill; a duplicate trigger for the same type no-ops.
self._running_types: set = set()
self._running_types: set[str] = set()
@classmethod
def get_instance(cls) -> "Autov3BackfillService":

4
py/services/backup_service.py
View File

@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
from __future__ import annotations
import asyncio

128
py/services/base_model_service.py
View File

@@ -2,7 +2,7 @@ from abc import ABC, abstractmethod
import asyncio
import re
import random
from typing import Any, Dict, List, Optional, Type, Union, TYPE_CHECKING
from typing import Any, Awaitable, Dict, List, Optional, Type, Union, TYPE_CHECKING, cast
import logging
import os
import time
@@ -70,24 +70,24 @@ class BaseModelService(ABC):
page: int,
page_size: int,
sort_by: str = "name",
folder: str = None,
folder_include: list = None,
folder_exclude: list = None,
search: str = None,
folder: str | None = None,
folder_include: list[str] | None = None,
folder_exclude: list[str] | None = None,
search: str | None = None,
fuzzy_search: bool = False,
base_models: list = None,
model_types: list = None,
base_models: list[str] | None = None,
model_types: list[str] | None = None,
tags: Optional[Dict[str, str]] = None,
auto_tags: Optional[Dict[str, str]] = None,
search_options: dict = None,
hash_filters: dict = None,
search_options: dict[str, Any] | None = None,
hash_filters: dict[str, Any] | None = None,
favorites_only: bool = False,
update_available_only: bool = False,
credit_required: Optional[bool] = None,
allow_selling_generated_content: Optional[bool] = None,
tag_logic: str = "any",
**kwargs,
) -> Dict:
) -> Dict[str, Any]:
"""Get paginated and filtered model data"""
overall_start = time.perf_counter()
@@ -178,8 +178,8 @@ class BaseModelService(ABC):
ufs = self.settings.get("version_grouping", "same_base")
group_by_base = ufs == "same_base"
model_groups: Dict[Any, List[Dict]] = {}
ungrouped_standalone: List[Dict] = []
model_groups: Dict[Any, List[Dict[str, Any]]] = {}
ungrouped_standalone: List[Dict[str, Any]] = []
for item in sorted_data:
mid = self._extract_group_key(item)
if mid is None:
@@ -249,7 +249,7 @@ class BaseModelService(ABC):
filter_duration = time.perf_counter() - t1
post_filter_count = len(filtered_data)
annotated_for_filter: Optional[List[Dict]] = None
annotated_for_filter: Optional[List[Dict[str, Any]]] = None
t2 = time.perf_counter()
if update_available_only:
annotated_for_filter = await self._annotate_update_flags(filtered_data)
@@ -296,11 +296,11 @@ class BaseModelService(ABC):
page: int,
page_size: int,
sort_by: str = "name",
search: str = None,
search: str | None = None,
fuzzy_search: bool = False,
search_options: dict = None,
search_options: dict[str, Any] | None = None,
**kwargs,
) -> Dict:
) -> Dict[str, Any]:
"""Get paginated excluded model data."""
excluded_paths = list(self.scanner.get_excluded_models())
excluded_entries: List[Dict[str, Any]] = []
@@ -326,7 +326,7 @@ class BaseModelService(ABC):
]
persist_current_cache = getattr(self.scanner, "_persist_current_cache", None)
if callable(persist_current_cache):
await persist_current_cache()
await cast(Awaitable[Any], persist_current_cache())
excluded_entries = self._sort_entries(excluded_entries, sort_by)
@@ -444,11 +444,11 @@ class BaseModelService(ABC):
return entry
async def _apply_hash_filters(
self, data: List[Dict], hash_filters: Dict
) -> List[Dict]:
self, data: List[Dict[str, Any]], hash_filters: Dict[str, Any]
) -> List[Dict[str, Any]]:
"""Apply hash-based filtering (SHA256 and AutoV3)."""
def matches_hash_set(item: Dict, hash_set: set) -> bool:
def matches_hash_set(item: Dict[str, Any], hash_set: set[str]) -> bool:
"""Check whether an item matches any hash in the set.
Compares the item's ``sha256`` field and its non-empty ``autov3``
@@ -476,18 +476,18 @@ class BaseModelService(ABC):
async def _apply_common_filters(
self,
data: List[Dict],
folder: str = None,
folder_include: list = None,
folder_exclude: list = None,
base_models: list = None,
model_types: list = None,
data: List[Dict[str, Any]],
folder: str | None = None,
folder_include: list[str] | None = None,
folder_exclude: list[str] | None = None,
base_models: list[str] | None = None,
model_types: list[str] | None = None,
tags: Optional[Dict[str, str]] = None,
auto_tags: Optional[Dict[str, str]] = None,
favorites_only: bool = False,
search_options: dict = None,
search_options: dict[str, Any] | None = None,
tag_logic: str = "any",
) -> List[Dict]:
) -> List[Dict[str, Any]]:
"""Apply common filters that work across all model types"""
normalized_options = self.search_strategy.normalize_options(search_options)
criteria = FilterCriteria(
@@ -506,24 +506,24 @@ class BaseModelService(ABC):
async def _apply_search_filters(
self,
data: List[Dict],
data: List[Dict[str, Any]],
search: str,
fuzzy_search: bool,
search_options: dict,
) -> List[Dict]:
search_options: dict[str, Any] | None,
) -> List[Dict[str, Any]]:
"""Apply search filtering"""
normalized_options = self.search_strategy.normalize_options(search_options)
return self.search_strategy.apply(
data, search, normalized_options, fuzzy_search
)
async def _apply_specific_filters(self, data: List[Dict], **kwargs) -> List[Dict]:
async def _apply_specific_filters(self, data: List[Dict[str, Any]], **kwargs) -> List[Dict[str, Any]]:
"""Apply model-specific filters - to be overridden by subclasses if needed"""
return data
async def _apply_credit_required_filter(
self, data: List[Dict], credit_required: bool
) -> List[Dict]:
self, data: List[Dict[str, Any]], credit_required: bool
) -> List[Dict[str, Any]]:
"""Apply credit required filtering based on license_flags.
Args:
@@ -553,8 +553,8 @@ class BaseModelService(ABC):
return filtered_data
async def _apply_allow_selling_filter(
self, data: List[Dict], allow_selling: bool
) -> List[Dict]:
self, data: List[Dict[str, Any]], allow_selling: bool
) -> List[Dict[str, Any]]:
"""Apply allow selling generated content filtering based on license_flags.
Args:
@@ -586,8 +586,8 @@ class BaseModelService(ABC):
async def _annotate_update_flags(
self,
items: List[Dict],
) -> List[Dict]:
items: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
"""Attach an update_available flag to each response item.
Items without a civitai model id default to False.
@@ -602,7 +602,7 @@ class BaseModelService(ABC):
item["update_available"] = False
return annotated
id_to_items: Dict[int, List[Dict]] = {}
id_to_items: Dict[int, List[Dict[str, Any]]] = {}
ordered_ids: List[int] = []
for item in annotated:
model_id = self._extract_model_id(item)
@@ -639,7 +639,7 @@ class BaseModelService(ABC):
record_method = getattr(self.update_service, "get_records_bulk", None)
if callable(record_method):
try:
records = await record_method(self.model_type, ordered_ids)
records = await cast(Awaitable[Any], record_method(self.model_type, ordered_ids))
resolved = {
model_id: record.has_update(hide_early_access=hide_early_access)
for model_id, record in records.items()
@@ -659,11 +659,11 @@ class BaseModelService(ABC):
bulk_method = getattr(self.update_service, "has_updates_bulk", None)
if callable(bulk_method):
try:
resolved = await bulk_method(
resolved = await cast(Awaitable[Any], bulk_method(
self.model_type,
ordered_ids,
hide_early_access=hide_early_access,
)
))
except Exception as exc:
logger.error(
"Failed to resolve update status in bulk for %s models (%s): %s",
@@ -725,7 +725,7 @@ class BaseModelService(ABC):
return annotated
@staticmethod
def _extract_hf_group_key(item: Dict) -> Optional[str]:
def _extract_hf_group_key(item: Dict[str, Any]) -> Optional[str]:
"""Extract `hf:{owner}/{repo}` from item's ``hf_url``, or None."""
hf_url = item.get("hf_url") if isinstance(item, dict) else None
if not hf_url or not isinstance(hf_url, str):
@@ -738,7 +738,7 @@ class BaseModelService(ABC):
return f"hf:{m.group(1)}"
@staticmethod
def _extract_group_key(item: Dict) -> Union[int, str, None]:
def _extract_group_key(item: Dict[str, Any]) -> Union[int, str, None]:
"""Return the group identity key: CivitAI modelId (int) or HF repo (str).
Preference order:
@@ -752,7 +752,7 @@ class BaseModelService(ABC):
return BaseModelService._extract_hf_group_key(item)
@staticmethod
def _extract_model_id(item: Dict) -> Optional[int]:
def _extract_model_id(item: Dict[str, Any]) -> Optional[int]:
civitai = item.get("civitai") if isinstance(item, dict) else None
if not isinstance(civitai, dict):
return None
@@ -765,7 +765,7 @@ class BaseModelService(ABC):
return None
@staticmethod
def _extract_version_id(item: Dict) -> Optional[int]:
def _extract_version_id(item: Dict[str, Any]) -> Optional[int]:
civitai = item.get("civitai") if isinstance(item, dict) else None
if not isinstance(civitai, dict):
return None
@@ -778,7 +778,7 @@ class BaseModelService(ABC):
return None
@staticmethod
def _extract_base_model(item: Dict) -> Optional[str]:
def _extract_base_model(item: Dict[str, Any]) -> Optional[str]:
value = item.get("base_model")
if value is None:
return None
@@ -830,7 +830,7 @@ class BaseModelService(ABC):
return highest_by_base
def _paginate(self, data: List[Dict], page: int, page_size: int) -> Dict:
def _paginate(self, data: List[Dict[str, Any]], page: int, page_size: int) -> Dict[str, Any]:
"""Apply pagination to filtered data"""
total_items = len(data)
start_idx = (page - 1) * page_size
@@ -845,7 +845,7 @@ class BaseModelService(ABC):
}
@abstractmethod
async def format_response(self, model_data: Dict) -> Optional[Dict]:
async def format_response(self, model_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""Format model data for API response - must be implemented by subclasses.
Subclasses should return None for corrupted entries so the handler
@@ -854,17 +854,17 @@ class BaseModelService(ABC):
pass
# Common service methods that delegate to scanner
async def get_top_tags(self, limit: int = 20) -> List[Dict]:
async def get_top_tags(self, limit: int = 20) -> List[Dict[str, Any]]:
"""Get top tags sorted by frequency"""
return await self.scanner.get_top_tags(limit)
async def search_tags(
self, query: str, limit: int = 50
) -> List[Dict]:
) -> List[Dict[str, Any]]:
"""Search tags by substring, sorted by frequency"""
return await self.scanner.search_tags(query, limit)
async def get_base_models(self, limit: int = 20) -> List[Dict]:
async def get_base_models(self, limit: int = 20) -> List[Dict[str, Any]]:
"""Get base models sorted by frequency"""
return await self.scanner.get_base_models(limit)
@@ -931,7 +931,7 @@ class BaseModelService(ABC):
"""Get model root directories"""
return self.scanner.get_model_roots()
def filter_civitai_data(self, data: Dict, minimal: bool = False) -> Dict:
def filter_civitai_data(self, data: Dict[str, Any], minimal: bool = False) -> Dict[str, Any]:
"""Filter relevant fields from CivitAI data"""
if not data:
return {}
@@ -957,7 +957,7 @@ class BaseModelService(ABC):
)
return {k: data[k] for k in fields if k in data}
async def get_folder_tree(self, model_root: str) -> Dict:
async def get_folder_tree(self, model_root: str) -> Dict[str, Any]:
"""Get hierarchical folder tree for a specific model root"""
cache = await self.scanner.get_cached_data()
@@ -986,7 +986,7 @@ class BaseModelService(ABC):
return tree
async def get_unified_folder_tree(self) -> Dict:
async def get_unified_folder_tree(self) -> Dict[str, Any]:
"""Get unified folder tree across all model roots"""
cache = await self.scanner.get_cached_data()
@@ -1015,7 +1015,7 @@ class BaseModelService(ABC):
return unified_tree
async def get_model_notes(self, model_name: str) -> Optional[dict]:
async def get_model_notes(self, model_name: str) -> Optional[dict[str, Any]]:
"""Get notes and file_path for a specific model file.
Supports both simple names (``OWSMianne_ANIMA_V1``) and full-path
@@ -1147,7 +1147,7 @@ class BaseModelService(ABC):
return {"civitai_url": None, "model_id": None, "version_id": None}
async def get_model_metadata(self, file_path: str) -> Optional[Dict]:
async def get_model_metadata(self, file_path: str) -> Optional[Dict[str, Any]]:
"""Load full metadata for a single model.
Listing/search endpoints return lightweight cache entries; this method performs
@@ -1243,7 +1243,7 @@ class BaseModelService(ABC):
return True
@staticmethod
def _relative_path_sort_key(relative_path: str, include_terms: List[str]) -> tuple:
def _relative_path_sort_key(relative_path: str, include_terms: List[str]) -> tuple[int, int, int, str]:
"""Sort paths by how well they satisfy the include tokens.
Sorts based on path without extension for consistent ordering.
@@ -1276,12 +1276,12 @@ class BaseModelService(ABC):
offset: int = 0,
*,
folder: Optional[str] = None,
folder_include: Optional[list] = None,
folder_exclude: Optional[list] = None,
base_models: Optional[list] = None,
model_types: Optional[list] = None,
tags: Optional[dict] = None,
auto_tags: Optional[dict] = None,
folder_include: Optional[list[str]] = None,
folder_exclude: Optional[list[str]] = None,
base_models: Optional[list[str]] = None,
model_types: Optional[list[str]] = None,
tags: Optional[dict[str, str]] = None,
auto_tags: Optional[dict[str, str]] = None,
tag_logic: str = "any",
credit_required: Optional[bool] = None,
allow_selling_generated_content: Optional[bool] = None,

6
py/services/checkpoint_scanner.py
View File

@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import asyncio
import json
import logging
@@ -427,7 +431,7 @@ class CheckpointScanner(ModelScanner):
roots.extend(config.extra_checkpoints_roots or [])
roots.extend(config.extra_unet_roots or [])
# Remove duplicates while preserving order
seen: set = set()
seen: set[str] = set()
unique_roots: List[str] = []
for root in roots:
if root not in seen:

56
py/services/checkpoint_service.py
View File

@@ -1,6 +1,6 @@
import os
import logging
from typing import Dict, Optional
from typing import Any, Dict, Optional
from .base_model_service import BaseModelService
from .auto_tag_service import extract_auto_tags
@@ -21,58 +21,58 @@ class CheckpointService(BaseModelService):
"""
super().__init__("checkpoint", scanner, CheckpointMetadata, update_service=update_service)
async def format_response(self, checkpoint_data: Dict) -> Optional[Dict]:
async def format_response(self, model_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""Format Checkpoint data for API response.
Returns None when the entry is missing critical fields (corrupted cache
row), so the handler layer can filter it out. See issue #730.
"""
# Guard against corrupted cache entries missing critical fields
file_path = checkpoint_data.get("file_path")
file_path = model_data.get("file_path")
if not file_path or not isinstance(file_path, str):
logger.warning(
"Skipping corrupted checkpoint entry (missing file_path): %s",
checkpoint_data.get("file_name", "<unknown>"),
model_data.get("file_name", "<unknown>"),
)
return None
# Get sub_type from cache entry (new canonical field)
sub_type = checkpoint_data.get("sub_type", "checkpoint")
sub_type = model_data.get("sub_type", "checkpoint")
file_name = checkpoint_data.get("file_name") or ""
model_name = checkpoint_data.get("model_name") or file_name
folder = checkpoint_data.get("folder") or ""
file_name = model_data.get("file_name") or ""
model_name = model_data.get("model_name") or file_name
folder = model_data.get("folder") or ""
return {
"model_name": model_name,
"file_name": file_name,
"preview_url": config.get_preview_static_url(checkpoint_data.get("preview_url", "")),
"preview_nsfw_level": checkpoint_data.get("preview_nsfw_level", 0),
"base_model": checkpoint_data.get("base_model", ""),
"preview_url": config.get_preview_static_url(model_data.get("preview_url", "")),
"preview_nsfw_level": model_data.get("preview_nsfw_level", 0),
"base_model": model_data.get("base_model", ""),
"folder": folder,
"sha256": checkpoint_data.get("sha256", ""),
"sha256": model_data.get("sha256", ""),
"file_path": file_path.replace(os.sep, "/"),
"file_size": checkpoint_data.get("size", 0),
"modified": checkpoint_data.get("modified", ""),
"tags": checkpoint_data.get("tags", []),
"from_civitai": checkpoint_data.get("from_civitai", True),
"usage_count": checkpoint_data.get("usage_count", 0),
"notes": checkpoint_data.get("notes", ""),
"file_size": model_data.get("size", 0),
"modified": model_data.get("modified", ""),
"tags": model_data.get("tags", []),
"from_civitai": model_data.get("from_civitai", True),
"usage_count": model_data.get("usage_count", 0),
"notes": model_data.get("notes", ""),
"sub_type": sub_type,
"favorite": checkpoint_data.get("favorite", False),
"exclude": bool(checkpoint_data.get("exclude", False)),
"update_available": bool(checkpoint_data.get("update_available", False)),
"skip_metadata_refresh": bool(checkpoint_data.get("skip_metadata_refresh", False)),
"civitai": self.filter_civitai_data(checkpoint_data.get("civitai", {}), minimal=True),
"auto_tags": checkpoint_data.get("auto_tags") or extract_auto_tags(checkpoint_data),
"version_count": checkpoint_data.get("version_count"),
"hf_url": checkpoint_data.get("hf_url", ""),
"favorite": model_data.get("favorite", False),
"exclude": bool(model_data.get("exclude", False)),
"update_available": bool(model_data.get("update_available", False)),
"skip_metadata_refresh": bool(model_data.get("skip_metadata_refresh", False)),
"civitai": self.filter_civitai_data(model_data.get("civitai", {}), minimal=True),
"auto_tags": model_data.get("auto_tags") or extract_auto_tags(model_data),
"version_count": model_data.get("version_count"),
"hf_url": model_data.get("hf_url", ""),
}
def find_duplicate_hashes(self) -> Dict:
def find_duplicate_hashes(self) -> Dict[str, Any]:
"""Find Checkpoints with duplicate SHA256 hashes"""
return self.scanner._hash_index.get_duplicate_hashes()
def find_duplicate_filenames(self) -> Dict:
def find_duplicate_filenames(self) -> Dict[str, Any]:
"""Find Checkpoints with conflicting filenames"""
return self.scanner._hash_index.get_duplicate_filenames()

77
py/services/civarchive_client.py
View File

@@ -1,8 +1,12 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import json
import logging
import asyncio
from copy import deepcopy
from typing import Optional, Dict, Tuple, List
from typing import Any, Optional, Dict, Tuple, List, cast
from .model_metadata_provider import CivArchiveModelMetadataProvider, ModelMetadataProviderManager
from .downloader import get_downloader
from .errors import RateLimitError
@@ -37,8 +41,8 @@ class CivArchiveClient:
async def _request_json(
self,
path: str,
params: Optional[Dict[str, str]] = None
) -> Tuple[Optional[Dict], Optional[str]]:
params: Optional[Dict[str, Any]] = None
) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Call CivArchive API and return JSON payload"""
success, payload = await self._make_request(path, params=params)
if not success:
@@ -52,12 +56,12 @@ class CivArchiveClient:
self,
path: str,
*,
params: Optional[Dict[str, str]] = None,
) -> Tuple[bool, Dict | str]:
params: Optional[Dict[str, Any]] = None,
) -> Tuple[bool, Dict[str, Any] | str]:
"""Wrapper around downloader.make_request that surfaces rate limits."""
downloader = await get_downloader()
kwargs: Dict[str, Dict[str, str]] = {}
kwargs: Dict[str, Dict[str, Any]] = {}
if params:
safe_params = {str(key): str(value) for key, value in params.items() if value is not None}
if safe_params:
@@ -73,10 +77,11 @@ class CivArchiveClient:
if payload.provider is None:
payload.provider = "civarchive_api"
raise payload
return success, payload
# RateLimitError is always raised above, so the returned payload is a dict or str.
return success, cast(Dict[str, Any] | str, payload)
@staticmethod
def _normalize_payload(payload: Dict) -> Dict:
def _normalize_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
"""Unwrap CivArchive responses that wrap content under a data key"""
if not isinstance(payload, dict):
return {}
@@ -86,12 +91,12 @@ class CivArchiveClient:
return payload
@staticmethod
def _split_context(payload: Dict) -> Tuple[Dict, Dict, List[Dict]]:
def _split_context(payload: Dict[str, Any]) -> Tuple[Dict[str, Any], Dict[str, Any], List[Dict[str, Any]]]:
"""Separate version payload from surrounding model context"""
data = CivArchiveClient._normalize_payload(payload)
context: Dict = {}
fallback_files: List[Dict] = []
version: Dict = {}
context: Dict[str, Any] = {}
fallback_files: List[Dict[str, Any]] = []
version: Dict[str, Any] = {}
for key, value in data.items():
if key in {"version", "model"}:
@@ -115,7 +120,7 @@ class CivArchiveClient:
return context, version, fallback_files
@staticmethod
def _ensure_list(value) -> List:
def _ensure_list(value: Any) -> List[Any]:
if isinstance(value, list):
return value
if value is None:
@@ -123,7 +128,7 @@ class CivArchiveClient:
return [value]
@staticmethod
def _build_model_info(context: Dict) -> Dict:
def _build_model_info(context: Dict[str, Any]) -> Dict[str, Any]:
tags = context.get("tags")
if not isinstance(tags, list):
tags = list(tags) if isinstance(tags, (set, tuple)) else ([] if tags is None else [tags])
@@ -136,7 +141,7 @@ class CivArchiveClient:
}
@staticmethod
def _build_creator_info(context: Dict) -> Dict:
def _build_creator_info(context: Dict[str, Any]) -> Dict[str, Any]:
username = context.get("creator_username") or context.get("username") or ""
image = context.get("creator_image") or context.get("creator_avatar") or ""
creator: Dict[str, Optional[str]] = {
@@ -150,7 +155,7 @@ class CivArchiveClient:
return creator
@staticmethod
def _transform_file_entry(file_data: Dict) -> Dict:
def _transform_file_entry(file_data: Dict[str, Any]) -> Dict[str, Any]:
mirrors = file_data.get("mirrors") or []
if not isinstance(mirrors, list):
mirrors = [mirrors]
@@ -165,7 +170,7 @@ class CivArchiveClient:
if not name and available_mirror:
name = available_mirror.get("filename")
transformed: Dict = {
transformed: Dict[str, Any] = {
"id": file_data.get("id"),
"sizeKB": file_data.get("sizeKB"),
"name": name,
@@ -216,23 +221,23 @@ class CivArchiveClient:
def _transform_files(
self,
files: Optional[List[Dict]],
fallback_files: Optional[List[Dict]] = None
) -> List[Dict]:
candidates: List[Dict] = []
files: Optional[List[Dict[str, Any]]],
fallback_files: Optional[List[Dict[str, Any]]] = None
) -> List[Dict[str, Any]]:
candidates: List[Dict[str, Any]] = []
if isinstance(files, list) and files:
candidates = files
elif isinstance(fallback_files, list):
candidates = fallback_files
transformed_files: List[Dict] = []
transformed_files: List[Dict[str, Any]] = []
for file_data in candidates:
if isinstance(file_data, dict):
transformed_files.append(self._transform_file_entry(file_data))
# Sort: .safetensors first, .ckpt second, others last
# so the backend fallback (no file_params) prefers safetensors
def _sort_key(f: Dict) -> int:
def _sort_key(f: Dict[str, Any]) -> int:
fname = f.get("name") or ""
if isinstance(fname, str):
lower = fname.lower()
@@ -247,10 +252,10 @@ class CivArchiveClient:
def _transform_version(
self,
context: Dict,
version: Dict,
fallback_files: Optional[List[Dict]] = None
) -> Optional[Dict]:
context: Dict[str, Any],
version: Dict[str, Any],
fallback_files: Optional[List[Dict[str, Any]]] = None
) -> Optional[Dict[str, Any]]:
if not version:
return None
@@ -291,7 +296,7 @@ class CivArchiveClient:
return version_copy
async def _resolve_version_from_files(self, payload: Dict) -> Optional[Dict]:
async def _resolve_version_from_files(self, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""Fallback to fetch version data when only file metadata is available"""
data = self._normalize_payload(payload)
files = data.get("files") or payload.get("files") or []
@@ -323,7 +328,7 @@ class CivArchiveClient:
return resolved
return None
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Find model by SHA256 hash value using CivArchive API"""
try:
payload, error = await self._request_json(f"/sha256/{model_hash.lower()}")
@@ -332,12 +337,12 @@ class CivArchiveClient:
return None, "Model not found"
return None, error
context, version_data, fallback_files = self._split_context(payload)
context, version_data, fallback_files = self._split_context(cast(Dict[str, Any], payload))
transformed = self._transform_version(context, version_data, fallback_files)
if transformed:
return transformed, None
resolved = await self._resolve_version_from_files(payload)
resolved = await self._resolve_version_from_files(cast(Dict[str, Any], payload))
if resolved:
return resolved, None
@@ -350,7 +355,7 @@ class CivArchiveClient:
logger.error(f"Error fetching CivArchive model by hash {model_hash[:10]}: {e}")
return None, str(e)
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
"""Get all versions of a model using CivArchive API"""
try:
payload, error = await self._request_json(f"/models/{model_id}")
@@ -364,7 +369,7 @@ class CivArchiveClient:
context, version_data, fallback_files = self._split_context(payload)
versions_meta = data.get("versions") or []
transformed_versions: List[Dict] = []
transformed_versions: List[Dict[str, Any]] = []
for meta in versions_meta:
if not isinstance(meta, dict):
continue
@@ -381,7 +386,7 @@ class CivArchiveClient:
if primary_version:
transformed_versions.insert(0, primary_version)
ordered_versions: List[Dict] = []
ordered_versions: List[Dict[str, Any]] = []
seen_ids = set()
for version in transformed_versions:
version_id = version.get("id")
@@ -402,7 +407,7 @@ class CivArchiveClient:
logger.error(f"Error fetching CivArchive model versions for {model_id}: {e}")
return None
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: int | str | None = None, version_id: int | str | None = None) -> Optional[Dict[str, Any]]:
"""Get specific model version using CivArchive API
Args:
@@ -459,7 +464,7 @@ class CivArchiveClient:
logger.error(f"Error fetching CivArchive model version via API {model_id}/{version_id}: {e}")
return None
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
""" Fetch model version metadata using a known bogus model lookup
CivArchive lacks a direct version lookup API, this uses a workaround (which we handle in the main model request now)

2
py/services/civitai_base_model_service.py
View File

@@ -283,7 +283,7 @@ class CivitaiBaseModelService:
return None
if isinstance(result, str):
data = json.loads(result)
data: Any = json.loads(result)
else:
data = result

75
py/services/civitai_client.py
View File

@@ -1,10 +1,14 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import asyncio
import copy
import logging
import os
import time
from collections import OrderedDict
from typing import Any, Optional, Dict, Tuple, List, Sequence
from typing import Any, Optional, Dict, Tuple, List, Sequence, cast
from .connectivity_guard import (
OFFLINE_FRIENDLY_MESSAGE,
is_expected_offline_error,
@@ -58,7 +62,7 @@ class CivitaiClient:
# Uses OrderedDict with LRU eviction at MAX_CACHE_ENTRIES to prevent
# unbounded growth in long-running server processes.
self._version_info_cache: OrderedDict[
str, Tuple[Optional[Dict], Optional[str]]
str, Tuple[Optional[Dict[str, Any]], Optional[str]]
] = OrderedDict()
self._MAX_CACHE_ENTRIES = 500
@@ -72,7 +76,7 @@ class CivitaiClient:
*,
use_auth: bool = False,
**kwargs,
) -> Tuple[bool, Dict | str]:
) -> Tuple[bool, Dict[str, Any] | str]:
"""Wrapper around downloader.make_request that surfaces rate limits,
with retry for transient server errors (5xx, Cloudflare 524, network flakiness)."""
@@ -86,7 +90,8 @@ class CivitaiClient:
**kwargs,
)
if success:
return True, result
# RateLimitError is raised below; a successful result is dict or str.
return True, cast(Dict[str, Any] | str, result)
if isinstance(result, RateLimitError):
if result.provider is None:
@@ -126,7 +131,7 @@ class CivitaiClient:
return False, "Unexpected error in _make_request"
@staticmethod
def _remove_comfy_metadata(model_version: Optional[Dict]) -> None:
def _remove_comfy_metadata(model_version: Optional[Dict[str, Any]]) -> None:
"""Remove Comfy-specific metadata from model version images."""
if not isinstance(model_version, dict):
return
@@ -173,7 +178,7 @@ class CivitaiClient:
async def get_model_by_hash(
self, model_hash: str
) -> Tuple[Optional[Dict], Optional[str]]:
) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
try:
success, version = await self._make_request(
"GET",
@@ -220,7 +225,7 @@ class CivitaiClient:
# Ensure directory exists
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with open(save_path, "wb") as f:
f.write(content)
f.write(content if isinstance(content, bytes) else content.encode("utf-8"))
return True
return False
except Exception as e:
@@ -275,7 +280,7 @@ class CivitaiClient:
return True
return False
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
"""Get all versions of a model with local availability info"""
try:
success, result = await self._make_request(
@@ -283,7 +288,7 @@ class CivitaiClient:
f"{self.base_url}/models/{model_id}",
use_auth=True,
)
if success:
if success and isinstance(result, dict):
# Also return model type along with versions
return {
"modelVersions": result.get("modelVersions", []),
@@ -317,7 +322,7 @@ class CivitaiClient:
async def get_model_versions_bulk(
self, model_ids: Sequence[int]
) -> Optional[Dict[int, Dict]]:
) -> Optional[Dict[int, Dict[str, Any]]]:
"""Fetch model metadata for multiple ids using the batch API."""
deduped: Dict[int, None] = {}
@@ -347,13 +352,13 @@ class CivitaiClient:
if not isinstance(items, list):
return {}
payload: Dict[int, Dict] = {}
payload: Dict[int, Dict[str, Any]] = {}
for item in items:
if not isinstance(item, dict):
continue
model_id = item.get("id")
try:
normalized_id = int(model_id)
normalized_id = int(cast(Any, model_id))
except (TypeError, ValueError):
continue
payload[normalized_id] = {
@@ -373,8 +378,8 @@ class CivitaiClient:
return None
async def get_model_version(
self, model_id: int = None, version_id: int = None
) -> Optional[Dict]:
self, model_id: int | None = None, version_id: int | None = None
) -> Optional[Dict[str, Any]]:
"""Get specific model version with additional metadata."""
try:
if model_id is None and version_id is not None:
@@ -392,7 +397,7 @@ class CivitaiClient:
logger.error(f"Error fetching model version: {e}")
return None
async def _get_version_by_id_only(self, version_id: int) -> Optional[Dict]:
async def _get_version_by_id_only(self, version_id: int) -> Optional[Dict[str, Any]]:
version = await self._fetch_version_by_id(version_id)
if version is None:
return None
@@ -411,7 +416,7 @@ class CivitaiClient:
async def _get_version_with_model_id(
self, model_id: int, version_id: Optional[int]
) -> Optional[Dict]:
) -> Optional[Dict[str, Any]]:
model_data = await self._fetch_model_data(model_id)
if not model_data:
return None
@@ -464,20 +469,20 @@ class CivitaiClient:
self._remove_comfy_metadata(version)
return version
async def _fetch_model_data(self, model_id: int) -> Optional[Dict]:
async def _fetch_model_data(self, model_id: int) -> Optional[Dict[str, Any]]:
success, data = await self._make_request(
"GET",
f"{self.base_url}/models/{model_id}",
use_auth=True,
)
if success:
if success and isinstance(data, dict):
return data
if is_expected_offline_error(data):
return None
logger.warning(f"Failed to fetch model data for model {model_id}")
return None
async def _fetch_version_by_id(self, version_id: Optional[int]) -> Optional[Dict]:
async def _fetch_version_by_id(self, version_id: Optional[int]) -> Optional[Dict[str, Any]]:
if version_id is None:
return None
@@ -486,7 +491,7 @@ class CivitaiClient:
f"{self.base_url}/model-versions/{version_id}",
use_auth=True,
)
if success:
if success and isinstance(version, dict):
return version
if is_expected_offline_error(version):
return None
@@ -494,7 +499,7 @@ class CivitaiClient:
logger.warning(f"Failed to fetch version by id {version_id}")
return None
async def _fetch_version_by_hash(self, model_hash: Optional[str]) -> Optional[Dict]:
async def _fetch_version_by_hash(self, model_hash: Optional[str]) -> Optional[Dict[str, Any]]:
if not model_hash:
return None
@@ -503,7 +508,7 @@ class CivitaiClient:
f"{self.base_url}/model-versions/by-hash/{model_hash}",
use_auth=True,
)
if success:
if success and isinstance(version, dict):
return version
if is_expected_offline_error(version):
return None
@@ -512,8 +517,8 @@ class CivitaiClient:
return None
def _select_target_version(
self, model_data: Dict, model_id: int, version_id: Optional[int]
) -> Optional[Dict]:
self, model_data: Dict[str, Any], model_id: int, version_id: Optional[int]
) -> Optional[Dict[str, Any]]:
model_versions = model_data.get("modelVersions", [])
if not model_versions:
logger.warning(f"No model versions found for model {model_id}")
@@ -532,7 +537,7 @@ class CivitaiClient:
return model_versions[0]
def _extract_primary_model_hash(self, version_entry: Dict) -> Optional[str]:
def _extract_primary_model_hash(self, version_entry: Dict[str, Any]) -> Optional[str]:
for file_info in version_entry.get("files", []):
if file_info.get("type") == "Model" and file_info.get("primary"):
hashes = file_info.get("hashes", {})
@@ -542,8 +547,8 @@ class CivitaiClient:
return None
def _build_version_from_model_data(
self, version_entry: Dict, model_id: int, model_data: Dict
) -> Dict:
self, version_entry: Dict[str, Any], model_id: int, model_data: Dict[str, Any]
) -> Dict[str, Any]:
version = copy.deepcopy(version_entry)
version.pop("index", None)
version["modelId"] = model_id
@@ -555,7 +560,7 @@ class CivitaiClient:
}
return version
def _enrich_version_with_model_data(self, version: Dict, model_data: Dict) -> None:
def _enrich_version_with_model_data(self, version: Dict[str, Any], model_data: Dict[str, Any]) -> None:
model_info = version.get("model")
if not isinstance(model_info, dict):
model_info = {}
@@ -571,7 +576,7 @@ class CivitaiClient:
async def get_model_version_info(
self, version_id: str
) -> Tuple[Optional[Dict], Optional[str]]:
) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Fetch model version metadata from Civitai
Args:
@@ -596,7 +601,7 @@ class CivitaiClient:
logger.debug("Resolving Civitai model version info: %s", url)
success, result = await self._make_request("GET", url, use_auth=True)
if success:
if success and isinstance(result, dict):
logger.debug("Successfully fetched model version info for: %s", version_id)
self._remove_comfy_metadata(result)
self._version_info_cache[version_id] = (result, None)
@@ -626,7 +631,7 @@ class CivitaiClient:
async def get_image_info(
self, image_id: str, source_url: str | None = None
) -> Optional[Dict]:
) -> Optional[Dict[str, Any]]:
"""Fetch image information from Civitai API
Args:
@@ -659,7 +664,7 @@ class CivitaiClient:
)
return None
if result and "items" in result and isinstance(result["items"], list):
if isinstance(result, dict) and "items" in result and isinstance(result["items"], list):
items = result["items"]
for item in items:
@@ -699,7 +704,7 @@ class CivitaiClient:
async def get_model_versions_by_hashes(
self, hashes: List[str]
) -> Optional[List[Dict]]:
) -> Optional[List[Dict[str, Any]]]:
"""Fetch full version details for up to 100 SHA256 hashes via the batch endpoint.
Uses POST /api/v1/model-versions/by-hash which returns full version
@@ -716,7 +721,7 @@ class CivitaiClient:
return []
BATCH_SIZE = 100
all_versions: List[Dict] = []
all_versions: List[Dict[str, Any]] = []
for start in range(0, len(hashes), BATCH_SIZE):
batch = hashes[start : start + BATCH_SIZE]
@@ -736,7 +741,7 @@ class CivitaiClient:
continue
if isinstance(result, list):
all_versions.extend(result)
all_versions.extend(cast(Any, result))
else:
logger.debug(
"Unexpected by-hash response type: %s", type(result)

2
py/services/download_coordinator.py
View File

@@ -18,7 +18,7 @@ class DownloadCoordinator:
self,
*,
ws_manager,
download_manager_factory: Callable[[], Awaitable],
download_manager_factory: Callable[[], Awaitable[Any]],
) -> None:
self._ws_manager = ws_manager
self._download_manager_factory = download_manager_factory

173
py/services/download_manager.py
View File

@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import copy
import logging
import os
@@ -8,7 +12,7 @@ import zipfile
from concurrent.futures import ThreadPoolExecutor
from collections import OrderedDict
import uuid
from typing import Dict, List, Optional, Set, Tuple
from typing import Any, Dict, List, Optional, Set, Tuple, cast
from urllib.parse import urlparse
from ..utils.models import LoraMetadata, CheckpointMetadata, EmbeddingMetadata
from ..utils.constants import (
@@ -121,7 +125,7 @@ class DownloadManager:
"delay": 0,
}
)
except DownloadInProgressError:
except DownloadInProgressError: # pyright: ignore[reportPossiblyUnboundVariable]
logger.info(
"Skipping automatic example images download for %s; another example images download is already running",
model_hash,
@@ -170,7 +174,7 @@ class DownloadManager:
logger.error("aria2 download failed for %s: %s", download_url, exc)
return False, str(exc)
download_kwargs = {
download_kwargs: Dict[str, Any] = {
"progress_callback": progress_callback,
"use_auth": use_auth,
}
@@ -204,16 +208,16 @@ class DownloadManager:
async def download_from_civitai(
self,
model_id: int = None,
model_version_id: int = None,
save_dir: str = None,
model_id: int | None = None,
model_version_id: int | None = None,
save_dir: str | None = None,
relative_path: str = "",
progress_callback=None,
use_default_paths: bool = False,
download_id: str = None,
source: str = None,
file_params: Dict = None,
) -> Dict:
download_id: str | None = None,
source: str | None = None,
file_params: Dict[str, Any] | None = None,
) -> Dict[str, Any]:
"""Download model from Civitai with task tracking and concurrency control
Args:
@@ -309,14 +313,14 @@ class DownloadManager:
async def _download_with_semaphore(
self,
task_id: str,
model_id: int,
model_version_id: int,
save_dir: str,
model_id: int | None,
model_version_id: int | None,
save_dir: str | None,
relative_path: str,
progress_callback=None,
use_default_paths: bool = False,
source: str = None,
file_params: Dict = None,
source: str | None = None,
file_params: Dict[str, Any] | None = None,
):
"""Execute download with semaphore to limit concurrency"""
# Update status to waiting
@@ -380,7 +384,8 @@ class DownloadManager:
# Use original download implementation
try:
# Check for cancellation before starting
if asyncio.current_task().cancelled():
current_task = asyncio.current_task()
if current_task is not None and current_task.cancelled():
raise asyncio.CancelledError()
result = await self._execute_original_download(
@@ -484,11 +489,11 @@ class DownloadManager:
# Schedule cleanup of download record after delay
asyncio.create_task(self._cleanup_download_record(task_id))
def _start_background_download_task(self, download_id: str, coroutine) -> asyncio.Task:
def _start_background_download_task(self, download_id: str, coroutine) -> asyncio.Task[Any]:
task = asyncio.create_task(coroutine)
self._download_tasks[download_id] = task
def _cleanup_done_task(done_task: asyncio.Task) -> None:
def _cleanup_done_task(done_task: asyncio.Task[Any]) -> None:
current_task = self._download_tasks.get(download_id)
if current_task is done_task:
self._download_tasks.pop(download_id, None)
@@ -530,7 +535,7 @@ class DownloadManager:
async def _cleanup_cancelled_download_files(
self,
download_id: str,
download_info: Optional[Dict],
download_info: Optional[Dict[str, Any]],
) -> None:
target_files = set()
persisted = await self._aria2_state_store.get(download_id)
@@ -603,13 +608,13 @@ class DownloadManager:
self,
download_id: str,
*,
extra: Optional[Dict] = None,
extra: Optional[Dict[str, Any]] = None,
) -> None:
info = self._active_downloads.get(download_id)
if not info:
return
payload = {
payload: Dict[str, Any] = {
"download_id": download_id,
"model_id": info.get("model_id"),
"model_version_id": info.get("model_version_id"),
@@ -631,7 +636,7 @@ class DownloadManager:
await self._aria2_state_store.upsert(download_id, payload)
def _build_restored_download_info(self, record: Dict, save_path: str) -> Dict:
def _build_restored_download_info(self, record: Dict[str, Any], save_path: str) -> Dict[str, Any]:
return {
"model_id": record.get("model_id"),
"model_version_id": record.get("model_version_id"),
@@ -653,8 +658,8 @@ class DownloadManager:
def _is_same_aria2_download_request(
self,
current_info: Optional[Dict],
persisted_record: Dict,
current_info: Optional[Dict[str, Any]],
persisted_record: Dict[str, Any],
) -> bool:
if not isinstance(current_info, dict):
return False
@@ -666,13 +671,15 @@ class DownloadManager:
return current_version_id == persisted_version_id
def _build_download_urls_from_file_info(self, file_info: Dict, source: str = None) -> List[str]:
def _build_download_urls_from_file_info(self, file_info: Dict[str, Any], source: str | None = None) -> List[str]:
mirrors = file_info.get("mirrors") or []
download_urls: List[str] = []
if mirrors:
for mirror in mirrors:
if mirror.get("deletedAt") is None and mirror.get("url"):
download_urls.append(normalize_civitai_download_url(mirror["url"]))
normalized_url = normalize_civitai_download_url(mirror["url"])
if normalized_url:
download_urls.append(normalized_url)
if source == "civarchive" and len(download_urls) > 1:
civitai_urls = [
@@ -688,7 +695,9 @@ class DownloadManager:
if not download_urls:
download_url = file_info.get("downloadUrl")
if download_url:
download_urls.append(normalize_civitai_download_url(download_url))
normalized_url = normalize_civitai_download_url(download_url)
if normalized_url:
download_urls.append(normalized_url)
return download_urls
@@ -696,8 +705,8 @@ class DownloadManager:
self,
*,
model_type: str,
version_info: Dict,
file_info: Dict,
version_info: Dict[str, Any],
file_info: Dict[str, Any],
save_path: str,
):
if model_type == "checkpoint":
@@ -706,7 +715,7 @@ class DownloadManager:
return EmbeddingMetadata.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) -> Optional[str]:
def _resolve_save_path_from_persisted_record(self, record: Dict[str, Any]) -> Optional[str]:
save_path = record.get("save_path") or record.get("file_path")
if isinstance(save_path, str) and save_path:
return os.path.abspath(save_path)
@@ -728,7 +737,7 @@ class DownloadManager:
return os.path.abspath(os.path.join(save_dir, file_name))
async def _resume_restored_aria2_download(self, download_id: str, record: Dict) -> Dict:
async def _resume_restored_aria2_download(self, download_id: str, record: Dict[str, Any]) -> Dict[str, Any]:
try:
if download_id in self._active_downloads:
self._active_downloads[download_id]["status"] = "downloading"
@@ -842,7 +851,7 @@ class DownloadManager:
self,
previous_download_id: str,
new_download_id: str,
persisted_record: Dict,
persisted_record: Dict[str, Any],
save_path: str,
) -> None:
aria2_downloader = await get_aria2_downloader()
@@ -938,7 +947,7 @@ class DownloadManager:
except Exception:
status_payload = None
if status_payload is not None:
if status_payload is not None and isinstance(gid, str):
remote_status = status_payload.get("status", "")
if remote_status in {"active", "waiting", "paused"}:
await aria2_downloader.restore_transfer(download_id, gid, save_path)
@@ -1115,17 +1124,17 @@ class DownloadManager:
async def _execute_original_download(
self,
model_id,
model_version_id,
save_dir,
relative_path,
model_id: int | None,
model_version_id: int | None,
save_dir: str | None,
relative_path: str,
progress_callback,
use_default_paths,
download_id=None,
transfer_backend="python",
source=None,
file_params=None,
):
use_default_paths: bool,
download_id: str | None = None,
transfer_backend: str = "python",
source: str | None = None,
file_params: Dict[str, Any] | None = None,
) -> Dict[str, Any]:
"""Wrapper for original download_from_civitai implementation"""
try:
# Check if model version already exists in library
@@ -1172,7 +1181,7 @@ class DownloadManager:
# Get version info based on the provided identifier
version_info = await metadata_provider.get_model_version(
model_id, model_version_id
cast(int, model_id), cast(int, model_version_id)
)
if not version_info:
@@ -1183,7 +1192,7 @@ class DownloadManager:
)
metadata_provider = await get_default_metadata_provider()
version_info = await metadata_provider.get_model_version(
model_id, model_version_id
cast(int, model_id), cast(int, model_version_id)
)
if not version_info:
@@ -1388,6 +1397,8 @@ class DownloadManager:
relative_path = self._calculate_relative_path(version_info, model_type)
# Update save directory with relative path if provided
if not save_dir:
return {"success": False, "error": "No save directory specified"}
if relative_path:
base_save_dir = save_dir
save_dir = os.path.join(save_dir, relative_path)
@@ -1561,6 +1572,11 @@ class DownloadManager:
version_info, file_info, save_path
)
logger.info(f"Creating EmbeddingMetadata for {file_name}")
else:
return {
"success": False,
"error": f'Unsupported model type "{model_type}"',
}
# 6. Start download process
if transfer_backend == "aria2" and download_id:
@@ -1580,7 +1596,7 @@ class DownloadManager:
},
)
execute_kwargs = {
execute_kwargs: Dict[str, Any] = {
"download_urls": download_urls,
"save_dir": save_dir,
"metadata": metadata,
@@ -1627,7 +1643,8 @@ class DownloadManager:
)
# If early_access_msg exists and download failed, replace error message
if "early_access_msg" in locals() and not result.get("success", False):
early_access_msg = locals().get("early_access_msg")
if early_access_msg and not result.get("success", False):
result["error"] = early_access_msg
return result
@@ -1652,7 +1669,7 @@ class DownloadManager:
self,
model_type: str,
model_id_value,
version_info: Dict,
version_info: Dict[str, Any],
fallback_version_id=None,
file_path: str | None = None,
) -> None:
@@ -1683,8 +1700,8 @@ class DownloadManager:
try:
await history_service.mark_downloaded(
model_type,
int(version_id),
model_id=int(resolved_model_id) if resolved_model_id is not None else None,
int(cast(Any, version_id)),
model_id=int(cast(Any, resolved_model_id)) if resolved_model_id is not None else None,
source="download",
file_path=file_path,
)
@@ -1701,7 +1718,7 @@ class DownloadManager:
self,
model_type: str,
model_id_value,
version_info: Dict,
version_info: Dict[str, Any],
fallback_version_id=None,
) -> None:
"""Ensure update tracking reflects a newly downloaded version."""
@@ -1725,7 +1742,7 @@ class DownloadManager:
if isinstance(model_info, dict):
resolved_model_id = model_info.get("id")
try:
resolved_model_id = int(resolved_model_id)
resolved_model_id = int(cast(Any, resolved_model_id))
except (TypeError, ValueError):
logger.debug(
"Skipping update sync; invalid model id: %s", resolved_model_id
@@ -1736,7 +1753,7 @@ class DownloadManager:
if version_id is None:
version_id = fallback_version_id
try:
version_id = int(version_id)
version_id = int(cast(Any, version_id))
except (TypeError, ValueError):
logger.debug(
"Skipping update sync; invalid version id for model %s: %s",
@@ -1773,7 +1790,7 @@ class DownloadManager:
for entry in local_versions or []:
vid = entry.get("versionId")
try:
version_ids.add(int(vid))
version_ids.add(int(cast(Any, vid)))
except (TypeError, ValueError):
continue
@@ -1795,7 +1812,7 @@ class DownloadManager:
)
def _calculate_relative_path(
self, version_info: Dict, model_type: str = "lora"
self, version_info: Dict[str, Any], model_type: str = "lora"
) -> str:
"""Calculate relative path using template from settings
@@ -1871,21 +1888,22 @@ class DownloadManager:
download_urls: List[str],
save_dir: str,
metadata,
version_info: Dict,
version_info: Dict[str, Any],
relative_path: str,
progress_callback=None,
model_type: str = "lora",
download_id: str = None,
download_id: str | None = None,
transfer_backend: Optional[str] = None,
) -> Dict:
) -> Dict[str, Any]:
"""Execute the actual download process including preview images and model files"""
metadata_entries: List = []
metadata_entries: List[Any] = []
metadata_files_for_cleanup: List[str] = []
extracted_paths: List[str] = []
metadata_path = ""
preview_targets: List[str] = []
preview_path: str | None = None
preview_nsfw_level = 0
save_path: str | None = None
transfer_backend = (transfer_backend or self._get_model_download_backend()).lower()
try:
resolved, save_path = await self._resolve_download_target_path(
@@ -1933,9 +1951,9 @@ class DownloadManager:
mature_threshold=mature_threshold,
)
preview_url = selected_image.get("url") if selected_image else None
preview_url = cast(Optional[str], selected_image.get("url")) if selected_image else None
media_type = (
(selected_image.get("type") or "").lower() if selected_image else ""
cast(str, selected_image.get("type") or "").lower() if selected_image else ""
)
def _extension_from_url(url: str, fallback: str) -> str:
@@ -1959,9 +1977,10 @@ class DownloadManager:
preview_url, media_type="video"
)
attempt_urls: List[str] = []
if rewritten:
if rewritten and rewritten_url:
attempt_urls.append(rewritten_url)
attempt_urls.append(preview_url)
if preview_url:
attempt_urls.append(preview_url)
seen_attempts = set()
for attempt in attempt_urls:
@@ -1978,7 +1997,7 @@ class DownloadManager:
rewritten_url, rewritten = rewrite_preview_url(
preview_url, media_type="image"
)
if rewritten:
if rewritten and rewritten_url:
preview_ext = _extension_from_url(preview_url, ".png")
preview_path = os.path.splitext(save_path)[0] + preview_ext
success, _ = await downloader.download_file(
@@ -2004,7 +2023,9 @@ class DownloadManager:
)
if success:
with open(temp_path, "wb") as temp_file_handle:
temp_file_handle.write(content)
temp_file_handle.write(
content if isinstance(content, bytes) else content.encode("utf-8")
)
preview_path = (
os.path.splitext(save_path)[0] + ".webp"
)
@@ -2056,6 +2077,8 @@ class DownloadManager:
last_error = None
for download_url in download_urls:
download_url = normalize_civitai_download_url(download_url)
if download_url is None:
continue
use_auth = download_url.startswith(CIVITAI_DOWNLOAD_URL_PREFIXES)
if transfer_backend == "aria2" and download_id:
await self._persist_aria2_state(
@@ -2239,7 +2262,7 @@ class DownloadManager:
entry, normalized_file_path, adjust_root
)
if adjusted_entry is not None:
entry = adjusted_entry
entry = cast(Any, adjusted_entry)
metadata_entries[index] = entry
metadata_file_path = (
@@ -2359,11 +2382,11 @@ class DownloadManager:
async def _build_metadata_entries(
self, base_metadata, file_paths: List[str]
) -> List:
) -> List[Any]:
if not file_paths:
return []
entries: List = []
entries: List[Any] = []
for index, file_path in enumerate(file_paths):
entry = base_metadata if index == 0 else copy.deepcopy(base_metadata)
# Update file paths without modifying size and modified timestamps
@@ -2406,7 +2429,7 @@ class DownloadManager:
return destination
def _distribute_preview_to_entries(
self, preview_path: str, entries: List
self, preview_path: str, entries: List[Any]
) -> List[str]:
if not preview_path or not entries:
return []
@@ -2465,7 +2488,7 @@ class DownloadManager:
progress_callback, normalized_snapshot, rounded_progress
)
async def cancel_download(self, download_id: str) -> Dict:
async def cancel_download(self, download_id: str) -> Dict[str, Any]:
"""Cancel an active download by download_id
Args:
@@ -2547,7 +2570,7 @@ class DownloadManager:
self._download_tasks.pop(download_id, None)
await self._aria2_state_store.remove(download_id)
async def skip_download(self, download_id: str) -> Dict:
async def skip_download(self, download_id: str) -> Dict[str, Any]:
"""Skip a download while preserving all partial files on disk.
Removes all in-memory tracking (asyncio task, semaphore, active/pause
@@ -2630,7 +2653,7 @@ class DownloadManager:
# Preserve aria2 state store entry so the partial download
# info survives restarts and can be resumed later
async def pause_download(self, download_id: str) -> Dict:
async def pause_download(self, download_id: str) -> Dict[str, Any]:
"""Pause an active download without losing progress."""
await self._restore_persisted_downloads()
@@ -2677,7 +2700,7 @@ class DownloadManager:
return {"success": True, "message": "Download paused successfully"}
async def resume_download(self, download_id: str) -> Dict:
async def resume_download(self, download_id: str) -> Dict[str, Any]:
"""Resume a previously paused download."""
await self._restore_persisted_downloads()
@@ -2694,7 +2717,7 @@ class DownloadManager:
self._pause_events[download_id] = pause_control
self._active_downloads[download_id] = self._build_restored_download_info(
persisted,
os.path.abspath(save_path),
os.path.abspath(cast(str, save_path)),
)
if pause_control.is_set():
@@ -2821,7 +2844,7 @@ class DownloadManager:
elif asyncio.iscoroutine(result):
await result
async def get_active_downloads(self) -> Dict:
async def get_active_downloads(self) -> Dict[str, Any]:
"""Get information about all active downloads
Returns:

4
py/services/downloaded_version_history_service.py
View File

@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
from __future__ import annotations
import asyncio

21
py/services/downloader.py
View File

@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
"""
Unified download manager for all HTTP/HTTPS downloads in the application.
@@ -20,7 +24,7 @@ from dataclasses import dataclass
from datetime import datetime, timedelta
from email.utils import parsedate_to_datetime
from urllib.parse import urlparse
from typing import Optional, Dict, Tuple, Callable, Union, Awaitable
from typing import Optional, Dict, Tuple, Callable, Union, Awaitable, Any, cast
from ..services.settings_manager import get_settings_manager
from .connectivity_guard import (
OFFLINE_COOLDOWN_ERROR,
@@ -204,6 +208,7 @@ class Downloader:
# Double check after acquiring lock
if self._session is None or self._should_refresh_session():
await self._create_session()
assert self._session is not None
return self._session
@property
@@ -231,7 +236,7 @@ class Downloader:
)
try:
timeout_value = float(raw_value)
timeout_value = float(cast(Any, raw_value))
except (TypeError, ValueError):
timeout_value = default_timeout
@@ -243,7 +248,7 @@ class Downloader:
raw_value = os.environ.get("COMFYUI_DOWNLOAD_MAX_RETRIES")
try:
retries = int(raw_value)
retries = int(cast(Any, raw_value))
except (TypeError, ValueError):
retries = default_retries
@@ -320,7 +325,7 @@ class Downloader:
# CA coverage across different Python environments (especially
# embedded/compatibility Python builds).
try:
import certifi # type: ignore[import-untyped]
import certifi # pyright: ignore[reportMissingTypeStubs]
ca_path = certifi.where()
ssl_context = ssl.create_default_context(cafile=ca_path)
@@ -330,7 +335,7 @@ class Downloader:
logger.debug("SSL: certifi unavailable; using system default CA bundle")
# Optimize TCP connection parameters
connector_kwargs = dict(
connector_kwargs: Dict[str, Any] = dict(
ssl=ssl_context,
limit=8, # Concurrent connections
ttl_dns_cache=300, # DNS cache timeout
@@ -890,7 +895,7 @@ class Downloader:
use_auth: bool = False,
custom_headers: Optional[Dict[str, str]] = None,
return_headers: bool = False,
) -> Tuple[bool, Union[bytes, str], Optional[Dict]]:
) -> Tuple[bool, Union[bytes, str], Optional[Dict[str, Any]]]:
"""
Download a file to memory (for small files like preview images)
@@ -976,7 +981,7 @@ class Downloader:
url: str,
use_auth: bool = False,
custom_headers: Optional[Dict[str, str]] = None,
) -> Tuple[bool, Union[Dict, str]]:
) -> Tuple[bool, Union[Dict[str, Any], str]]:
"""
Get response headers without downloading the full content
@@ -1036,7 +1041,7 @@ class Downloader:
use_auth: bool = False,
custom_headers: Optional[Dict[str, str]] = None,
**kwargs,
) -> Tuple[bool, Union[Dict, str]]:
) -> Tuple[bool, Union[Dict[str, Any], str, RateLimitError]]:
"""
Make a generic HTTP request and return JSON response

2
py/services/embedding_scanner.py
View File

@@ -27,7 +27,7 @@ class EmbeddingScanner(ModelScanner):
roots.extend(config.embeddings_roots or [])
roots.extend(config.extra_embeddings_roots or [])
# Remove duplicates while preserving order
seen: set = set()
seen: set[str] = set()
unique_roots: List[str] = []
for root in roots:
if root and root not in seen:

56
py/services/embedding_service.py
View File

@@ -1,6 +1,6 @@
import os
import logging
from typing import Dict, Optional
from typing import Any, Dict, Optional
from .base_model_service import BaseModelService
from .auto_tag_service import extract_auto_tags
@@ -21,58 +21,58 @@ class EmbeddingService(BaseModelService):
"""
super().__init__("embedding", scanner, EmbeddingMetadata, update_service=update_service)
async def format_response(self, embedding_data: Dict) -> Optional[Dict]:
async def format_response(self, model_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""Format Embedding data for API response.
Returns None when the entry is missing critical fields (corrupted cache
row), so the handler layer can filter it out. See issue #730.
"""
# Guard against corrupted cache entries missing critical fields
file_path = embedding_data.get("file_path")
file_path = model_data.get("file_path")
if not file_path or not isinstance(file_path, str):
logger.warning(
"Skipping corrupted embedding entry (missing file_path): %s",
embedding_data.get("file_name", "<unknown>"),
model_data.get("file_name", "<unknown>"),
)
return None
# Get sub_type from cache entry (new canonical field)
sub_type = embedding_data.get("sub_type", "embedding")
sub_type = model_data.get("sub_type", "embedding")
file_name = embedding_data.get("file_name") or ""
model_name = embedding_data.get("model_name") or file_name
folder = embedding_data.get("folder") or ""
file_name = model_data.get("file_name") or ""
model_name = model_data.get("model_name") or file_name
folder = model_data.get("folder") or ""
return {
"model_name": model_name,
"file_name": file_name,
"preview_url": config.get_preview_static_url(embedding_data.get("preview_url", "")),
"preview_nsfw_level": embedding_data.get("preview_nsfw_level", 0),
"base_model": embedding_data.get("base_model", ""),
"preview_url": config.get_preview_static_url(model_data.get("preview_url", "")),
"preview_nsfw_level": model_data.get("preview_nsfw_level", 0),
"base_model": model_data.get("base_model", ""),
"folder": folder,
"sha256": embedding_data.get("sha256", ""),
"sha256": model_data.get("sha256", ""),
"file_path": file_path.replace(os.sep, "/"),
"file_size": embedding_data.get("size", 0),
"modified": embedding_data.get("modified", ""),
"tags": embedding_data.get("tags", []),
"from_civitai": embedding_data.get("from_civitai", True),
# "usage_count": embedding_data.get("usage_count", 0), # TODO: Enable when embedding usage tracking is implemented
"notes": embedding_data.get("notes", ""),
"file_size": model_data.get("size", 0),
"modified": model_data.get("modified", ""),
"tags": model_data.get("tags", []),
"from_civitai": model_data.get("from_civitai", True),
# "usage_count": model_data.get("usage_count", 0), # TODO: Enable when embedding usage tracking is implemented
"notes": model_data.get("notes", ""),
"sub_type": sub_type,
"favorite": embedding_data.get("favorite", False),
"exclude": bool(embedding_data.get("exclude", False)),
"update_available": bool(embedding_data.get("update_available", False)),
"skip_metadata_refresh": bool(embedding_data.get("skip_metadata_refresh", False)),
"civitai": self.filter_civitai_data(embedding_data.get("civitai", {}), minimal=True),
"auto_tags": embedding_data.get("auto_tags") or extract_auto_tags(embedding_data),
"version_count": embedding_data.get("version_count"),
"hf_url": embedding_data.get("hf_url", ""),
"favorite": model_data.get("favorite", False),
"exclude": bool(model_data.get("exclude", False)),
"update_available": bool(model_data.get("update_available", False)),
"skip_metadata_refresh": bool(model_data.get("skip_metadata_refresh", False)),
"civitai": self.filter_civitai_data(model_data.get("civitai", {}), minimal=True),
"auto_tags": model_data.get("auto_tags") or extract_auto_tags(model_data),
"version_count": model_data.get("version_count"),
"hf_url": model_data.get("hf_url", ""),
}
def find_duplicate_hashes(self) -> Dict:
def find_duplicate_hashes(self) -> Dict[str, Any]:
"""Find Embeddings with duplicate SHA256 hashes"""
return self.scanner._hash_index.get_duplicate_hashes()
def find_duplicate_filenames(self) -> Dict:
def find_duplicate_filenames(self) -> Dict[str, Any]:
"""Find Embeddings with conflicting filenames"""
return self.scanner._hash_index.get_duplicate_filenames()

2
py/services/example_images_cleanup_service.py
View File

@@ -35,7 +35,7 @@ class CleanupResult:
def to_dict(self) -> Dict[str, object]:
"""Convert the dataclass to a serialisable dictionary."""
data = {
data: Dict[str, object] = {
"success": self.success,
"checked_folders": self.checked_folders,
"moved_empty_folders": self.moved_empty_folders,

18
py/services/lora_scanner.py
View File

@@ -1,10 +1,12 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import logging
from typing import List
from ..utils.models import LoraMetadata
from ..config import config
from .model_scanner import ModelScanner
from .model_hash_index import ModelHashIndex # Changed from LoraHashIndex to ModelHashIndex
import sys
logger = logging.getLogger(__name__)
@@ -15,8 +17,10 @@ class LoraScanner(ModelScanner):
def __init__(self):
# Define supported file extensions
file_extensions = {'.safetensors'}
# Initialize parent class with ModelHashIndex
from .model_hash_index import ModelHashIndex
super().__init__(
model_type="lora",
model_class=LoraMetadata,
@@ -26,11 +30,13 @@ class LoraScanner(ModelScanner):
def get_model_roots(self) -> List[str]:
"""Get lora root directories (including extra paths)"""
from ..config import config
roots: List[str] = []
roots.extend(config.loras_roots or [])
roots.extend(config.extra_loras_roots or [])
# Remove duplicates while preserving order
seen: set = set()
seen: set[str] = set()
unique_roots: List[str] = []
for root in roots:
if root and root not in seen:
@@ -68,8 +74,12 @@ class LoraScanner(ModelScanner):
test_hash = next(iter(self._hash_index._hash_to_path.keys()))
test_path = self._hash_index.get_path(test_hash)
logger.debug(f"\nTest lookup by hash: {test_hash[:8]}... -> {test_path}")
if test_path is None:
return
# Also test reverse lookup
test_hash_result = self._hash_index.get_hash(test_path)
if test_hash_result is None:
return
logger.debug(f"Test reverse lookup: {test_path} -> {test_hash_result[:8]}...\n\n")

82
py/services/lora_service.py
View File

@@ -1,7 +1,7 @@
import logging
import json
import os
from typing import Dict, List, Optional
from typing import Any, Dict, List, Optional
from .base_model_service import BaseModelService
from .model_query import resolve_sub_type
@@ -24,7 +24,7 @@ class LoraService(BaseModelService):
"""
super().__init__("lora", scanner, LoraMetadata, update_service=update_service)
async def format_response(self, lora_data: Dict) -> Optional[Dict]:
async def format_response(self, model_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""Format LoRA data for API response.
Returns None when the entry is missing critical fields (corrupted cache
@@ -32,56 +32,56 @@ class LoraService(BaseModelService):
whole listing request. See issue #730.
"""
# Guard against corrupted cache entries missing critical fields
file_path = lora_data.get("file_path")
file_path = model_data.get("file_path")
if not file_path or not isinstance(file_path, str):
logger.warning(
"Skipping corrupted LoRA entry (missing file_path): %s",
lora_data.get("file_name", "<unknown>"),
model_data.get("file_name", "<unknown>"),
)
return None
# Resolve sub_type using priority: sub_type > model_type > civitai.model.type > default
# Normalize to lowercase for consistent API responses
sub_type = resolve_sub_type(lora_data).lower()
sub_type = resolve_sub_type(model_data).lower()
file_name = lora_data.get("file_name") or ""
model_name = lora_data.get("model_name") or file_name
folder = lora_data.get("folder") or ""
file_name = model_data.get("file_name") or ""
model_name = model_data.get("model_name") or file_name
folder = model_data.get("folder") or ""
return {
"model_name": model_name,
"file_name": file_name,
"preview_url": config.get_preview_static_url(
lora_data.get("preview_url", "")
model_data.get("preview_url", "")
),
"preview_nsfw_level": lora_data.get("preview_nsfw_level", 0),
"base_model": lora_data.get("base_model", ""),
"preview_nsfw_level": model_data.get("preview_nsfw_level", 0),
"base_model": model_data.get("base_model", ""),
"folder": folder,
"sha256": lora_data.get("sha256", ""),
"sha256": model_data.get("sha256", ""),
"file_path": file_path.replace(os.sep, "/"),
"file_size": lora_data.get("size", 0),
"modified": lora_data.get("modified", ""),
"tags": lora_data.get("tags", []),
"from_civitai": lora_data.get("from_civitai", True),
"usage_count": lora_data.get("usage_count", 0),
"usage_tips": lora_data.get("usage_tips", ""),
"notes": lora_data.get("notes", ""),
"favorite": lora_data.get("favorite", False),
"exclude": bool(lora_data.get("exclude", False)),
"update_available": bool(lora_data.get("update_available", False)),
"file_size": model_data.get("size", 0),
"modified": model_data.get("modified", ""),
"tags": model_data.get("tags", []),
"from_civitai": model_data.get("from_civitai", True),
"usage_count": model_data.get("usage_count", 0),
"usage_tips": model_data.get("usage_tips", ""),
"notes": model_data.get("notes", ""),
"favorite": model_data.get("favorite", False),
"exclude": bool(model_data.get("exclude", False)),
"update_available": bool(model_data.get("update_available", False)),
"skip_metadata_refresh": bool(
lora_data.get("skip_metadata_refresh", False)
model_data.get("skip_metadata_refresh", False)
),
"sub_type": sub_type,
"civitai": self.filter_civitai_data(
lora_data.get("civitai", {}), minimal=True
model_data.get("civitai", {}), minimal=True
),
"auto_tags": lora_data.get("auto_tags") or extract_auto_tags(lora_data),
"version_count": lora_data.get("version_count"),
"hf_url": lora_data.get("hf_url", ""),
"auto_tags": model_data.get("auto_tags") or extract_auto_tags(model_data),
"version_count": model_data.get("version_count"),
"hf_url": model_data.get("hf_url", ""),
}
async def _apply_specific_filters(self, data: List[Dict], **kwargs) -> List[Dict]:
async def _apply_specific_filters(self, data: List[Dict[str, Any]], **kwargs) -> List[Dict[str, Any]]:
"""Apply LoRA-specific filters"""
# Handle first_letter filter for LoRAs
first_letter = kwargs.get("first_letter")
@@ -152,7 +152,7 @@ class LoraService(BaseModelService):
return data
def _filter_by_first_letter(self, data: List[Dict], letter: str) -> List[Dict]:
def _filter_by_first_letter(self, data: List[Dict[str, Any]], letter: str) -> List[Dict[str, Any]]:
"""Filter data by first letter of model name
Special handling:
@@ -307,7 +307,7 @@ class LoraService(BaseModelService):
return None
@staticmethod
def get_recommended_strength_from_lora_data(lora_data: Dict) -> Optional[float]:
def get_recommended_strength_from_lora_data(lora_data: Dict[str, Any]) -> Optional[float]:
"""Parse usage_tips JSON and extract recommended model strength."""
try:
usage_tips = lora_data.get("usage_tips", "")
@@ -320,7 +320,7 @@ class LoraService(BaseModelService):
@staticmethod
def get_recommended_clip_strength_from_lora_data(
lora_data: Dict,
lora_data: Dict[str, Any],
) -> Optional[float]:
"""Parse usage_tips JSON and extract recommended clip strength."""
try:
@@ -332,7 +332,7 @@ class LoraService(BaseModelService):
except (json.JSONDecodeError, TypeError, AttributeError):
return None
async def get_lora_metadata_by_filename(self, filename: str) -> Optional[Dict]:
async def get_lora_metadata_by_filename(self, filename: str) -> Optional[Dict[str, Any]]:
"""Return cached raw metadata for a LoRA matching the given filename."""
cache = await self.scanner.get_cached_data(force_refresh=False)
@@ -357,11 +357,11 @@ class LoraService(BaseModelService):
return None
def find_duplicate_hashes(self) -> Dict:
def find_duplicate_hashes(self) -> Dict[str, Any]:
"""Find LoRAs with duplicate SHA256 hashes"""
return self.scanner._hash_index.get_duplicate_hashes()
def find_duplicate_filenames(self) -> Dict:
def find_duplicate_filenames(self) -> Dict[str, Any]:
"""Find LoRAs with conflicting filenames"""
return self.scanner._hash_index.get_duplicate_filenames()
@@ -373,8 +373,8 @@ class LoraService(BaseModelService):
use_same_clip_strength: bool = True,
clip_strength_min: float = 0.0,
clip_strength_max: float = 1.0,
locked_loras: Optional[List[Dict]] = None,
pool_config: Optional[Dict] = None,
locked_loras: Optional[List[Dict[str, Any]]] = None,
pool_config: Optional[Dict[str, Any]] = None,
count_mode: str = "fixed",
count_min: int = 3,
count_max: int = 7,
@@ -382,7 +382,7 @@ class LoraService(BaseModelService):
recommended_strength_scale_min: float = 0.5,
recommended_strength_scale_max: float = 1.0,
seed: Optional[int] = None,
) -> List[Dict]:
) -> List[Dict[str, Any]]:
"""
Get random LoRAs with specified strength ranges.
@@ -513,8 +513,8 @@ class LoraService(BaseModelService):
return result_loras
async def _apply_pool_filters(
self, available_loras: List[Dict], pool_config: Dict
) -> List[Dict]:
self, available_loras: List[Dict[str, Any]], pool_config: Dict[str, Any]
) -> List[Dict[str, Any]]:
"""
Apply pool_config filters to available LoRAs.
@@ -671,8 +671,8 @@ class LoraService(BaseModelService):
return available_loras
async def get_cycler_list(
self, pool_config: Optional[Dict] = None, sort_by: str = "filename"
) -> List[Dict]:
self, pool_config: Optional[Dict[str, Any]] = None, sort_by: str = "filename"
) -> List[Dict[str, Any]]:
"""
Get filtered and sorted LoRA list for cycling.

6
py/services/metadata_service.py
View File

@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import os
import logging
from .model_metadata_provider import (
@@ -170,7 +174,7 @@ def _wrap_provider_with_rate_limit(provider_name: str | None, provider: ModelMet
return RateLimitRetryingProvider(provider, label=provider_name)
async def get_metadata_provider(provider_name: str = None):
async def get_metadata_provider(provider_name: str | None = None):
"""Get a specific metadata provider or default provider with rate-limit handling."""
provider_manager = await ModelMetadataProviderManager.get_instance()

14
py/services/metadata_sync_service.py
View File

@@ -6,7 +6,7 @@ import json
import logging
import os
from datetime import datetime
from typing import Any, Awaitable, Callable, Dict, Iterable, Optional
from typing import Any, Awaitable, Callable, Dict, Iterable, Optional, Protocol
from ..services.settings_manager import SettingsManager
from ..utils.civitai_utils import resolve_license_payload
@@ -18,14 +18,14 @@ from .errors import RateLimitError
logger = logging.getLogger(__name__)
class MetadataProviderProtocol:
class MetadataProviderProtocol(Protocol):
"""Subset of metadata provider interface consumed by the sync service."""
async def get_model_by_hash(self, sha256: str) -> tuple[Optional[Dict[str, Any]], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> tuple[Optional[Dict[str, Any]], Optional[str]]:
...
async def get_model_version(
self, model_id: int, model_version_id: Optional[int]
self, model_id: Any = None, version_id: Any = None
) -> Optional[Dict[str, Any]]:
...
@@ -39,8 +39,8 @@ class MetadataSyncService:
metadata_manager,
preview_service,
settings: SettingsManager,
default_metadata_provider_factory: Callable[[], Awaitable[MetadataProviderProtocol]],
metadata_provider_selector: Callable[[str], Awaitable[MetadataProviderProtocol]],
default_metadata_provider_factory: Callable[..., Awaitable[MetadataProviderProtocol]],
metadata_provider_selector: Callable[..., Awaitable[MetadataProviderProtocol]],
) -> None:
self._metadata_manager = metadata_manager
self._preview_service = preview_service
@@ -492,7 +492,7 @@ class MetadataSyncService:
if not file_paths:
raise ValueError("No file paths provided for verification")
results = {
results: Dict[str, Any] = {
"verified_as_duplicates": True,
"mismatched_files": [],
"new_hash_map": {},

37
py/services/model_cache.py
View File

@@ -31,17 +31,22 @@ DISPLAY_NAME_MODES = {"model_name", "file_name"}
class ModelCache:
"""Cache structure for model data with extensible sorting."""
raw_data: List[Dict]
raw_data: List[Dict[str, Any]]
folders: List[str]
version_index: Dict[int, Dict] = field(default_factory=dict)
version_index: Dict[int, Dict[str, Any]] = field(default_factory=dict)
model_id_index: Dict[int, List[Dict[str, Any]]] = field(default_factory=dict)
name_display_mode: str = "model_name"
_lock: Any = field(init=False, repr=False, default=None)
# Cache for last sort: (sort_key, order, seed) -> sorted list
_last_sort: Tuple[Optional[str], str, Optional[str]] = field(
init=False, repr=False, default=(None, "asc", None)
)
_last_sorted_data: List[Dict[str, Any]] = field(
init=False, repr=False, default_factory=list
)
def __post_init__(self):
self._lock = asyncio.Lock()
# Cache for last sort: (sort_key, order, seed) -> sorted list
self._last_sort: Tuple[Optional[str], str, Optional[str]] = (None, "asc", None)
self._last_sorted_data: List[Dict] = []
self._normalize_raw_data()
self.name_display_mode = self._normalize_display_mode(self.name_display_mode)
# Default sort on init
@@ -64,7 +69,7 @@ class ModelCache:
return ""
return str(value)
def _normalize_item(self, item: Dict) -> None:
def _normalize_item(self, item: Dict[str, Any]) -> None:
"""Ensure core metadata fields are present and string typed."""
if not isinstance(item, dict):
@@ -80,7 +85,7 @@ class ModelCache:
for item in self.raw_data:
self._normalize_item(item)
def _get_display_name(self, item: Dict) -> str:
def _get_display_name(self, item: Dict[str, Any]) -> str:
"""Return the value used for name-based sorting based on display settings."""
if self.name_display_mode == "file_name":
@@ -114,7 +119,7 @@ class ModelCache:
for item in self.raw_data:
self.add_to_version_index(item)
def add_to_version_index(self, item: Dict) -> None:
def add_to_version_index(self, item: Dict[str, Any]) -> None:
"""Register a cache item in the version/model indexes if possible."""
civitai_data = item.get('civitai') if isinstance(item, dict) else None
@@ -143,7 +148,7 @@ class ModelCache:
else:
versions.append(descriptor)
def remove_from_version_index(self, item: Dict) -> None:
def remove_from_version_index(self, item: Dict[str, Any]) -> None:
"""Remove a cache item from the version/model indexes if present."""
civitai_data = item.get('civitai') if isinstance(item, dict) else None
@@ -177,7 +182,7 @@ class ModelCache:
def _build_version_descriptor(
self,
item: Dict,
item: Dict[str, Any],
civitai_data: Dict[str, Any],
version_id: int,
) -> Optional[Dict[str, Any]]:
@@ -204,8 +209,8 @@ class ModelCache:
async def resort(self):
"""Resort cached data according to last sort mode if set"""
async with self._lock:
if self._last_sort[0] is not None:
sort_key, order, seed = self._last_sort
sort_key, order, seed = self._last_sort
if sort_key is not None:
sorted_data = self._sort_data(self.raw_data, sort_key, order, seed)
self._last_sorted_data = sorted_data
# Update folder list
@@ -219,7 +224,7 @@ class ModelCache:
self.folders = sorted(list(all_folders), key=lambda x: x.lower())
self.rebuild_version_index()
def _sort_data(self, data: List[Dict], sort_key: str, order: str, seed: Optional[str] = None) -> List[Dict]:
def _sort_data(self, data: List[Dict[str, Any]], sort_key: str, order: str, seed: Optional[str] = None) -> List[Dict[str, Any]]:
"""Sort data by sort_key and order"""
start_time = time.perf_counter()
reverse = (order == 'desc')
@@ -293,7 +298,7 @@ class ModelCache:
logger.debug("ModelCache._sort_data(%s, %s) for %d items took %.3fs", sort_key, order, len(data), duration)
return result
async def get_sorted_data(self, sort_key: str = 'name', order: str = 'asc', seed: Optional[str] = None) -> List[Dict]:
async def get_sorted_data(self, sort_key: str = 'name', order: str = 'asc', seed: Optional[str] = None) -> List[Dict[str, Any]]:
"""Get sorted data by sort_key and order, using cache if possible"""
async with self._lock:
cache_key = (sort_key, order, seed)
@@ -321,8 +326,8 @@ class ModelCache:
self.name_display_mode = normalized
if self._last_sort[0] == 'name':
sort_key, order, seed = self._last_sort
sort_key, order, seed = self._last_sort
if sort_key == 'name':
self._last_sorted_data = self._sort_data(self.raw_data, sort_key, order, seed)
async def update_preview_url(self, file_path: str, preview_url: str, preview_nsfw_level: int) -> bool:

4
py/services/model_file_service.py
View File

@@ -41,7 +41,7 @@ class AutoOrganizeResult:
def to_dict(self) -> Dict[str, Any]:
"""Convert result to dictionary"""
result = {
result: Dict[str, Any] = {
'success': self.status != 'error',
'status': self.status,
'message': f'Auto-organize {self.operation_type} completed: {self.success_count} moved, {self.skipped_count} skipped, {self.failure_count} failed out of {self.total} total',
@@ -418,6 +418,8 @@ class ModelFileService:
"""Calculate the target directory for a model"""
if is_flat_structure:
file_path = model.get('file_path')
if not isinstance(file_path, str):
return None
current_dir = os.path.dirname(file_path)
# Check if already in root directory

3
py/services/model_hash_index.py
View File

@@ -35,6 +35,7 @@ class ModelHashIndex:
# Track duplicates by filename - FIXED LOGIC
is_re_registration = False
existing_hash: Optional[str] = None
if filename in self._filename_to_hash:
existing_hash = self._filename_to_hash[filename]
existing_path = self._hash_to_path.get(existing_hash)
@@ -101,7 +102,7 @@ class ModelHashIndex:
"""Extract filename without extension from path"""
return os.path.splitext(os.path.basename(file_path))[0]
def remove_by_path(self, file_path: str, hash_val: str = None) -> None:
def remove_by_path(self, file_path: str, hash_val: Optional[str] = None) -> None:
"""Remove entry by file path"""
filename = self._get_filename_from_path(file_path)

13
py/services/model_lifecycle_service.py
View File

@@ -4,7 +4,7 @@ from __future__ import annotations
import logging
import os
from typing import Any, Awaitable, Callable, Dict, Iterable, List, Mapping, Optional, TYPE_CHECKING
from typing import Any, Awaitable, Callable, Dict, Iterable, List, Mapping, Optional, TYPE_CHECKING, cast
from ..services.service_registry import ServiceRegistry
from ..utils.constants import PREVIEW_EXTENSIONS
@@ -87,8 +87,8 @@ class ModelLifecycleService:
scanner,
metadata_manager,
metadata_loader: Callable[[str], Awaitable[Dict[str, object]]],
recipe_scanner_factory: Callable[[], Awaitable] | None = None,
update_service: "ModelUpdateService" | None = None,
recipe_scanner_factory: Callable[[], Awaitable[Any]] | None = None,
update_service: Optional["ModelUpdateService"] = None,
) -> None:
self._scanner = scanner
self._metadata_manager = metadata_manager
@@ -146,7 +146,7 @@ class ModelLifecycleService:
persist_current_cache = getattr(self._scanner, "_persist_current_cache", None)
if callable(persist_current_cache):
await persist_current_cache()
await cast(Awaitable[Any], persist_current_cache())
return {"success": True, "deleted_files": deleted_files}
@@ -252,7 +252,7 @@ class ModelLifecycleService:
persist_current_cache = getattr(self._scanner, "_persist_current_cache", None)
if callable(persist_current_cache):
await persist_current_cache()
await cast(Awaitable[Any], persist_current_cache())
message = f"Model {os.path.basename(file_path)} excluded"
return {"success": True, "message": message}
@@ -357,7 +357,8 @@ class ModelLifecycleService:
if os.path.exists(metadata_path):
metadata = await self._metadata_loader(metadata_path)
hash_value = metadata.get("sha256") if isinstance(metadata, dict) else None
raw_hash = metadata.get("sha256") if isinstance(metadata, dict) else None
hash_value = raw_hash if isinstance(raw_hash, str) else None
renamed_files: List[str] = []
new_metadata_path: Optional[str] = None

104
py/services/model_metadata_provider.py
View File

@@ -10,7 +10,7 @@ from .errors import RateLimitError, ResourceNotFoundError
try:
from bs4 import BeautifulSoup
except ImportError as exc:
BeautifulSoup = None # type: ignore[assignment]
BeautifulSoup = None # pyright: ignore[reportAssignmentType]
_BS4_IMPORT_ERROR = exc
else:
_BS4_IMPORT_ERROR = None
@@ -18,7 +18,7 @@ else:
try:
import aiosqlite
except ImportError as exc:
aiosqlite = None # type: ignore[assignment]
aiosqlite = None # pyright: ignore[reportAssignmentType]
_AIOSQLITE_IMPORT_ERROR = exc
else:
_AIOSQLITE_IMPORT_ERROR = None
@@ -105,24 +105,24 @@ class ModelMetadataProvider(ABC):
"""Base abstract class for all model metadata providers"""
@abstractmethod
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Find model by hash value"""
pass
@abstractmethod
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
"""Get all versions of a model with their details"""
pass
async def get_model_versions_bulk(
self, model_ids: Sequence[int]
) -> Optional[Dict[int, Dict]]:
) -> Optional[Dict[int, Dict[str, Any]]]:
"""Fetch model versions for multiple model ids when supported."""
raise NotImplementedError
async def get_model_versions_by_hashes(
self, hashes: List[str]
) -> Optional[List[Dict]]:
) -> Optional[List[Dict[str, Any]]]:
"""Fetch full version details for multiple SHA256 hashes.
Used specifically to retrieve ``usageControl`` which is only
@@ -133,17 +133,17 @@ class ModelMetadataProvider(ABC):
raise NotImplementedError
@abstractmethod
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
"""Get specific model version with additional metadata"""
pass
@abstractmethod
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Fetch model version metadata"""
pass
@abstractmethod
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict]:
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
"""Fetch one page of models owned by the specified user.
Returns ``{"items": [...], "nextCursor": <str|None>}`` on success,
@@ -161,29 +161,29 @@ class CivitaiModelMetadataProvider(ModelMetadataProvider):
def __init__(self, civitai_client):
self.client = civitai_client
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
return await self.client.get_model_by_hash(model_hash)
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
return await self.client.get_model_versions(model_id)
async def get_model_versions_bulk(
self, model_ids: Sequence[int]
) -> Optional[Dict[int, Dict]]:
) -> Optional[Dict[int, Dict[str, Any]]]:
return await self.client.get_model_versions_bulk(model_ids)
async def get_model_versions_by_hashes(
self, hashes: List[str]
) -> Optional[List[Dict]]:
) -> Optional[List[Dict[str, Any]]]:
return await self.client.get_model_versions_by_hashes(hashes)
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
return await self.client.get_model_version(model_id, version_id)
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
return await self.client.get_model_version_info(version_id)
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict]:
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
return await self.client.get_user_models(username, cursor)
async def get_creator_model_count(self, username: str) -> Optional[int]:
@@ -195,19 +195,19 @@ class CivArchiveModelMetadataProvider(ModelMetadataProvider):
def __init__(self, civarchive_client):
self.client = civarchive_client
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
return await self.client.get_model_by_hash(model_hash)
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
return await self.client.get_model_versions(model_id)
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
return await self.client.get_model_version(model_id, version_id)
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
return await self.client.get_model_version_info(version_id)
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict]:
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
"""Not supported by CivArchive provider"""
return None
@@ -218,7 +218,7 @@ class SQLiteModelMetadataProvider(ModelMetadataProvider):
self.db_path = db_path
self._aiosqlite = _require_aiosqlite()
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Find model by hash value from SQLite database"""
async with self._aiosqlite.connect(self.db_path) as db:
# Look up in model_files table to get model_id and version_id
@@ -243,7 +243,7 @@ class SQLiteModelMetadataProvider(ModelMetadataProvider):
result = await self._get_version_with_model_data(db, model_id, version_id)
return result, None if result else "Error retrieving model data"
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
"""Get all versions of a model from SQLite database"""
async with self._aiosqlite.connect(self.db_path) as db:
db.row_factory = self._aiosqlite.Row
@@ -299,7 +299,7 @@ class SQLiteModelMetadataProvider(ModelMetadataProvider):
'name': model_name
}
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
"""Get specific model version with additional metadata from SQLite database"""
if not model_id and not version_id:
return None
@@ -339,7 +339,7 @@ class SQLiteModelMetadataProvider(ModelMetadataProvider):
# Now we have both model_id and version_id, get the full data
return await self._get_version_with_model_data(db, model_id, version_id)
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Fetch model version metadata from SQLite database"""
async with self._aiosqlite.connect(self.db_path) as db:
db.row_factory = self._aiosqlite.Row
@@ -358,11 +358,11 @@ class SQLiteModelMetadataProvider(ModelMetadataProvider):
version_data = await self._get_version_with_model_data(db, model_id, version_id)
return version_data, None
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict]:
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
"""Listing models by username is not supported for archive database"""
return None
async def _get_version_with_model_data(self, db, model_id, version_id) -> Optional[Dict]:
async def _get_version_with_model_data(self, db, model_id, version_id) -> Optional[Dict[str, Any]]:
"""Helper to build version data with model information"""
# Get version details
version_query = "SELECT name, base_model, data FROM model_versions WHERE id = ? AND model_id = ?"
@@ -485,7 +485,7 @@ class FallbackMetadataProvider(ModelMetadataProvider):
jitter_ratio=self._rate_limit_jitter_ratio,
)
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
for provider, label in self._iter_providers():
try:
result, error = await self._call_with_rate_limit(
@@ -507,7 +507,7 @@ class FallbackMetadataProvider(ModelMetadataProvider):
continue
return None, "Model not found"
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
not_found_confirmed = False
for provider, label in self._iter_providers():
try:
@@ -538,7 +538,7 @@ class FallbackMetadataProvider(ModelMetadataProvider):
continue
return None
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
for provider, label in self._iter_providers():
try:
result = await self._call_with_rate_limit(
@@ -561,7 +561,7 @@ class FallbackMetadataProvider(ModelMetadataProvider):
continue
return None
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
for provider, label in self._iter_providers():
try:
result, error = await self._call_with_rate_limit(
@@ -585,7 +585,7 @@ class FallbackMetadataProvider(ModelMetadataProvider):
async def get_model_versions_by_hashes(
self, hashes: List[str]
) -> Optional[List[Dict]]:
) -> Optional[List[Dict[str, Any]]]:
for provider, label in self._iter_providers():
try:
result = await self._call_with_rate_limit(
@@ -613,7 +613,7 @@ class FallbackMetadataProvider(ModelMetadataProvider):
continue
return None
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict]:
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
for provider, label in self._iter_providers():
try:
result = await self._call_with_rate_limit(
@@ -681,14 +681,14 @@ class RateLimitRetryingProvider(ModelMetadataProvider):
def __getattr__(self, item):
return getattr(self._provider, item)
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_model_by_hash,
model_hash,
)
async def get_model_versions(self, model_id: str) -> Optional[Dict]:
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_model_versions,
@@ -698,7 +698,7 @@ class RateLimitRetryingProvider(ModelMetadataProvider):
async def get_model_versions_bulk(
self,
model_ids: Sequence[int],
) -> Optional[Dict[int, Dict]]:
) -> Optional[Dict[int, Dict[str, Any]]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_model_versions_bulk,
@@ -707,14 +707,14 @@ class RateLimitRetryingProvider(ModelMetadataProvider):
async def get_model_versions_by_hashes(
self, hashes: List[str]
) -> Optional[List[Dict]]:
) -> Optional[List[Dict[str, Any]]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_model_versions_by_hashes,
hashes,
)
async def get_model_version(self, model_id: int = None, version_id: int = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None) -> Optional[Dict[str, Any]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_model_version,
@@ -722,14 +722,14 @@ class RateLimitRetryingProvider(ModelMetadataProvider):
version_id,
)
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_model_version_info,
version_id,
)
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict]:
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
return await self._rate_limit_helper.run(
self._label,
self._provider.get_user_models,
@@ -762,12 +762,12 @@ class ModelMetadataProviderManager:
if is_default or self.default_provider is None:
self.default_provider = name
async def get_model_by_hash(self, model_hash: str, provider_name: str = None) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_by_hash(self, model_hash: str, provider_name: Optional[str] = None) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Find model by hash using specified or default provider"""
provider = self._get_provider(provider_name)
return await provider.get_model_by_hash(model_hash)
async def get_model_versions(self, model_id: str, provider_name: str = None) -> Optional[Dict]:
async def get_model_versions(self, model_id: str, provider_name: Optional[str] = None) -> Optional[Dict[str, Any]]:
"""Get model versions using specified or default provider"""
provider = self._get_provider(provider_name)
return await provider.get_model_versions(model_id)
@@ -775,8 +775,8 @@ class ModelMetadataProviderManager:
async def get_model_versions_bulk(
self,
model_ids: Sequence[int],
provider_name: str = None,
) -> Optional[Dict[int, Dict]]:
provider_name: Optional[str] = None,
) -> Optional[Dict[int, Dict[str, Any]]]:
"""Fetch model versions for multiple model ids when supported by provider."""
provider = self._get_provider(provider_name)
try:
@@ -784,12 +784,12 @@ class ModelMetadataProviderManager:
except NotImplementedError:
return None
async def get_model_version(self, model_id: int = None, version_id: int = None, provider_name: str = None) -> Optional[Dict]:
async def get_model_version(self, model_id: Optional[int] = None, version_id: Optional[int] = None, provider_name: Optional[str] = None) -> Optional[Dict[str, Any]]:
"""Get specific model version using specified or default provider"""
provider = self._get_provider(provider_name)
return await provider.get_model_version(model_id, version_id)
async def get_model_version_info(self, version_id: str, provider_name: str = None) -> Tuple[Optional[Dict], Optional[str]]:
async def get_model_version_info(self, version_id: str, provider_name: Optional[str] = None) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
"""Fetch model version info using specified or default provider"""
provider = self._get_provider(provider_name)
return await provider.get_model_version_info(version_id)
@@ -797,8 +797,8 @@ class ModelMetadataProviderManager:
async def get_model_versions_by_hashes(
self,
hashes: List[str],
provider_name: str = None,
) -> Optional[List[Dict]]:
provider_name: Optional[str] = None,
) -> Optional[List[Dict[str, Any]]]:
provider = self._get_provider(provider_name)
try:
return await provider.get_model_versions_by_hashes(hashes)
@@ -808,19 +808,19 @@ class ModelMetadataProviderManager:
async def get_user_models(
self,
username: str,
provider_name: str = None,
provider_name: Optional[str] = None,
cursor: Optional[str] = None,
) -> Optional[Dict]:
) -> Optional[Dict[str, Any]]:
"""Fetch one page of models owned by the specified user"""
provider = self._get_provider(provider_name)
return await provider.get_user_models(username, cursor)
async def get_creator_model_count(self, username: str, provider_name: str = None) -> Optional[int]:
async def get_creator_model_count(self, username: str, provider_name: Optional[str] = None) -> Optional[int]:
"""Best-effort published model count for the specified user"""
provider = self._get_provider(provider_name)
return await provider.get_creator_model_count(username)
def _get_provider(self, provider_name: str = None) -> ModelMetadataProvider:
def _get_provider(self, provider_name: Optional[str] = None) -> ModelMetadataProvider:
"""Get provider by name or default provider"""
if provider_name:
if provider_name not in self.providers:

3
py/services/model_query.py
View File

@@ -12,6 +12,7 @@ from typing import (
Tuple,
Protocol,
Callable,
cast,
)
from ..utils.constants import NSFW_LEVELS
@@ -309,7 +310,7 @@ class ModelFilterSet:
else:
include_tags.add(normalized)
else:
include_tags = {tag.strip().lower() for tag in tag_filters if tag}
include_tags = {tag.strip().lower() for tag in cast(Iterable[Any], tag_filters) if tag}
if include_tags:
tag_logic = criteria.tag_logic.lower() if criteria.tag_logic else "any"

64
py/services/model_scanner.py
View File

@@ -5,7 +5,7 @@ import asyncio
import time
import shutil
from dataclasses import dataclass
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Set, Type, Union
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Sequence, Set, Type, Union, cast
from ..utils.models import BaseModelMetadata, autov3_from_civitai_files
from ..config import config
@@ -29,7 +29,7 @@ logger = logging.getLogger(__name__)
class CacheBuildResult:
"""Represents the outcome of scanning model files for cache building."""
raw_data: List[Dict]
raw_data: List[Dict[str, Any]]
hash_index: ModelHashIndex
tags_count: Dict[str, int]
excluded_models: List[str]
@@ -59,7 +59,7 @@ class ModelScanner:
lock = cls._get_lock()
async with lock:
if cls not in cls._instances:
cls._instances[cls] = cls()
cls._instances[cls] = cls() # pyright: ignore[reportCallIssue]
return cls._instances[cls]
def __init__(self, model_type: str, model_class: Type[BaseModelMetadata], file_extensions: Set[str], hash_index: Optional[ModelHashIndex] = None):
@@ -78,7 +78,7 @@ class ModelScanner:
self.model_type = model_type
self.model_class = model_class
self.file_extensions = file_extensions
self._cache = None
self._cache: Any = None
self._hash_index = hash_index or ModelHashIndex()
self._tags_count = {} # Dictionary to store tag counts
self._is_initializing = False # Flag to track initialization state
@@ -183,7 +183,7 @@ class ModelScanner:
is_mapping = isinstance(source, Mapping)
def get_value(key: str, default: Any = None) -> Any:
if is_mapping:
if isinstance(source, Mapping):
return source.get(key, default)
sentinel = object()
@@ -772,7 +772,7 @@ class ModelScanner:
else:
await self._reconcile_cache()
return self._cache
return cast(ModelCache, self._cache)
async def _initialize_cache(self) -> None:
"""Initialize or refresh the cache"""
@@ -932,6 +932,8 @@ class ModelScanner:
)
continue
model_data = validation_result.entry
if model_data is None:
continue
self._ensure_license_flags(model_data)
# Add to cache
@@ -992,8 +994,8 @@ class ModelScanner:
self._cache.raw_data = [item for item in self._cache.raw_data if item['file_path'] not in missing_files]
dedup_removed = 0
seen_paths: set = set()
deduped: list = []
seen_paths: set[str] = set()
deduped: list[Dict[str, Any]] = []
for item in reversed(self._cache.raw_data):
path = item.get('file_path', '')
if path not in seen_paths:
@@ -1108,7 +1110,7 @@ class ModelScanner:
*,
hash_index: Optional[ModelHashIndex] = None,
excluded_models: Optional[List[str]] = None
) -> Dict:
) -> Optional[Dict[str, Any]]:
"""Process a single model file and return its metadata"""
hash_index = hash_index or self._hash_index
excluded_models = excluded_models if excluded_models is not None else self._excluded_models
@@ -1132,7 +1134,7 @@ class ModelScanner:
file_name = os.path.splitext(os.path.basename(file_path))[0]
file_info['name'] = file_name
metadata = self.model_class.from_civitai_info(version_info, file_info, file_path)
metadata = cast(Any, self.model_class).from_civitai_info(version_info, file_info, file_path)
metadata.preview_url = find_preview_file(file_name, os.path.dirname(file_path))
await MetadataManager.save_metadata(file_path, metadata)
logger.info(f"Created metadata from .civitai.info for {file_path} (Reason: .civitai.info was found but .metadata.json was missing)")
@@ -1169,6 +1171,8 @@ class ModelScanner:
if metadata is None:
metadata = await self._create_default_metadata(file_path)
assert metadata is not None
# Hook: allow subclasses to adjust metadata
metadata = self.adjust_metadata(metadata, file_path, root_path)
@@ -1296,7 +1300,7 @@ class ModelScanner:
async def _sync_download_history(
self,
raw_data: List[Mapping[str, Any]],
raw_data: Sequence[Mapping[str, Any]],
*,
source: str,
) -> None:
@@ -1345,7 +1349,7 @@ class ModelScanner:
) -> CacheBuildResult:
"""Collect metadata for all model files."""
raw_data: List[Dict] = []
raw_data: List[Dict[str, Any]] = []
hash_index = ModelHashIndex()
tags_count: Dict[str, int] = {}
excluded_models: List[str] = []
@@ -1409,6 +1413,8 @@ class ModelScanner:
)
continue
result = validation_result.entry
if result is None:
continue
self._ensure_license_flags(result)
raw_data.append(result)
@@ -1448,7 +1454,7 @@ class ModelScanner:
excluded_models=excluded_models
)
async def add_model_to_cache(self, metadata_dict: Dict, folder: str = '') -> bool:
async def add_model_to_cache(self, metadata_dict: Dict[str, Any], folder: str = '') -> bool:
"""Add a model to the cache
Args:
@@ -1461,7 +1467,8 @@ class ModelScanner:
try:
if self._cache is None:
await self.get_cached_data()
assert self._cache is not None
# Update folder in metadata
metadata_dict['folder'] = folder
@@ -1496,7 +1503,7 @@ class ModelScanner:
logger.error(f"Error adding model to cache: {e}")
return False
async def move_model(self, source_path: str, target_path: str) -> Optional[str]:
async def move_model(self, source_path: str, target_path: str) -> Optional[Dict[str, Any]]:
"""Move a model and its associated files to a new location
Args:
@@ -1530,7 +1537,7 @@ class ModelScanner:
# Check for filename conflicts and auto-rename if necessary
from ..utils.models import BaseModelMetadata
final_filename = BaseModelMetadata.generate_unique_filename(
target_path, base_name, file_ext, get_source_hash
target_path, base_name, file_ext, lambda: get_source_hash() or ""
)
target_file = os.path.join(target_path, final_filename).replace(os.sep, '/')
@@ -1578,7 +1585,7 @@ class ModelScanner:
logger.error(f"Error moving associated file {source_file}: {e}")
# Handle metadata file specially to update paths
if source_metadata and os.path.exists(source_metadata):
if source_metadata and moved_metadata_path and os.path.exists(source_metadata):
try:
shutil.move(source_metadata, moved_metadata_path)
metadata = await self._update_metadata_paths(moved_metadata_path, target_file)
@@ -1596,7 +1603,7 @@ class ModelScanner:
logger.error(f"Error moving model: {e}", exc_info=True)
return None
async def _update_metadata_paths(self, metadata_path: str, model_path: str) -> Dict:
async def _update_metadata_paths(self, metadata_path: str, model_path: str) -> Optional[Dict[str, Any]]:
"""Update file paths in metadata file"""
try:
with open(metadata_path, 'r', encoding='utf-8') as f:
@@ -1622,7 +1629,7 @@ class ModelScanner:
logger.error(f"Error updating metadata paths: {e}", exc_info=True)
return None
async def update_single_model_cache(self, original_path: str, new_path: str, metadata: Dict, recalculate_type: bool = False) -> Union[bool, Dict]:
async def update_single_model_cache(self, original_path: str, new_path: str, metadata: Optional[Dict[str, Any]], recalculate_type: bool = False) -> Union[bool, Dict[str, Any]]:
"""Update cache after a model has been moved or modified"""
cache = await self.get_cached_data()
@@ -1645,6 +1652,7 @@ class ModelScanner:
]
cache_modified = bool(existing_item) or bool(metadata)
cache_entry: Optional[Dict[str, Any]] = None
if metadata:
normalized_new_path = new_path.replace(os.sep, '/')
@@ -1695,7 +1703,9 @@ class ModelScanner:
if cache_modified:
await self._persist_current_cache()
return cache_entry if metadata else True
if metadata and cache_entry is not None:
return cache_entry
return True
async def sync_cache_from_metadata(
self, file_path: str, metadata_dict: Dict[str, Any]
@@ -1820,8 +1830,8 @@ class ModelScanner:
existing_entry.update(desired_entry)
# ---- Incremental tag count update ----
new_tags: set = set(desired_entry.get("tags") or [])
old_tag_set: set = set(old_tags)
new_tags: set[str] = set(desired_entry.get("tags") or [])
old_tag_set: set[str] = set(old_tags)
for tag in old_tag_set - new_tags:
current = self._tags_count.get(tag, 0)
if current <= 1:
@@ -2020,7 +2030,7 @@ class ModelScanner:
return None
async def get_top_tags(self, limit: int = 20) -> List[Dict[str, any]]:
async def get_top_tags(self, limit: int = 20) -> List[Dict[str, Any]]:
"""Get top tags sorted by count. If limit is 0, return all tags."""
await self.get_cached_data()
@@ -2036,7 +2046,7 @@ class ModelScanner:
async def search_tags(
self, query: str, limit: int = 50
) -> List[Dict[str, any]]:
) -> List[Dict[str, Any]]:
"""Search tags by case-insensitive substring match, sorted by count.
If query is empty, behaves like get_top_tags (returns top ``limit``
@@ -2059,7 +2069,7 @@ class ModelScanner:
return matched
return matched[:limit]
async def get_base_models(self, limit: int = 20) -> List[Dict[str, any]]:
async def get_base_models(self, limit: int = 20) -> List[Dict[str, Any]]:
"""Get base models sorted by count. If limit is 0, return all."""
cache = await self.get_cached_data()
@@ -2140,7 +2150,7 @@ class ModelScanner:
await self._persist_current_cache()
return updated
async def bulk_delete_models(self, file_paths: List[str]) -> Dict:
async def bulk_delete_models(self, file_paths: List[str]) -> Dict[str, Any]:
"""Delete multiple models and update cache in a batch operation
Args:
@@ -2338,7 +2348,7 @@ class ModelScanner:
logger.error(f"Error checking model version existence: {e}")
return False
async def get_model_versions_by_id(self, model_id: int) -> List[Dict]:
async def get_model_versions_by_id(self, model_id: int) -> List[Dict[str, Any]]:
"""Get all versions of a model by its ID
Args:

12
py/services/model_service_factory.py
View File

@@ -6,13 +6,13 @@ logger = logging.getLogger(__name__)
class ModelServiceFactory:
"""Factory for managing model services and routes"""
_services: Dict[str, Type] = {}
_routes: Dict[str, Type] = {}
_services: Dict[str, Type[Any]] = {}
_routes: Dict[str, Type[Any]] = {}
_initialized_services: Dict[str, Any] = {}
_initialized_routes: Dict[str, Any] = {}
@classmethod
def register_model_type(cls, model_type: str, service_class: Type, route_class: Type):
def register_model_type(cls, model_type: str, service_class: Type[Any], route_class: Type[Any]):
"""Register a new model type with its service and route classes
Args:
@@ -24,7 +24,7 @@ class ModelServiceFactory:
cls._routes[model_type] = route_class
@classmethod
def get_service_class(cls, model_type: str) -> Type:
def get_service_class(cls, model_type: str) -> Type[Any]:
"""Get service class for a model type
Args:
@@ -41,7 +41,7 @@ class ModelServiceFactory:
return cls._services[model_type]
@classmethod
def get_route_class(cls, model_type: str) -> Type:
def get_route_class(cls, model_type: str) -> Type[Any]:
"""Get route class for a model type
Args:
@@ -87,7 +87,7 @@ class ModelServiceFactory:
logger.error(f"Failed to setup routes for {model_type}: {e}", exc_info=True)
@classmethod
def get_registered_types(cls) -> list:
def get_registered_types(cls) -> list[str]:
"""Get list of all registered model types
Returns:

53
py/services/model_update_service.py
View File

@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
"""Service for tracking remote model version updates."""
from __future__ import annotations
@@ -336,9 +340,9 @@ class ModelUpdateService:
return
try:
from .persistent_model_cache import get_persistent_cache
from .persistent_model_cache import PersistentModelCache
legacy_path = get_persistent_cache(self._library_name).get_database_path()
legacy_path = PersistentModelCache.get_default(self._library_name).get_database_path()
except Exception:
return
@@ -735,7 +739,7 @@ class ModelUpdateService:
)
results: Dict[int, ModelUpdateRecord] = {}
prefetched: Dict[int, Mapping] = {}
prefetched: Dict[int, Mapping[Any, Any]] = {}
fetch_targets: List[int] = []
if metadata_provider and local_versions:
@@ -834,7 +838,7 @@ class ModelUpdateService:
model_id: int,
version_ids: Sequence[int],
*,
version_info: Optional[Mapping] = None,
version_info: Optional[Mapping[str, Any]] = None,
) -> ModelUpdateRecord:
"""Persist a new set of in-library version identifiers."""
@@ -954,7 +958,11 @@ class ModelUpdateService:
records = self._get_records_bulk(model_type, normalized_ids)
return {
model_id: records.get(model_id).has_update(hide_early_access=hide_early_access) if records.get(model_id) else False
model_id: (
records[model_id].has_update(hide_early_access=hide_early_access)
if model_id in records
else False
)
for model_id in normalized_ids
}
@@ -980,7 +988,7 @@ class ModelUpdateService:
metadata_provider,
*,
force_refresh: bool = False,
prefetched_response: Optional[Mapping] = None,
prefetched_response: Optional[Mapping[str, Any]] = None,
all_local_version_ids: Optional[Sequence[int]] = None,
) -> Optional[ModelUpdateRecord]:
normalized_local = self._normalize_sequence(local_versions)
@@ -1010,7 +1018,7 @@ class ModelUpdateService:
fallback_attempted = False
fallback_error_message: Optional[str] = None
mark_model_as_ignored = False
response: Optional[Mapping] = None
response: Optional[Mapping[str, Any]] = None
if metadata_provider and should_fetch:
response = prefetched_response
if response is None:
@@ -1122,7 +1130,7 @@ class ModelUpdateService:
async def _enrich_version_entries(
self,
metadata_provider,
responses_by_model_id: Dict[int, Mapping],
responses_by_model_id: Dict[int, Mapping[Any, Any]],
) -> None:
"""Enrich version entries with ``usageControl`` via batch hash endpoint.
@@ -1151,7 +1159,7 @@ class ModelUpdateService:
all_hashes = list(version_ids_by_hash.keys())
BATCH_SIZE = 100
enrichment: Dict[int, Dict] = {}
enrichment: Dict[int, Dict[str, Any]] = {}
try:
for start in range(0, len(all_hashes), BATCH_SIZE):
batch = all_hashes[start : start + BATCH_SIZE]
@@ -1208,7 +1216,7 @@ class ModelUpdateService:
version["earlyAccessEndsAt"] = extra["earlyAccessEndsAt"]
@staticmethod
def _collect_hashes_from_response(response: Mapping) -> Dict[int, str]:
def _collect_hashes_from_response(response: Mapping[str, Any]) -> Dict[int, str]:
"""Extract ``{version_id: sha256}`` from a model-level API response.
Returns an empty dict if the response structure is unexpected.
@@ -1229,7 +1237,7 @@ class ModelUpdateService:
return result
@staticmethod
def _extract_sha256_from_version_entry(entry: Mapping) -> Optional[str]:
def _extract_sha256_from_version_entry(entry: Mapping[str, Any]) -> Optional[str]:
"""Return the SHA256 hash from the primary model file of a version entry."""
files = entry.get("files")
if not isinstance(files, list):
@@ -1253,22 +1261,19 @@ class ModelUpdateService:
self,
metadata_provider,
model_ids: Sequence[int],
) -> Dict[int, Mapping]:
) -> Dict[int, Mapping[Any, Any]]:
"""Fetch model metadata in batches of up to 100 ids."""
BATCH_SIZE = 100
normalized = self._normalize_sequence(model_ids)
if not normalized:
provider = metadata_provider
if not normalized or provider is None:
return {}
aggregated: Dict[int, Mapping] = {}
aggregated: Dict[int, Mapping[Any, Any]] = {}
total_ids = len(normalized)
total_batches = (total_ids + BATCH_SIZE - 1) // BATCH_SIZE
provider_name = (
metadata_provider.__class__.__name__
if metadata_provider is not None
else "unknown"
)
provider_name = provider.__class__.__name__
for batch_index, start in enumerate(range(0, total_ids, BATCH_SIZE), start=1):
chunk = normalized[start : start + BATCH_SIZE]
logger.info(
@@ -1279,7 +1284,7 @@ class ModelUpdateService:
provider_name,
)
try:
response = await metadata_provider.get_model_versions_bulk(chunk)
response = await provider.get_model_versions_bulk(chunk)
except RateLimitError:
raise
if response is None:
@@ -1356,7 +1361,7 @@ class ModelUpdateService:
model_type: Optional[str] = None,
model_id: Optional[int] = None,
last_checked_at: Optional[float] = None,
version_info: Optional[Mapping] = None,
version_info: Optional[Mapping[str, Any]] = None,
) -> ModelUpdateRecord:
local_set = set(normalized_local)
# When folder-filtering, also consider versions in other folders
@@ -1578,7 +1583,7 @@ class ModelUpdateService:
if not isinstance(files, Iterable):
return None
def parse_size(entry: Mapping) -> Optional[int]:
def parse_size(entry: Mapping[str, Any]) -> Optional[int]:
size_kb = entry.get("sizeKB")
if size_kb is None:
return None
@@ -1664,8 +1669,8 @@ class ModelUpdateService:
return {}
ids = list(model_ids)
status_rows: list = []
version_rows: list = []
status_rows: list[sqlite3.Row] = []
version_rows: list[sqlite3.Row] = []
with self._connect() as conn:
for start in range(0, len(ids), self._SQLITE_MAX_VARIABLES):

40
py/services/persistent_model_cache.py
View File

@@ -4,7 +4,7 @@ import os
import sqlite3
import threading
from dataclasses import dataclass, field
from typing import Dict, List, Mapping, Optional, Sequence, Tuple
from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration
@@ -15,7 +15,7 @@ logger = logging.getLogger(__name__)
class PersistedCacheData:
"""Lightweight structure returned by the persistent cache."""
raw_data: List[Dict]
raw_data: List[Dict[str, Any]]
hash_rows: List[Tuple[str, str]]
excluded_models: List[str]
autov3_hash_rows: List[Tuple[str, str]] = field(default_factory=list)
@@ -70,8 +70,8 @@ class PersistentModelCache:
self._db_path = db_path or self._resolve_default_path(self._library_name)
self._db_lock = threading.Lock()
self._schema_initialized = False
directory = os.path.dirname(self._db_path)
try:
directory = os.path.dirname(self._db_path)
if directory:
os.makedirs(directory, exist_ok=True)
except Exception as exc: # pragma: no cover - defensive guard
@@ -134,7 +134,7 @@ class PersistentModelCache:
logger.warning("Failed to load persisted cache for %s: %s", model_type, exc)
return None
raw_data: List[Dict] = []
raw_data: List[Dict[str, Any]] = []
for row in rows:
file_path: str = row["file_path"]
trained_words = []
@@ -145,7 +145,7 @@ class PersistentModelCache:
trained_words = []
creator_username = row["civitai_creator_username"]
civitai: Optional[Dict] = None
civitai: Optional[Dict[str, Any]] = None
civitai_has_data = any(
row[col] is not None
for col in ("civitai_id", "civitai_model_id", "civitai_model_type", "civitai_name")
@@ -223,7 +223,7 @@ class PersistentModelCache:
autov3_hash_rows=autov3_pairs,
)
def save_cache(self, model_type: str, raw_data: Sequence[Dict], hash_index: Dict[str, List[str]], excluded_models: Sequence[str], autov3_hash_index: Optional[Dict[str, List[str]]] = None) -> None:
def save_cache(self, model_type: str, raw_data: Sequence[Dict[str, Any]], hash_index: Dict[str, List[str]], excluded_models: Sequence[str], autov3_hash_index: Optional[Dict[str, List[str]]] = None) -> None:
if not self.is_enabled():
return
if not self._schema_initialized:
@@ -238,7 +238,7 @@ class PersistentModelCache:
conn.execute("BEGIN")
model_rows = [self._prepare_model_row(model_type, item) for item in raw_data]
model_map: Dict[str, Tuple] = {
model_map: Dict[str, Tuple[Any, ...]] = {
row[1]: row for row in model_rows if row[1] # row[1] is file_path
}
@@ -279,8 +279,8 @@ class PersistentModelCache:
to_remove_models,
)
insert_rows: List[Tuple] = []
update_rows: List[Tuple] = []
insert_rows: List[Tuple[Any, ...]] = []
update_rows: List[Tuple[Any, ...]] = []
for file_path, row in model_map.items():
existing = existing_model_map.get(file_path)
@@ -312,11 +312,11 @@ class PersistentModelCache:
"SELECT file_path, tag FROM model_tags WHERE model_type = ?",
(model_type,),
).fetchall()
existing_tags: Dict[str, set] = {}
existing_tags: Dict[str, set[str]] = {}
for row in existing_tags_rows:
existing_tags.setdefault(row["file_path"], set()).add(row["tag"])
new_tags: Dict[str, set] = {}
new_tags: Dict[str, set[str]] = {}
for item in raw_data:
file_path = item.get("file_path")
if not file_path:
@@ -355,14 +355,14 @@ class PersistentModelCache:
"SELECT sha256, file_path FROM hash_index WHERE model_type = ?",
(model_type,),
).fetchall()
existing_hash_map: Dict[str, set] = {}
existing_hash_map: Dict[str, set[str]] = {}
for row in existing_hash_rows:
sha_value = (row["sha256"] or "").lower()
if not sha_value:
continue
existing_hash_map.setdefault(sha_value, set()).add(row["file_path"])
new_hash_map: Dict[str, set] = {}
new_hash_map: Dict[str, set[str]] = {}
for sha_value, paths in hash_index.items():
normalized_sha = (sha_value or "").lower()
if not normalized_sha:
@@ -401,14 +401,14 @@ class PersistentModelCache:
"SELECT autov3, file_path FROM autov3_index WHERE model_type = ?",
(model_type,),
).fetchall()
existing_autov3_map: Dict[str, set] = {}
existing_autov3_map: Dict[str, set[str]] = {}
for row in existing_autov3_rows:
autov3_value = (row["autov3"] or "").lower()
if not autov3_value:
continue
existing_autov3_map.setdefault(autov3_value, set()).add(row["file_path"])
new_autov3_map: Dict[str, set] = {}
new_autov3_map: Dict[str, set[str]] = {}
for autov3_value, paths in autov3_hash_index.items():
normalized_autov3 = (autov3_value or "").lower()
if not normalized_autov3:
@@ -600,7 +600,7 @@ class PersistentModelCache:
conn.row_factory = sqlite3.Row
return conn
def _prepare_model_row(self, model_type: str, item: Dict) -> Tuple:
def _prepare_model_row(self, model_type: str, item: Dict[str, Any]) -> Tuple[Any, ...]:
civitai = item.get("civitai") or {}
trained_words = civitai.get("trainedWords")
if isinstance(trained_words, str):
@@ -675,8 +675,8 @@ class PersistentModelCache:
def update_single_model(
self,
model_type: str,
new_item: Dict,
old_item: Optional[Dict] = None,
new_item: Dict[str, Any],
old_item: Optional[Dict[str, Any]] = None,
) -> None:
"""Update a single model row in the persistent cache.
@@ -715,8 +715,8 @@ class PersistentModelCache:
conn.execute(self._insert_model_sql(), row)
# --- tags ---
new_tags: set = set(new_item.get("tags") or [])
old_tags: set = set(old_item.get("tags") or []) if old_item else set()
new_tags: set[str] = set(new_item.get("tags") or [])
old_tags: set[str] = set(old_item.get("tags") or []) if old_item else set()
tags_to_delete = old_tags - new_tags
tags_to_insert = new_tags - old_tags

20
py/services/persistent_recipe_cache.py
View File

@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
"""SQLite-based persistent cache for recipe metadata.
This module provides fast recipe cache persistence using SQLite, enabling
@@ -13,7 +17,7 @@ import os
import sqlite3
import threading
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set, Tuple
from typing import Any, Dict, List, Optional, Set, Tuple
from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration
@@ -24,7 +28,7 @@ logger = logging.getLogger(__name__)
class PersistedRecipeData:
"""Lightweight structure returned by the persistent recipe cache."""
raw_data: List[Dict]
raw_data: List[Dict[str, Any]]
file_stats: Dict[str, Tuple[float, int]] # json_path -> (mtime, size)
image_id_map: Dict[str, str] = field(default_factory=dict)
"""Precomputed mapping of civitai image_id → recipe_id."""
@@ -63,8 +67,8 @@ class PersistentRecipeCache:
self._db_path = db_path or self._resolve_default_path(self._library_name)
self._db_lock = threading.Lock()
self._schema_initialized = False
directory = os.path.dirname(self._db_path)
try:
directory = os.path.dirname(self._db_path)
if directory:
os.makedirs(directory, exist_ok=True)
except Exception as exc:
@@ -140,7 +144,7 @@ class PersistentRecipeCache:
logger.warning("Failed to load persisted recipe cache: %s", exc)
return None
raw_data: List[Dict] = []
raw_data: List[Dict[str, Any]] = []
file_stats: Dict[str, Tuple[float, int]] = {}
for row in rows:
@@ -162,7 +166,7 @@ class PersistentRecipeCache:
def save_cache(
self,
recipes: List[Dict],
recipes: List[Dict[str, Any]],
json_paths: Optional[Dict[str, str]] = None,
image_id_map: Optional[Dict[str, str]] = None,
) -> None:
@@ -251,7 +255,7 @@ class PersistentRecipeCache:
except Exception:
return {}
def update_recipe(self, recipe: Dict, json_path: Optional[str] = None) -> None:
def update_recipe(self, recipe: Dict[str, Any], json_path: Optional[str] = None) -> None:
"""Update or insert a single recipe in the cache.
Args:
@@ -439,7 +443,7 @@ class PersistentRecipeCache:
conn.row_factory = sqlite3.Row
return conn
def _prepare_recipe_row(self, recipe: Dict, json_path: str) -> Tuple:
def _prepare_recipe_row(self, recipe: Dict[str, Any], json_path: str) -> Tuple[Any, ...]:
"""Convert a recipe dict to a row tuple for SQLite insertion."""
loras = recipe.get("loras")
loras_json = json.dumps(loras) if loras else None
@@ -486,7 +490,7 @@ class PersistentRecipeCache:
tags_json,
)
def _row_to_recipe(self, row: sqlite3.Row) -> Dict:
def _row_to_recipe(self, row: sqlite3.Row) -> Dict[str, Any]:
"""Convert a SQLite row to a recipe dictionary."""
loras = []
if row["loras_json"]:

4
py/services/preview_asset_service.py
View File

@@ -22,7 +22,7 @@ class PreviewAssetService:
self,
*,
metadata_manager,
downloader_factory: Callable[[], Awaitable],
downloader_factory: Callable[[], Awaitable[Any]],
exif_utils,
) -> None:
self._metadata_manager = metadata_manager
@@ -69,6 +69,8 @@ class PreviewAssetService:
if not preview_url:
return
preview_url = str(preview_url)
def extension_from_url(url: str, fallback: str) -> str:
try:
parsed = urlparse(url)

25
py/services/recipe_cache.py
View File

@@ -1,5 +1,5 @@
import asyncio
from typing import Iterable, List, Dict, Optional
from typing import Any, Iterable, List, Dict, Optional
from dataclasses import dataclass, field
from natsort import natsorted
@@ -8,12 +8,13 @@ from natsort import natsorted
class RecipeCache:
"""Cache structure for Recipe data"""
raw_data: List[Dict]
sorted_by_name: List[Dict]
sorted_by_date: List[Dict]
raw_data: List[Dict[str, Any]]
sorted_by_name: List[Dict[str, Any]]
sorted_by_date: List[Dict[str, Any]]
folders: List[str] | None = None
folder_tree: Dict | None = None
folder_tree: Dict[str, Any] | None = None
image_id_map: Dict[str, str] = field(default_factory=dict)
_lock: Any = field(init=False, repr=False, default=None)
"""Mapping of civitai image_id → recipe_id, precomputed at cache build time.
Built once during cache initialization (O(n)) so that
@@ -40,7 +41,7 @@ class RecipeCache:
)
async def update_recipe_metadata(
self, recipe_id: str, metadata: Dict, *, resort: bool = True
self, recipe_id: str, metadata: Dict[str, Any], *, resort: bool = True
) -> bool:
"""Update metadata for a specific recipe in all cached data
@@ -60,7 +61,7 @@ class RecipeCache:
return True
return False # Recipe not found
async def add_recipe(self, recipe_data: Dict, *, resort: bool = False) -> None:
async def add_recipe(self, recipe_data: Dict[str, Any], *, resort: bool = False) -> None:
"""Add a new recipe to the cache."""
async with self._lock:
@@ -70,7 +71,7 @@ class RecipeCache:
async def remove_recipe(
self, recipe_id: str, *, resort: bool = False
) -> Optional[Dict]:
) -> Optional[Dict[str, Any]]:
"""Remove a recipe from the cache by ID.
Args:
@@ -91,7 +92,7 @@ class RecipeCache:
async def bulk_remove(
self, recipe_ids: Iterable[str], *, resort: bool = False
) -> List[Dict]:
) -> List[Dict[str, Any]]:
"""Remove multiple recipes from the cache."""
id_set = {str(recipe_id) for recipe_id in recipe_ids}
@@ -111,7 +112,7 @@ class RecipeCache:
return removed
async def replace_recipe(
self, recipe_id: str, new_data: Dict, *, resort: bool = False
self, recipe_id: str, new_data: Dict[str, Any], *, resort: bool = False
) -> bool:
"""Replace cached data for a recipe."""
@@ -124,7 +125,7 @@ class RecipeCache:
return True
return False
async def get_recipe(self, recipe_id: str) -> Optional[Dict]:
async def get_recipe(self, recipe_id: str) -> Optional[Dict[str, Any]]:
"""Return a shallow copy of a cached recipe."""
async with self._lock:
@@ -133,7 +134,7 @@ class RecipeCache:
return dict(recipe)
return None
async def snapshot(self) -> List[Dict]:
async def snapshot(self) -> List[Dict[str, Any]]:
"""Return a copy of all cached recipes."""
async with self._lock:

4
py/services/recipe_fts_index.py
View File

@@ -58,8 +58,8 @@ class RecipeFTSIndex:
self._warned_not_ready = False
# Ensure directory exists
directory = os.path.dirname(self._db_path)
try:
directory = os.path.dirname(self._db_path)
if directory:
os.makedirs(directory, exist_ok=True)
except Exception as exc:
@@ -509,7 +509,7 @@ class RecipeFTSIndex:
(recipe_id,)
)
def _prepare_fts_row(self, recipe: Dict[str, Any]) -> tuple:
def _prepare_fts_row(self, recipe: Dict[str, Any]) -> tuple[str, str, str, str, str, str, str]:
"""Prepare a row tuple for FTS insertion."""
recipe_id = str(recipe.get('id', ''))
title = str(recipe.get('title', ''))

156
py/services/recipe_scanner.py
View File

@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
from __future__ import annotations
import asyncio
@@ -5,28 +9,20 @@ import json
import logging
import os
import time
from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple
from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Union, cast
from ..config import config
from .recipe_cache import RecipeCache
from .recipe_fts_index import RecipeFTSIndex
from .persistent_recipe_cache import (
PersistentRecipeCache,
get_persistent_recipe_cache,
PersistedRecipeData,
)
from .service_registry import ServiceRegistry
from .lora_scanner import LoraScanner
from .metadata_service import get_default_metadata_provider
from .checkpoint_scanner import CheckpointScanner
from .settings_manager import get_settings_manager
from .recipes.errors import RecipeNotFoundError
from ..utils.civitai_utils import extract_civitai_image_id
from ..utils.utils import calculate_recipe_fingerprint
from natsort import natsorted
import sys
import re
from ..recipes.merger import GenParamsMerger
from ..recipes.enrichment import RecipeEnricher
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from .lora_scanner import LoraScanner
from .checkpoint_scanner import CheckpointScanner
from .recipe_fts_index import RecipeFTSIndex
from .persistent_recipe_cache import PersistentRecipeCache, PersistedRecipeData
logger = logging.getLogger(__name__)
@@ -48,8 +44,12 @@ class RecipeScanner:
if cls._instance is None:
if not lora_scanner:
# Get lora scanner from service registry if not provided
from .service_registry import ServiceRegistry
lora_scanner = await ServiceRegistry.get_lora_scanner()
if not checkpoint_scanner:
from .service_registry import ServiceRegistry
checkpoint_scanner = await ServiceRegistry.get_checkpoint_scanner()
cls._instance = cls(lora_scanner, checkpoint_scanner)
return cls._instance
@@ -77,17 +77,18 @@ class RecipeScanner:
if not hasattr(self, "_initialized"):
self._cache: Optional[RecipeCache] = None
self._initialization_lock = asyncio.Lock()
self._initialization_task: Optional[asyncio.Task] = None
self._initialization_task: Optional[asyncio.Task[Any]] = None
self._is_initializing = False
self._mutation_lock = asyncio.Lock()
self._post_scan_task: Optional[asyncio.Task] = None
self._resort_tasks: Set[asyncio.Task] = set()
self._post_scan_task: Optional[asyncio.Task[Any]] = None
self._resort_tasks: Set[asyncio.Task[Any]] = set()
self._cancel_requested = False
# FTS index for fast search
self._fts_index: Optional[RecipeFTSIndex] = None
self._fts_index_task: Optional[asyncio.Task] = None
self._fts_index_task: Optional[asyncio.Task[Any]] = None
# Persistent cache for fast startup
self._persistent_cache: Optional[PersistentRecipeCache] = None
self._civitai_client: Any = None # Lazily initialized from registry
self._json_path_map: Dict[str, str] = {} # recipe_id -> json_path
if lora_scanner:
self._lora_scanner = lora_scanner
@@ -123,6 +124,8 @@ class RecipeScanner:
# Reset persistent cache instance for new library
self._persistent_cache = None
self._json_path_map = {}
from .persistent_recipe_cache import PersistentRecipeCache
PersistentRecipeCache.clear_instances()
self._cache = None
@@ -140,6 +143,8 @@ class RecipeScanner:
async def _get_civitai_client(self):
"""Lazily initialize CivitaiClient from registry"""
if self._civitai_client is None:
from .service_registry import ServiceRegistry
self._civitai_client = await ServiceRegistry.get_civitai_client()
return self._civitai_client
@@ -157,7 +162,7 @@ class RecipeScanner:
return self._cancel_requested
async def repair_all_recipes(
self, progress_callback: Optional[Callable[[Dict], Any]] = None
self, progress_callback: Optional[Callable[[Dict[str, Any]], Any]] = None
) -> Dict[str, Any]:
"""Repair all recipes by enrichment with Civitai and embedded metadata.
@@ -339,6 +344,8 @@ class RecipeScanner:
# 3. Use Enricher to repair/enrich
try:
from ..recipes.enrichment import RecipeEnricher
updated = await RecipeEnricher.enrich_recipe(recipe, civitai_client)
except Exception as e:
logger.error(f"Error enriching recipe {recipe.get('id')}: {e}")
@@ -490,6 +497,7 @@ class RecipeScanner:
3. Fall back to full directory scan if cache miss or reconciliation fails
4. Persist results for next startup
"""
loop = None
try:
# Ensure cache exists to avoid None reference errors
if self._cache is None:
@@ -507,6 +515,8 @@ class RecipeScanner:
# Initialize persistent cache
if self._persistent_cache is None:
from .persistent_recipe_cache import get_persistent_recipe_cache
self._persistent_cache = get_persistent_recipe_cache()
recipes_dir = self.recipes_dir
@@ -592,13 +602,14 @@ class RecipeScanner:
return self._cache if hasattr(self, "_cache") else None
finally:
# Clean up the event loop
loop.close()
if loop is not None:
loop.close()
def _reconcile_recipe_cache(
self,
persisted: PersistedRecipeData,
recipes_dir: str,
) -> Tuple[List[Dict], bool, Dict[str, str]]:
) -> Tuple[List[Dict[str, Any]], bool, Dict[str, str]]:
"""Reconcile persisted cache with current filesystem state.
Args:
@@ -608,7 +619,7 @@ class RecipeScanner:
Returns:
Tuple of (recipes list, changed flag, json_paths dict).
"""
recipes: List[Dict] = []
recipes: List[Dict[str, Any]] = []
json_paths: Dict[str, str] = {}
changed = False
@@ -625,12 +636,12 @@ class RecipeScanner:
continue
# Build recipe_id -> recipe lookup (O(n) instead of O(n²))
recipe_by_id: Dict[str, Dict] = {
recipe_by_id: Dict[str, Dict[str, Any]] = {
str(r.get("id", "")): r for r in persisted.raw_data if r.get("id")
}
# Build json_path -> recipe lookup from file_stats (O(m))
persisted_by_path: Dict[str, Dict] = {}
persisted_by_path: Dict[str, Dict[str, Any]] = {}
for json_path in persisted.file_stats.keys():
basename = os.path.basename(json_path)
if basename.lower().endswith(".recipe.json"):
@@ -696,7 +707,7 @@ class RecipeScanner:
def _backfill_source_path_if_needed(
self,
recipes: List[Dict],
recipes: List[Dict[str, Any]],
json_paths: Dict[str, str],
) -> bool:
"""Backfill source_path from recipe JSON files if missing from cache.
@@ -724,7 +735,7 @@ class RecipeScanner:
def _full_directory_scan_sync(
self, recipes_dir: str
) -> Tuple[List[Dict], Dict[str, str]]:
) -> Tuple[List[Dict[str, Any]], Dict[str, str]]:
"""Perform a full synchronous directory scan for recipes.
Args:
@@ -733,7 +744,7 @@ class RecipeScanner:
Returns:
Tuple of (recipes list, json_paths dict).
"""
recipes: List[Dict] = []
recipes: List[Dict[str, Any]] = []
json_paths: Dict[str, str] = {}
# Get all recipe JSON files
@@ -756,7 +767,7 @@ class RecipeScanner:
return recipes, json_paths
def _load_recipe_file_sync(self, recipe_path: str) -> Optional[Dict]:
def _load_recipe_file_sync(self, recipe_path: str) -> Optional[Dict[str, Any]]:
"""Load a single recipe file synchronously.
Args:
@@ -835,6 +846,8 @@ class RecipeScanner:
def _sort_cache_sync(self) -> None:
"""Sort cache data synchronously."""
if self._cache is None:
return
try:
# Sort by name
self._cache.sorted_by_name = natsorted(
@@ -868,6 +881,8 @@ class RecipeScanner:
source = recipe.get("source_path")
if not source:
continue
from ..utils.civitai_utils import extract_civitai_image_id
image_id = extract_civitai_image_id(source)
if image_id and image_id not in mapping:
recipe_id = recipe.get("id")
@@ -950,6 +965,8 @@ class RecipeScanner:
return
try:
from .recipe_fts_index import RecipeFTSIndex
self._fts_index = RecipeFTSIndex()
# Check if existing index is valid
@@ -987,7 +1004,7 @@ class RecipeScanner:
_build_fts(), name="recipe_fts_index_build"
)
def _search_with_fts(self, search: str, search_options: Dict) -> Optional[Set[str]]:
def _search_with_fts(self, search: str, search_options: Dict[str, Any]) -> Optional[Set[str]]:
"""Search recipes using FTS index if available.
Args:
@@ -1002,7 +1019,7 @@ class RecipeScanner:
return None
# Build the set of fields to search based on search_options
fields: Set[str] = set()
fields: Optional[Set[str]] = set()
if search_options.get("title", True):
fields.add("title")
if search_options.get("tags", True):
@@ -1033,12 +1050,12 @@ class RecipeScanner:
return None
def _update_fts_index_for_recipe(
self, recipe: Dict[str, Any], operation: str = "add"
self, recipe: Union[Dict[str, Any], str], operation: str = "add"
) -> None:
"""Update FTS index for a single recipe (add, update, or remove).
Args:
recipe: The recipe dictionary.
recipe: The recipe dictionary, or a recipe ID string for removal.
operation: One of 'add', 'update', or 'remove'.
"""
if not self._fts_index or not self._fts_index.is_ready():
@@ -1053,7 +1070,7 @@ class RecipeScanner:
)
self._fts_index.remove_recipe(recipe_id)
elif operation in ("add", "update"):
self._fts_index.update_recipe(recipe)
self._fts_index.update_recipe(cast(Dict[str, Any], recipe))
except Exception as exc:
logger.debug("Failed to update FTS index for recipe: %s", exc)
@@ -1071,6 +1088,8 @@ class RecipeScanner:
if value in (None, ""):
continue
from ..recipes.merger import GenParamsMerger
normalized_key = GenParamsMerger.NORMALIZATION_MAPPING.get(key, key)
if normalized_key not in GenParamsMerger.ALLOWED_KEYS:
continue
@@ -1130,7 +1149,8 @@ class RecipeScanner:
def _schedule_resort(self, *, name_only: bool = False) -> None:
"""Schedule a background resort of the recipe cache."""
if not self._cache:
cache = self._cache
if not cache:
return
# Keep folder metadata up to date alongside sort order
@@ -1138,7 +1158,7 @@ class RecipeScanner:
async def _resort_wrapper() -> None:
try:
await self._cache.resort(name_only=name_only)
await cache.resort(name_only=name_only)
except Exception as exc: # pragma: no cover - defensive logging
logger.error(
"Recipe Scanner: error resorting cache: %s", exc, exc_info=True
@@ -1164,10 +1184,10 @@ class RecipeScanner:
except Exception:
return ""
def _build_folder_tree(self, folders: list[str]) -> dict:
def _build_folder_tree(self, folders: list[str]) -> Dict[str, Any]:
"""Build a nested folder tree structure from relative folder paths."""
tree: dict[str, dict] = {}
tree: dict[str, Dict[str, Any]] = {}
for folder in folders:
if not folder:
continue
@@ -1208,18 +1228,20 @@ class RecipeScanner:
cache = await self.get_cached_data()
self._update_folder_metadata(cache)
return cache.folders
return cache.folders or []
async def get_folder_tree(self) -> dict:
async def get_folder_tree(self) -> Dict[str, Any]:
"""Return a hierarchical tree of recipe folders for sidebar navigation."""
cache = await self.get_cached_data()
self._update_folder_metadata(cache)
return cache.folder_tree
return cache.folder_tree or {}
@property
def recipes_dir(self) -> str:
"""Get path to recipes directory"""
from .settings_manager import get_settings_manager
custom_recipes_dir = get_settings_manager().get("recipes_path", "")
if isinstance(custom_recipes_dir, str) and custom_recipes_dir.strip():
recipes_dir = os.path.abspath(
@@ -1242,7 +1264,7 @@ class RecipeScanner:
# If cache is already initialized and no refresh is needed, return it immediately
if self._cache is not None and not force_refresh:
self._update_folder_metadata()
return self._cache
return cast(RecipeCache, self._cache)
# If another initialization is already in progress, wait for it to complete
if self._is_initializing and not force_refresh:
@@ -1293,7 +1315,7 @@ class RecipeScanner:
self._schedule_post_scan_enrichment()
self._schedule_fts_index_build()
return self._cache
return cast(RecipeCache, self._cache)
except Exception as e:
logger.error(
@@ -1344,6 +1366,8 @@ class RecipeScanner:
source = recipe_data.get("source_path")
if source:
from ..utils.civitai_utils import extract_civitai_image_id
image_id = extract_civitai_image_id(source)
if image_id:
recipe_id_value = recipe_data.get("id")
@@ -1410,7 +1434,7 @@ class RecipeScanner:
self._persistent_cache.save_image_id_map(cache.image_id_map)
return len(removed)
async def scan_all_recipes(self) -> List[Dict]:
async def scan_all_recipes(self) -> List[Dict[str, Any]]:
"""Scan all recipe JSON files and return metadata"""
recipes = []
recipes_dir = self.recipes_dir
@@ -1436,7 +1460,7 @@ class RecipeScanner:
return recipes
async def _load_recipe_file(self, recipe_path: str) -> Optional[Dict]:
async def _load_recipe_file(self, recipe_path: str) -> Optional[Dict[str, Any]]:
"""Load recipe data from a JSON file"""
try:
with open(recipe_path, "r", encoding="utf-8") as f:
@@ -1517,6 +1541,8 @@ class RecipeScanner:
# Calculate and update fingerprint if missing
if "loras" in recipe_data and "fingerprint" not in recipe_data:
from ..utils.utils import calculate_recipe_fingerprint
fingerprint = calculate_recipe_fingerprint(recipe_data["loras"])
recipe_data["fingerprint"] = fingerprint
@@ -1548,7 +1574,7 @@ class RecipeScanner:
with open(recipe_path, "w", encoding="utf-8") as file_obj:
json.dump(recipe_data, file_obj, indent=4, ensure_ascii=False)
async def _update_lora_information(self, recipe_data: Dict) -> bool:
async def _update_lora_information(self, recipe_data: Dict[str, Any]) -> bool:
"""Update LoRA information with hash and file_name
Returns:
@@ -1575,14 +1601,14 @@ class RecipeScanner:
if isinstance(model_version_id, int) and model_version_id > 0:
# Try to find in lora cache first
hash_from_cache = await self._find_hash_in_lora_cache(
model_version_id
str(model_version_id)
)
if hash_from_cache:
lora["hash"] = hash_from_cache
metadata_updated = True
else:
# If not in cache, fetch from Civitai
result = await self._get_hash_from_civitai(model_version_id)
result = await self._get_hash_from_civitai(str(model_version_id))
if isinstance(result, tuple):
hash_from_civitai, is_deleted = result
if hash_from_civitai:
@@ -1645,14 +1671,16 @@ class RecipeScanner:
logger.error(f"Error finding hash in lora cache: {e}")
return None
async def _get_hash_from_civitai(self, model_version_id: str) -> Optional[str]:
async def _get_hash_from_civitai(self, model_version_id: str) -> Tuple[Optional[str], bool]:
"""Get hash from Civitai API"""
try:
# Get metadata provider instead of civitai client directly
from .metadata_service import get_default_metadata_provider
metadata_provider = await get_default_metadata_provider()
if not metadata_provider:
logger.error("Failed to get metadata provider")
return None
return None, False
version_info, error_msg = await metadata_provider.get_model_version_info(
model_version_id
@@ -1733,7 +1761,7 @@ class RecipeScanner:
return version_index.get(normalized_id)
async def _determine_base_model(self, loras: List[Dict]) -> Optional[str]:
async def _determine_base_model(self, loras: List[Dict[str, Any]]) -> Optional[str]:
"""Determine the most common base model among LoRAs"""
base_models = {}
@@ -1956,11 +1984,11 @@ class RecipeScanner:
page: int,
page_size: int,
sort_by: str = "date",
search: str = None,
filters: dict = None,
search_options: dict = None,
lora_hash: str = None,
checkpoint_hash: str = None,
search: Optional[str] = None,
filters: Optional[Dict[str, Any]] = None,
search_options: Optional[Dict[str, Any]] = None,
lora_hash: Optional[str] = None,
checkpoint_hash: Optional[str] = None,
bypass_filters: bool = True,
folder: str | None = None,
recursive: bool = True,
@@ -2220,7 +2248,7 @@ class RecipeScanner:
return result
async def get_recipe_by_id(self, recipe_id: str) -> dict:
async def get_recipe_by_id(self, recipe_id: str) -> Optional[Dict[str, Any]]:
"""Get a single recipe by ID with all metadata and formatted URLs
Args:
@@ -2312,7 +2340,7 @@ class RecipeScanner:
return self._normalize_recipe_gen_params(recipe_data)
def _format_file_url(self, file_path: str) -> str:
def _format_file_url(self, file_path: Optional[str]) -> str:
"""Format file path as URL for serving in web UI"""
if not file_path:
return "/loras_static/images/no-preview.png"
@@ -2360,7 +2388,7 @@ class RecipeScanner:
return None
async def update_recipe_metadata(self, recipe_id: str, metadata: dict) -> bool:
async def update_recipe_metadata(self, recipe_id: str, metadata: Dict[str, Any]) -> bool:
"""Update recipe metadata (like title and tags) in both file system and cache
Args:
@@ -2465,6 +2493,8 @@ class RecipeScanner:
lora_entry["modelVersionName"] = civitai_info.get("name", "")
lora_entry["modelVersionId"] = civitai_info.get("id")
from ..utils.utils import calculate_recipe_fingerprint
recipe_data["fingerprint"] = calculate_recipe_fingerprint(
recipe_data.get("loras", [])
)
@@ -2696,7 +2726,7 @@ class RecipeScanner:
return file_updated_count, cache_updated_count
async def find_recipes_by_fingerprint(self, fingerprint: str) -> list:
async def find_recipes_by_fingerprint(self, fingerprint: str) -> List[Dict[str, Any]]:
"""Find recipes with a matching fingerprint
Args:
@@ -2727,7 +2757,7 @@ class RecipeScanner:
return matching_recipes
async def find_all_duplicate_recipes(self) -> dict:
async def find_all_duplicate_recipes(self) -> Dict[str, List[Any]]:
"""Find all recipe duplicates based on fingerprints
Returns:
@@ -2753,7 +2783,7 @@ class RecipeScanner:
return duplicate_groups
async def find_duplicate_recipes_by_source(self) -> dict:
async def find_duplicate_recipes_by_source(self) -> Dict[str, List[Any]]:
"""Find all recipe duplicates based on source_path (Civitai image URLs)
Returns:

5
py/services/recipes/analysis_service.py
View File

@@ -101,6 +101,7 @@ class RecipeAnalysisService:
temp_path = None
metadata: Optional[dict[str, Any]] = None
image_info: Optional[dict[str, Any]] = None
is_video = False
extension = ".jpg" # Default
@@ -413,7 +414,7 @@ class RecipeAnalysisService:
error_msg = "This image does not contain any generation metadata (prompt, models, or parameters)"
else:
error_msg = "No parser found for this image"
payload = {"error": error_msg, "loras": []}
payload: dict[str, Any] = {"error": error_msg, "loras": []}
if include_image_base64 and image_path:
payload["image_base64"] = self._encode_file(image_path)
payload["is_video"] = is_video
@@ -494,7 +495,7 @@ class RecipeAnalysisService:
getattr(tensor_image, "dtype", None),
)
import torch # type: ignore[import-not-found]
import torch # pyright: ignore[reportMissingImports]
if isinstance(tensor_image, torch.Tensor):
image_np = tensor_image.cpu().numpy()

8
py/services/recipes/persistence_service.py
View File

@@ -9,7 +9,7 @@ import shutil
import time
import uuid
from dataclasses import dataclass
from typing import Any, Dict, Iterable, Optional
from typing import Any, Awaitable, Dict, Iterable, Optional, cast
from ...config import config
from ...recipes.constants import GEN_PARAM_KEYS
@@ -72,6 +72,8 @@ class RecipePersistenceService:
f"Missing required fields: {', '.join(missing_fields)}"
)
assert metadata is not None
resolved_image_bytes = self._resolve_image_bytes(image_bytes, image_base64)
recipes_dir = target_dir or recipe_scanner.recipes_dir
os.makedirs(recipes_dir, exist_ok=True)
@@ -650,7 +652,9 @@ class RecipePersistenceService:
for candidate in candidates:
try:
checkpoint_info = await lookup(candidate)
checkpoint_info = await cast(
Awaitable[Any], lookup(candidate)
)
except Exception as exc:
self._logger.debug(
"Failed to lookup checkpoint %s while saving widget recipe: %s",

4
py/services/server_i18n.py
View File

@@ -55,7 +55,7 @@ class ServerI18nManager:
logger.warning(f"Locale {locale} not found, using 'en'")
self.current_locale = 'en'
def get_translation(self, key: str, params: Dict[str, Any] = None, **kwargs) -> str:
def get_translation(self, key: str, params: Dict[str, Any] | None = None, **kwargs) -> str:
"""Get translation for a key with optional parameters (supports both dict and keyword args)"""
# Merge kwargs into params for convenience
if params is None:
@@ -100,7 +100,7 @@ class ServerI18nManager:
return value
def get_available_locales(self) -> list:
def get_available_locales(self) -> list[str]:
"""Get list of available locales"""
return list(self.translations.keys())

4
py/services/service_registry.py
View File

@@ -1,3 +1,7 @@
# pyright: reportImportCycles=false
# Lazy (function-local) imports still count as static edges in basedpyright's
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
# import cycles. Breaking them would require an architectural refactor.
import asyncio
import logging
from typing import Optional, Dict, Any, TypeVar, Type

33
py/services/settings_manager.py
View File

@@ -12,6 +12,7 @@ from threading import Lock
from typing import (
Any,
Awaitable,
Coroutine,
Dict,
Iterable,
List,
@@ -411,7 +412,12 @@ class SettingsManager:
needs_library_bootstrap = not isinstance(libraries, dict) or not libraries
if not needs_library_bootstrap and top_level_has_paths and len(libraries) == 1:
if (
not needs_library_bootstrap
and top_level_has_paths
and isinstance(libraries, Mapping)
and len(libraries) == 1
):
only_library_payload = next(iter(libraries.values()))
if isinstance(only_library_payload, Mapping):
folder_payload = only_library_payload.get("folder_paths")
@@ -455,6 +461,9 @@ class SettingsManager:
):
seed_library_name = target_name
if not isinstance(libraries, dict) or not libraries:
return
sanitized_libraries: Dict[str, Dict[str, Any]] = {}
changed = False
for name, data in libraries.items():
@@ -594,7 +603,7 @@ class SettingsManager:
return payload
def _normalize_folder_paths(
self, folder_paths: Mapping[str, Iterable[str]]
self, folder_paths: Mapping[str, Any]
) -> Dict[str, List[str]]:
normalized: Dict[str, List[str]] = {}
for key, values in folder_paths.items():
@@ -623,7 +632,7 @@ class SettingsManager:
candidate_values = [values]
else:
try:
candidate_values = list(values) # type: ignore[arg-type]
candidate_values = list(values) # pyright: ignore[reportArgumentType]
except TypeError:
continue
@@ -656,7 +665,7 @@ class SettingsManager:
def _validate_folder_paths(
self,
library_name: str,
folder_paths: Mapping[str, Iterable[str]],
folder_paths: Mapping[str, Any],
) -> None:
"""Ensure folder paths do not overlap with other libraries.
@@ -1119,7 +1128,7 @@ class SettingsManager:
return []
if isinstance(value, str):
candidates: Iterable[str] = (
candidates: Iterable[Any] = (
value.replace("\n", ",").replace(";", ",").split(",")
)
elif isinstance(value, Sequence) and not isinstance(
@@ -1167,7 +1176,7 @@ class SettingsManager:
return []
if isinstance(value, str):
candidates: Iterable[str] = (
candidates: Iterable[Any] = (
value.replace("\n", ",").replace(";", ",").split(",")
)
elif isinstance(value, Sequence) and not isinstance(
@@ -1207,7 +1216,7 @@ class SettingsManager:
return []
if isinstance(value, str):
candidates: Iterable[str] = (
candidates: Iterable[Any] = (
value.replace("\n", ",").replace(";", ",").split(",")
)
elif isinstance(value, Sequence) and not isinstance(
@@ -1595,11 +1604,11 @@ class SettingsManager:
if key == "folder_paths" and isinstance(value, Mapping):
active_name = self.get_active_library_name()
self._validate_folder_paths(active_name, value)
self._update_active_library_entry(folder_paths=value) # type: ignore[arg-type]
self._update_active_library_entry(folder_paths=value) # pyright: ignore[reportArgumentType]
elif key == "extra_folder_paths" and isinstance(value, Mapping):
active_name = self.get_active_library_name()
self._validate_folder_paths(active_name, value)
self._update_active_library_entry(extra_folder_paths=value) # type: ignore[arg-type]
self._update_active_library_entry(extra_folder_paths=value) # pyright: ignore[reportArgumentType]
elif key == "default_lora_root":
self._update_active_library_entry(default_lora_root=str(value))
elif key == "default_checkpoint_root":
@@ -1752,12 +1761,12 @@ class SettingsManager:
"""Trigger cache resorting when the model name display preference updates."""
try:
from .service_registry import ServiceRegistry # type: ignore
from .service_registry import ServiceRegistry # pyright: ignore[reportImportCycles]
except Exception: # pragma: no cover - registry optional in some contexts
return
display_mode = value if isinstance(value, str) else "model_name"
pending: List[Tuple[Optional[asyncio.AbstractEventLoop], Awaitable[Any]]] = []
pending: List[Tuple[Optional[asyncio.AbstractEventLoop], Coroutine[Any, Any, Any]]] = []
def _resolve_service_loop(service: Any) -> Optional[asyncio.AbstractEventLoop]:
loop = getattr(service, "loop", None)
@@ -2118,7 +2127,7 @@ class SettingsManager:
logger.debug("Failed to apply library settings to config: %s", exc)
try:
from .service_registry import ServiceRegistry # type: ignore
from .service_registry import ServiceRegistry # pyright: ignore[reportImportCycles]
for service_name in (
"lora_scanner",

11
py/services/tag_fts_index.py
View File

@@ -18,7 +18,7 @@ import sqlite3
import threading
import time
from pathlib import Path
from typing import Dict, List, Optional, Set
from typing import Any, Dict, List, Optional, Set
from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration
@@ -87,10 +87,11 @@ class TagFTSIndex:
self._indexing_in_progress = False
self._schema_initialized = False
self._warned_not_ready = False
self._needs_rebuild = False
# Ensure directory exists
directory = os.path.dirname(self._db_path)
try:
directory = os.path.dirname(self._db_path)
if directory:
os.makedirs(directory, exist_ok=True)
except Exception as exc:
@@ -358,7 +359,7 @@ class TagFTSIndex:
finally:
self._indexing_in_progress = False
def _insert_batch(self, conn: sqlite3.Connection, rows: List[tuple]) -> None:
def _insert_batch(self, conn: sqlite3.Connection, rows: List[tuple[str, int, int, str]]) -> None:
"""Insert a batch of rows into the database.
Each row is a tuple of (tag_name, category, post_count, aliases).
@@ -443,7 +444,7 @@ class TagFTSIndex:
categories: Optional[List[int]] = None,
limit: int = 20,
offset: int = 0,
) -> List[Dict]:
) -> List[Dict[str, Any]]:
"""Search tags using FTS5 with prefix matching.
Supports alias search: if the query matches an alias rather than
@@ -530,7 +531,7 @@ class TagFTSIndex:
categories: Optional[List[int]],
limit: int,
offset: int,
) -> tuple[str, list[object]]:
) -> tuple[str, list[int | str]]:
"""Build the SQL statement and params for a tag search."""
# Escape special LIKE characters and add wildcard
query_escaped = (

3
py/services/tag_update_service.py
View File

@@ -28,7 +28,8 @@ class TagUpdateService:
metadata_path = f"{base}.metadata.json"
metadata = await metadata_loader(metadata_path)
existing_tags = list(metadata.get("tags", []))
raw_tags = metadata.get("tags", [])
existing_tags = list(raw_tags) if isinstance(raw_tags, list) else []
existing_lower = [tag.lower() for tag in existing_tags]
tags_added: List[str] = []

2
py/services/use_cases/auto_organize_use_case.py
View File

@@ -13,9 +13,11 @@ class AutoOrganizeLockProvider(Protocol):
def is_auto_organize_running(self) -> bool:
"""Return ``True`` when an auto-organize operation is in-flight."""
...
async def get_auto_organize_lock(self) -> asyncio.Lock:
"""Return the asyncio lock guarding auto-organize operations."""
...
class AutoOrganizeInProgressError(RuntimeError):

2
py/services/use_cases/bulk_metadata_refresh_use_case.py
View File

@@ -81,7 +81,7 @@ class BulkMetadataRefreshUseCase:
async def emit(status: str, **extra: Any) -> None:
if progress_callback is None:
return
payload = {
payload: Dict[str, Any] = {
"status": status,
"total": total_models,
"processed": processed,

11
py/services/use_cases/example_images/import_example_images_use_case.py
View File

@@ -5,9 +5,10 @@ from __future__ import annotations
import os
import tempfile
from contextlib import suppress
from typing import Any, Dict, List
from typing import Any, Dict, List, cast
from aiohttp import web
from aiohttp.multipart import BodyPartReader
from ....utils.example_images_processor import (
ExampleImagesImportError,
@@ -35,7 +36,8 @@ class ImportExampleImagesUseCase:
if request.content_type and "multipart/form-data" in request.content_type:
reader = await request.multipart()
first_field = await reader.next()
first_field_raw = await reader.next()
first_field = cast(BodyPartReader, first_field_raw) if first_field_raw is not None else None
if first_field and first_field.name == "model_hash":
model_hash = await first_field.text()
else:
@@ -43,7 +45,8 @@ class ImportExampleImagesUseCase:
if first_field is not None:
await self._collect_upload_file(first_field, files_to_import, temp_files)
async for field in reader:
async for raw_field in reader:
field = cast(BodyPartReader, raw_field)
if field.name == "model_hash" and not model_hash:
model_hash = await field.text()
elif field.name == "files":
@@ -53,6 +56,8 @@ class ImportExampleImagesUseCase:
model_hash = data.get("model_hash")
files_to_import = list(data.get("file_paths", []))
if not model_hash:
raise ImportExampleImagesValidationError("Missing model_hash parameter")
result = await self._processor.import_images(model_hash, files_to_import)
return result
except ExampleImagesValidationError as exc:

30
py/services/websocket_manager.py
View File

@@ -1,6 +1,6 @@
import logging
from aiohttp import web
from typing import Set, Dict, Optional
from typing import Set, Dict, Optional, Any
from uuid import uuid4
import asyncio
from datetime import datetime, timedelta
@@ -15,13 +15,13 @@ class WebSocketManager:
self._init_websockets: Set[web.WebSocketResponse] = set() # New set for initialization progress clients
self._download_websockets: Dict[str, web.WebSocketResponse] = {} # New dict for download-specific clients
# Add progress tracking dictionary
self._download_progress: Dict[str, Dict] = {}
self._download_progress: Dict[str, Dict[str, Any]] = {}
# Cache last initialization progress payloads
self._last_init_progress: Dict[str, Dict] = {}
self._last_init_progress: Dict[str, Dict[str, Any]] = {}
# Add auto-organize progress tracking
self._auto_organize_progress: Optional[Dict] = None
self._auto_organize_progress: Optional[Dict[str, Any]] = None
# Add recipe repair progress tracking
self._recipe_repair_progress: Optional[Dict] = None
self._recipe_repair_progress: Optional[Dict[str, Any]] = None
self._auto_organize_lock = asyncio.Lock()
async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
@@ -95,7 +95,7 @@ class WebSocketManager:
self.cleanup_download_progress(download_id)
logger.debug(f"Delayed cleanup completed for download {download_id}")
async def broadcast(self, data: Dict):
async def broadcast(self, data: Dict[str, Any]):
"""Broadcast message to all connected clients"""
if not self._websockets:
return
@@ -106,7 +106,7 @@ class WebSocketManager:
except Exception as e:
logger.error(f"Error sending progress: {e}")
async def broadcast_init_progress(self, data: Dict):
async def broadcast_init_progress(self, data: Dict[str, Any]):
"""Broadcast initialization progress to connected clients"""
payload = dict(data) if data else {}
@@ -145,7 +145,7 @@ class WebSocketManager:
except Exception as e:
logger.debug(f'Error sending cached initialization progress: {e}')
def _get_init_progress_key(self, data: Dict) -> str:
def _get_init_progress_key(self, data: Dict[str, Any]) -> str:
"""Return a stable key for caching initialization progress payloads"""
page_type = data.get('pageType')
if page_type:
@@ -155,7 +155,7 @@ class WebSocketManager:
return f'scanner:{scanner_type}'
return 'global'
async def broadcast_download_progress(self, download_id: str, data: Dict):
async def broadcast_download_progress(self, download_id: str, data: Dict[str, Any]):
"""Send progress update to specific download client"""
progress_entry = {
'progress': data.get('progress', 0),
@@ -183,7 +183,7 @@ class WebSocketManager:
except Exception as e:
logger.error(f"Error sending download progress: {e}")
async def broadcast_auto_organize_progress(self, data: Dict):
async def broadcast_auto_organize_progress(self, data: Dict[str, Any]):
"""Broadcast auto-organize progress to connected clients"""
# Store progress data in memory
self._auto_organize_progress = data
@@ -191,7 +191,7 @@ class WebSocketManager:
# Broadcast via WebSocket
await self.broadcast(data)
async def broadcast_recipe_repair_progress(self, data: Dict):
async def broadcast_recipe_repair_progress(self, data: Dict[str, Any]):
"""Broadcast recipe repair progress to connected clients"""
# Store progress data in memory
self._recipe_repair_progress = data
@@ -199,7 +199,7 @@ class WebSocketManager:
# Broadcast via WebSocket
await self.broadcast(data)
def get_auto_organize_progress(self) -> Optional[Dict]:
def get_auto_organize_progress(self) -> Optional[Dict[str, Any]]:
"""Get current auto-organize progress"""
return self._auto_organize_progress
@@ -207,7 +207,7 @@ class WebSocketManager:
"""Clear auto-organize progress data"""
self._auto_organize_progress = None
def get_recipe_repair_progress(self) -> Optional[Dict]:
def get_recipe_repair_progress(self) -> Optional[Dict[str, Any]]:
"""Get current recipe repair progress"""
return self._recipe_repair_progress
@@ -234,7 +234,7 @@ class WebSocketManager:
"""Get the auto-organize lock"""
return self._auto_organize_lock
def get_download_progress(self, download_id: str) -> Optional[Dict]:
def get_download_progress(self, download_id: str) -> Optional[Dict[str, Any]]:
"""Get progress information for a specific download"""
return self._download_progress.get(download_id)
@@ -255,7 +255,7 @@ class WebSocketManager:
self._download_progress.pop(download_id, None)
logger.debug(f"Cleaned up old download progress for {download_id}")
async def broadcast_cache_health_warning(self, report: 'HealthReport', page_type: str = None):
async def broadcast_cache_health_warning(self, report: 'HealthReport', page_type: Optional[str] = None):
"""
Broadcast cache health warning to frontend.

4
py/services/websocket_progress_callback.py
View File

@@ -24,6 +24,6 @@ class WebSocketProgressCallback(ProgressCallback):
class WebSocketBroadcastCallback:
"""Generic WebSocket progress callback broadcasting to all clients."""
async def on_progress(self, progress_data: Dict[str, Any]) -> None:
async def on_progress(self, payload: Dict[str, Any]) -> None:
"""Send the provided payload to all connected clients."""
await ws_manager.broadcast(progress_data)
await ws_manager.broadcast(payload)