mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 11:11:26 -03:00
Compare commits
103 Commits
v1.1.8
...
27027c4497
| Author | SHA1 | Date | |
|---|---|---|---|
| 27027c4497 | |||
| 86c85c08ec | |||
| 196c8ffc3e | |||
| cfc95ee02a | |||
| 479fa36997 | |||
| 3e1216e9bc | |||
| 007883b7d1 | |||
| dc9200a12c | |||
| d2f955266d | |||
| 8e724538bd | |||
| 6fcdeb799d | |||
| 97b9b1f62b | |||
| 4bf9a4b640 | |||
| c5088772e8 | |||
| 56acefbd6c | |||
| 5ab06c4aae | |||
| c11f4b5c68 | |||
| 86376284f4 | |||
| 2b8a2fc7d8 | |||
| f26e1b41c8 | |||
| c1671af99f | |||
| ac7707d0f6 | |||
| 381cd710a2 | |||
| ad0d18cb79 | |||
| 7980ee77d0 | |||
| 916b8bb327 | |||
| 87e3d4dea9 | |||
| 76a913f5e0 | |||
| d8c192e647 | |||
| c453437620 | |||
| 720fa6d909 | |||
| b4f71089f4 | |||
| 83e6657ead | |||
| 7ea6df4111 | |||
| d9ab92602a | |||
| 5ffadaed31 | |||
| 24f5f7df5d | |||
| daf01fb1d6 | |||
| 0f11b6def9 | |||
| 7df83f44b8 | |||
| 169fa7bed6 | |||
| 027b504fe8 | |||
| 186ef4da78 | |||
| dc674098e7 | |||
| 9087b4b07c | |||
| 8e45c22d7a | |||
| 191c4e03cd | |||
| ab4154c57d | |||
| 28e93d12ff | |||
| 75e63c758b | |||
| 823f71f269 | |||
| 042dd4088d | |||
| eaa791a9eb | |||
| 2228627ff4 | |||
| 4c647ad9c8 | |||
| 8ca3e6c33f | |||
| dd6bdbf297 | |||
| b47dde87e4 | |||
| 99e65cccd8 | |||
| 3bdacb8f46 | |||
| b4f9c224d3 | |||
| 5ec0399c81 | |||
| b464fdc333 | |||
| 53825500db | |||
| f2ac790752 | |||
| 0d8805cdee | |||
| 656e24ac9b | |||
| 6718b37403 | |||
| c9e5e784fc | |||
| f92f958682 | |||
| f63fab0676 | |||
| cfc4903c0c | |||
| a527a847fe | |||
| 91b0bf8933 | |||
| 66d1c96783 | |||
| 986128076e | |||
| 1de0a53241 | |||
| 0ec7eaf606 | |||
| d9fcb0e92b | |||
| f49b4ba4db | |||
| 84e708328b | |||
| 125bed3f09 | |||
| 077e70169d | |||
| e6dc169a05 | |||
| f34c02756d | |||
| 1e4c315481 | |||
| a8283a0d00 | |||
| 55896669fc | |||
| e341e0b9d2 | |||
| e6538c83bb | |||
| 92e1285ea5 | |||
| 2aabd1d90e | |||
| 7b8b778f83 | |||
| 7c8dc57d55 | |||
| fe95fae5f2 | |||
| ce8a95abf7 | |||
| c8e7e543d6 | |||
| a9dbb15ffa | |||
| cf64043f7d | |||
| ccaff92c18 | |||
| 585b5c922a | |||
| ea80c2224c | |||
| 8b0f56c1a6 |
@@ -25,6 +25,7 @@ model_cache/
|
||||
reasonix.toml
|
||||
.reasonix/
|
||||
.codegraph/
|
||||
.playwright-mcp/
|
||||
|
||||
# Vue widgets development cache (but keep build output)
|
||||
vue-widgets/node_modules/
|
||||
|
||||
@@ -31,7 +31,7 @@ COVERAGE_FILE=coverage/backend/.coverage pytest \
|
||||
--cov-report=xml:coverage/backend/coverage.xml
|
||||
```
|
||||
|
||||
### Frontend Development (Standalone Web UI)
|
||||
### Frontend Development (LoRA Manager Web UI)
|
||||
|
||||
```bash
|
||||
npm install
|
||||
@@ -137,7 +137,13 @@ npm run test:coverage # Generate coverage report
|
||||
- Dual mode: ComfyUI plugin (folder_paths) vs standalone (settings.json)
|
||||
- Detection: `os.environ.get("LORA_MANAGER_STANDALONE", "0") == "1"`
|
||||
- Run `python scripts/sync_translation_keys.py` after adding UI strings to `locales/en.json`
|
||||
- Symlinks require normalized paths
|
||||
- Symlinks require normalized paths.
|
||||
**Business paths vs real paths**: All stored paths and operation routing use the
|
||||
original paths as they appear under configured model roots — symlinks are NOT
|
||||
resolved. `os.path.realpath` is only for scanner dedup and the symlink cache.
|
||||
Any path passed to `os.remove`/`os.rename`/`shutil.move` or validated by a
|
||||
containment check MUST use the business path (i.e. `os.path.abspath`, not
|
||||
`realpath`).
|
||||
|
||||
## Git / Commit Messages
|
||||
|
||||
@@ -148,9 +154,9 @@ npm run test:coverage # Generate coverage report
|
||||
|
||||
## Frontend UI Architecture
|
||||
|
||||
### 1. Standalone Web UI
|
||||
### 1. LoRA Manager Web UI
|
||||
- Location: `./static/` and `./templates/`
|
||||
- Tech: Vanilla JS + CSS, served by standalone server
|
||||
- Tech: Vanilla JS + CSS, served by the hosting server (ComfyUI app in plugin mode, `standalone.py` in standalone mode)
|
||||
- Tests via npm in root directory
|
||||
|
||||
### 2. ComfyUI Custom Node Widgets
|
||||
|
||||
+10
@@ -17,6 +17,8 @@ try: # pragma: no cover - import fallback for pytest collection
|
||||
from .py.nodes.lora_cycler import LoraCyclerLM
|
||||
from .py.nodes.lora_info import LoraInfoLM
|
||||
from .py.nodes.lora_syntax_to_path import LoraSyntaxToPath
|
||||
from .py.nodes.create_hook_lora import CreateHookLoraLM
|
||||
from .py.nodes.metadata_overwrite import MetadataOverwriteLM
|
||||
from .py.metadata_collector import init as init_metadata_collector
|
||||
except (
|
||||
ImportError
|
||||
@@ -62,6 +64,12 @@ except (
|
||||
LoraSyntaxToPath = importlib.import_module(
|
||||
"py.nodes.lora_syntax_to_path"
|
||||
).LoraSyntaxToPath
|
||||
CreateHookLoraLM = importlib.import_module(
|
||||
"py.nodes.create_hook_lora"
|
||||
).CreateHookLoraLM
|
||||
MetadataOverwriteLM = importlib.import_module(
|
||||
"py.nodes.metadata_overwrite"
|
||||
).MetadataOverwriteLM
|
||||
init_metadata_collector = importlib.import_module("py.metadata_collector").init
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
@@ -83,6 +91,8 @@ NODE_CLASS_MAPPINGS = {
|
||||
LoraCyclerLM.NAME: LoraCyclerLM,
|
||||
LoraInfoLM.NAME: LoraInfoLM,
|
||||
LoraSyntaxToPath.NAME: LoraSyntaxToPath,
|
||||
CreateHookLoraLM.NAME: CreateHookLoraLM,
|
||||
MetadataOverwriteLM.NAME: MetadataOverwriteLM,
|
||||
}
|
||||
|
||||
WEB_DIRECTORY = "./web/comfyui"
|
||||
|
||||
+313
-291
File diff suppressed because it is too large
Load Diff
@@ -39,6 +39,7 @@ These fields are present in all model metadata files.
|
||||
| `metadata_source` | string\|null | ❌ No | ✅ Yes | Last provider that supplied metadata (see below) |
|
||||
| `last_checked_at` | float | ❌ No (default: `0`) | ✅ Yes | Unix timestamp of last metadata check |
|
||||
| `hash_status` | string | ❌ No (default: `"completed"`) | ✅ Yes | Hash calculation status: `"pending"`, `"calculating"`, `"completed"`, `"failed"` |
|
||||
| `autov3` | string\|null | ❌ No | ✅ Yes | CivitAI AutoV3 hash (first 12 chars, lowercase hex) sourced from the safetensors embedded metadata (`sshs_model_hash` / `modelspec.hash_sha256`). **Absent** = not yet checked (may be backfilled later); **`null`** = checked but unavailable (header has no recognized hash); **12-char hex string** = value |
|
||||
|
||||
---
|
||||
|
||||
@@ -287,6 +288,7 @@ These fields are automatically synchronized with the filesystem:
|
||||
- `preview_url` — Updated if preview file is moved/removed
|
||||
- `sha256` — Updated during hash calculation (when `hash_status="pending"`)
|
||||
- `hash_status` — Updated during hash calculation
|
||||
- `autov3` — Set when metadata is first created (from safetensors header); may be backfilled later for entries where it is absent
|
||||
- `last_checked_at` — Timestamp of scan
|
||||
- `metadata_source` — Set based on metadata provider
|
||||
|
||||
@@ -345,6 +347,7 @@ These fields can be edited by users at any time through the Lora Manager UI or b
|
||||
| `metadata_source` | `null` |
|
||||
| `last_checked_at` | `0` |
|
||||
| `hash_status` | `"completed"` |
|
||||
| `autov3` | absent (not checked) or `null` (checked, no value) |
|
||||
| `usage_tips` | `"{}"` (LoRA only) |
|
||||
| `model_type` | `"checkpoint"` or `"embedding"` (not present in LoRA models) |
|
||||
|
||||
@@ -354,6 +357,7 @@ These fields can be edited by users at any time through the Lora Manager UI or b
|
||||
|
||||
| Version | Date | Changes |
|
||||
|---------|------|---------|
|
||||
| 1.1 | 2026-08 | Added `autov3` field (CivitAI AutoV3 hash with three-state semantics) |
|
||||
| 1.0 | 2026-03 | Initial schema documentation |
|
||||
|
||||
---
|
||||
|
||||
+2248
-2205
File diff suppressed because it is too large
Load Diff
+45
-2
@@ -449,6 +449,12 @@
|
||||
"compact": "7 (1080p), 8 (2K), 10 (4K)"
|
||||
},
|
||||
"displayDensityWarning": "Warning: Higher densities may cause performance issues on systems with limited resources.",
|
||||
"recipesLayout": "Recipes Layout",
|
||||
"recipesLayoutHelp": "Choose how recipe cards are arranged: a uniform grid or a masonry (Pinterest-style) layout that preserves each image's aspect ratio.",
|
||||
"recipesLayoutOptions": {
|
||||
"grid": "Grid",
|
||||
"masonry": "Masonry"
|
||||
},
|
||||
"showFolderSidebar": "Show Folder Sidebar",
|
||||
"showFolderSidebarHelp": "Toggle the folder navigation sidebar on model pages. When disabled, the sidebar and hover area stay hidden.",
|
||||
"cardInfoDisplay": "Card Info Display",
|
||||
@@ -678,6 +684,7 @@
|
||||
"deepseek": "DeepSeek",
|
||||
"groq": "Groq",
|
||||
"openrouter": "OpenRouter",
|
||||
"google": "Gemini",
|
||||
"opencode-go": "OpenCode Go",
|
||||
"custom": "Custom (OpenAI-compatible)"
|
||||
},
|
||||
@@ -714,7 +721,9 @@
|
||||
"versionsCount": "Local Versions",
|
||||
"versionsCountDesc": "Most versions first",
|
||||
"versionsCountAsc": "Fewest versions first",
|
||||
"versionIdDesc": "Newest version first"
|
||||
"versionIdDesc": "Newest version first",
|
||||
"random": "Random",
|
||||
"randomAction": "Randomize (shuffle)"
|
||||
},
|
||||
"refresh": {
|
||||
"title": "Refresh model list",
|
||||
@@ -771,6 +780,8 @@
|
||||
"deleteAll": "Delete Selected",
|
||||
"downloadMissingLoras": "Download Missing LoRAs",
|
||||
"downloadExamples": "Download Example Images",
|
||||
"downloadMissingExamples": "Download Missing",
|
||||
"reprocessExamples": "Re-process All",
|
||||
"clear": "Clear Selection",
|
||||
"skipMetadataRefreshCount": "Skip ({count} models)",
|
||||
"resumeMetadataRefreshCount": "Resume ({count} models)",
|
||||
@@ -806,6 +817,8 @@
|
||||
"sendToWorkflowReplace": "Send to Workflow (Replace)",
|
||||
"openExamples": "Open Examples Folder",
|
||||
"downloadExamples": "Download Example Images",
|
||||
"downloadMissingExamples": "Download Missing",
|
||||
"reprocessExamples": "Re-process All",
|
||||
"replacePreview": "Replace Preview",
|
||||
"setContentRating": "Set Content Rating",
|
||||
"moveToFolder": "Move to Folder",
|
||||
@@ -1548,6 +1561,7 @@
|
||||
"empty": "No version history available for this model yet.",
|
||||
"error": "Failed to load versions.",
|
||||
"missingModelId": "This model is missing a Civitai model id.",
|
||||
"hfGroupInfo": "This is a HuggingFace model group. Open the library to see all versions in the grid.",
|
||||
"confirm": {
|
||||
"delete": "Delete this version from your library?"
|
||||
},
|
||||
@@ -1574,6 +1588,21 @@
|
||||
"downloadCsv": "Download CSV",
|
||||
"columnModelName": "Model Name",
|
||||
"columnError": "Error"
|
||||
},
|
||||
"downloadBatchSummary": {
|
||||
"title": "Batch Download Summary",
|
||||
"statSuccess": "Success",
|
||||
"statFailed": "Failed",
|
||||
"statTotal": "Total",
|
||||
"successMessage": "All {count} models downloaded successfully",
|
||||
"completedWithErrors": "Completed with errors",
|
||||
"failed": "Download failed",
|
||||
"failedItems": "Failed Items ({count})",
|
||||
"columnName": "Model Name",
|
||||
"columnError": "Error",
|
||||
"close": "Close",
|
||||
"copyReport": "Copy Report",
|
||||
"retryFailed": "Retry Failed ({count})"
|
||||
}
|
||||
},
|
||||
"modelTags": {
|
||||
@@ -1751,6 +1780,12 @@
|
||||
"checkingMessage": "Please wait while we check for the latest version.",
|
||||
"showNotifications": "Show update notifications",
|
||||
"latestBadge": "Latest",
|
||||
"latestMain": "Latest main",
|
||||
"channel": "Update Channel",
|
||||
"channels": {
|
||||
"release": "Release",
|
||||
"nightly": "Nightly"
|
||||
},
|
||||
"updateProgress": {
|
||||
"preparing": "Preparing update...",
|
||||
"installing": "Installing update...",
|
||||
@@ -1771,6 +1806,15 @@
|
||||
"warning": "Warning: Nightly builds may contain experimental features and could be unstable.",
|
||||
"enable": "Enable Nightly Updates"
|
||||
},
|
||||
"channelSwitch": {
|
||||
"nightlyTitle": "Switch to Nightly Channel",
|
||||
"nightlyMessage": "Switching to Nightly will initialize a Git repository and track the latest main branch commits. Updates will be more frequent but may be unstable. You can switch back to Release at any time.",
|
||||
"releaseTitle": "Switch to Release Channel",
|
||||
"releaseMessage": "Switching to Release will checkout the latest stable release tag. You can switch back to Nightly at any time.",
|
||||
"switching": "Switching to {channel} channel...",
|
||||
"completed": "Successfully switched to {channel} channel",
|
||||
"failed": "Failed to switch channel"
|
||||
},
|
||||
"banners": {
|
||||
"recent": "Recent messages",
|
||||
"empty": "No recent banners yet.",
|
||||
@@ -2015,7 +2059,6 @@
|
||||
"presetNameTooLong": "Preset name must be {max} characters or less",
|
||||
"presetNameInvalidChars": "Preset name contains invalid characters",
|
||||
"presetNameExists": "A preset with this name already exists",
|
||||
"maxPresetsReached": "Maximum {max} presets allowed. Delete one to add more.",
|
||||
"presetNotFound": "Preset not found",
|
||||
"invalidPreset": "Invalid preset data",
|
||||
"deletePresetFailed": "Failed to delete preset",
|
||||
|
||||
+2248
-2205
File diff suppressed because it is too large
Load Diff
+2248
-2205
File diff suppressed because it is too large
Load Diff
+2248
-2205
File diff suppressed because it is too large
Load Diff
+2248
-2205
File diff suppressed because it is too large
Load Diff
+2248
-2205
File diff suppressed because it is too large
Load Diff
+2248
-2205
File diff suppressed because it is too large
Load Diff
+2248
-2205
File diff suppressed because it is too large
Load Diff
+2248
-2205
File diff suppressed because it is too large
Load Diff
+15
-10
@@ -1,9 +1,13 @@
|
||||
# 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 platform
|
||||
import posixpath
|
||||
import threading
|
||||
from pathlib import Path
|
||||
import folder_paths # type: ignore
|
||||
import folder_paths # pyright: ignore[reportMissingImports]
|
||||
from typing import Any, Dict, Iterable, List, Mapping, Optional, Set, Tuple
|
||||
import logging
|
||||
import json
|
||||
@@ -90,7 +94,7 @@ def _resolve_valid_default_root(
|
||||
|
||||
|
||||
def _normalize_folder_paths_for_comparison(
|
||||
folder_paths: Mapping[str, Iterable[str]],
|
||||
folder_paths: Mapping[str, Any],
|
||||
) -> Dict[str, Set[str]]:
|
||||
"""Normalize folder paths for comparison across libraries."""
|
||||
|
||||
@@ -482,7 +486,7 @@ class Config:
|
||||
import ctypes
|
||||
|
||||
FILE_ATTRIBUTE_REPARSE_POINT = 0x400
|
||||
attrs = ctypes.windll.kernel32.GetFileAttributesW(str(path)) # type: ignore[attr-defined]
|
||||
attrs = ctypes.windll.kernel32.GetFileAttributesW(str(path)) # pyright: ignore[reportAttributeAccessIssue]
|
||||
return attrs != -1 and (attrs & FILE_ATTRIBUTE_REPARSE_POINT)
|
||||
except Exception as e:
|
||||
logger.error(f"Error checking Windows reparse point: {e}")
|
||||
@@ -491,7 +495,7 @@ class Config:
|
||||
logger.error(f"Error checking link status for {path}: {e}")
|
||||
return False
|
||||
|
||||
def _entry_is_symlink(self, entry: os.DirEntry) -> bool:
|
||||
def _entry_is_symlink(self, entry: os.DirEntry[str]) -> bool:
|
||||
"""Check if a directory entry is a symlink, including Windows junctions."""
|
||||
if entry.is_symlink():
|
||||
return True
|
||||
@@ -500,7 +504,7 @@ class Config:
|
||||
import ctypes
|
||||
|
||||
FILE_ATTRIBUTE_REPARSE_POINT = 0x400
|
||||
attrs = ctypes.windll.kernel32.GetFileAttributesW(entry.path) # type: ignore[attr-defined]
|
||||
attrs = ctypes.windll.kernel32.GetFileAttributesW(entry.path) # pyright: ignore[reportAttributeAccessIssue]
|
||||
return attrs != -1 and (attrs & FILE_ATTRIBUTE_REPARSE_POINT)
|
||||
except Exception:
|
||||
pass
|
||||
@@ -1126,8 +1130,8 @@ class Config:
|
||||
|
||||
def _apply_library_paths(
|
||||
self,
|
||||
folder_paths: Mapping[str, Iterable[str]],
|
||||
extra_folder_paths: Optional[Mapping[str, Iterable[str]]] = None,
|
||||
folder_paths: Mapping[str, Any],
|
||||
extra_folder_paths: Optional[Mapping[str, Any]] = None,
|
||||
recipes_path: str = "",
|
||||
) -> None:
|
||||
self._path_mappings.clear()
|
||||
@@ -1432,12 +1436,13 @@ class Config:
|
||||
# ('_lm_config_cache') that is NEVER removed from sys.modules (its key does
|
||||
# NOT start with 'py.'), so it survives re-imports of py.* modules.
|
||||
_CONFIG_SENTINEL = "_lm_config_cache"
|
||||
config: Config
|
||||
if _CONFIG_SENTINEL in _sys.modules:
|
||||
# Re-import: reuse the existing singleton from the sentinel.
|
||||
config: Config = _sys.modules[_CONFIG_SENTINEL].config # type: ignore[valid-type]
|
||||
config = _sys.modules[_CONFIG_SENTINEL].config
|
||||
else:
|
||||
config: Config = Config()
|
||||
config = Config()
|
||||
# Register the sentinel so re-imports of py.config find us.
|
||||
_sentinel_mod = _types.ModuleType(_CONFIG_SENTINEL)
|
||||
_sentinel_mod.config = config
|
||||
setattr(_sentinel_mod, "config", config)
|
||||
_sys.modules[_CONFIG_SENTINEL] = _sentinel_mod
|
||||
|
||||
+1
-1
@@ -14,7 +14,7 @@ standalone_mode = (
|
||||
if not standalone_mode:
|
||||
setup_logging()
|
||||
|
||||
from server import PromptServer # type: ignore
|
||||
from server import PromptServer # pyright: ignore[reportMissingImports]
|
||||
|
||||
from .config import config
|
||||
from .services.model_service_factory import (
|
||||
|
||||
@@ -22,7 +22,7 @@ if not standalone_mode:
|
||||
|
||||
logger.info("ComfyUI Metadata Collector initialized")
|
||||
|
||||
def get_metadata(prompt_id=None): # type: ignore[no-redef]
|
||||
def get_metadata(prompt_id=None): # pyright: ignore[reportRedeclaration]
|
||||
"""Helper function to get metadata from the registry"""
|
||||
registry = MetadataRegistry()
|
||||
return registry.get_metadata(prompt_id)
|
||||
@@ -31,6 +31,6 @@ else:
|
||||
def init():
|
||||
logger.info("ComfyUI Metadata Collector disabled in standalone mode")
|
||||
|
||||
def get_metadata(prompt_id=None): # type: ignore[no-redef]
|
||||
def get_metadata(prompt_id=None): # pyright: ignore[reportRedeclaration]
|
||||
"""Dummy implementation for standalone mode"""
|
||||
return {}
|
||||
|
||||
@@ -1,5 +1,11 @@
|
||||
"""Constants used by the metadata collector"""
|
||||
|
||||
# Sentinel value for clip_skip to distinguish "unconnected / widget default"
|
||||
# from "user wired value 0". Both ComfyUI CLIPSetLastLayer (-24..-1) and
|
||||
# A1111 conventions treat 0 as meaningless for clip skipping, but users may
|
||||
# explicitly wire 0 to the overwrite node to express "no clip skip / default".
|
||||
CLIP_SKIP_SENTINEL = -25
|
||||
|
||||
# Metadata categories
|
||||
MODELS = "models"
|
||||
PROMPTS = "prompts"
|
||||
@@ -9,6 +15,14 @@ EMBEDDINGS = "embeddings"
|
||||
SIZE = "size"
|
||||
IMAGES = "images"
|
||||
IS_SAMPLER = "is_sampler" # New constant to mark sampler nodes
|
||||
OVERWRITE = "overwrite" # Manual metadata overwrite from MetadataOverwriteLM node
|
||||
|
||||
# Field names that the MetadataOverwriteLM node and its extractor share
|
||||
METADATA_OVERWRITE_FIELDS = (
|
||||
"prompt", "negative_prompt", "seed", "steps", "cfg_scale",
|
||||
"sampler", "scheduler", "model", "loras", "size",
|
||||
"clip_skip", "additional_data",
|
||||
)
|
||||
|
||||
# Complete list of categories to track
|
||||
METADATA_CATEGORIES = [MODELS, PROMPTS, SAMPLING, LORAS, EMBEDDINGS, SIZE, IMAGES]
|
||||
METADATA_CATEGORIES = [MODELS, PROMPTS, SAMPLING, LORAS, EMBEDDINGS, SIZE, IMAGES, OVERWRITE]
|
||||
|
||||
@@ -16,7 +16,7 @@ class MetadataHook:
|
||||
execution = None
|
||||
try:
|
||||
# Try direct import first
|
||||
import execution # type: ignore
|
||||
import execution # pyright: ignore[reportMissingImports]
|
||||
except ImportError:
|
||||
# Try to locate from system modules
|
||||
for module_name in sys.modules:
|
||||
@@ -83,7 +83,8 @@ class MetadataHook:
|
||||
|
||||
# Record inputs before execution
|
||||
if node_id is not None:
|
||||
registry.record_node_execution(node_id, class_type, input_data_all, None)
|
||||
return_types = getattr(obj, 'RETURN_TYPES', None)
|
||||
registry.record_node_execution(node_id, class_type, input_data_all, None, return_types=return_types)
|
||||
except Exception as e:
|
||||
logger.error(f"Error collecting metadata (pre-execution): {str(e)}")
|
||||
|
||||
@@ -114,7 +115,8 @@ class MetadataHook:
|
||||
|
||||
# Record outputs after execution
|
||||
if node_id is not None:
|
||||
registry.update_node_execution(node_id, class_type, results)
|
||||
return_types = getattr(obj, 'RETURN_TYPES', None)
|
||||
registry.update_node_execution(node_id, class_type, results, return_types=return_types)
|
||||
except Exception as e:
|
||||
logger.error(f"Error collecting metadata (post-execution): {str(e)}")
|
||||
|
||||
@@ -135,10 +137,13 @@ class MetadataHook:
|
||||
# Store the dynprompt reference for node lookups
|
||||
if hasattr(prompt, 'original_prompt'):
|
||||
registry.set_current_prompt(prompt)
|
||||
|
||||
|
||||
# Store extra_data for accessing full workflow node properties
|
||||
registry.set_extra_data(extra_data)
|
||||
|
||||
# Execute the original function
|
||||
return original_execute(*args, **kwargs)
|
||||
|
||||
|
||||
# Replace the functions
|
||||
execution._map_node_over_list = map_node_over_list_with_metadata
|
||||
execution.execute = execute_with_prompt_tracking
|
||||
@@ -163,7 +168,8 @@ class MetadataHook:
|
||||
class_type = obj.__class__.__name__
|
||||
node_id = unique_id
|
||||
if node_id is not None:
|
||||
registry.record_node_execution(node_id, class_type, input_data_all, None)
|
||||
return_types = getattr(obj, 'RETURN_TYPES', None)
|
||||
registry.record_node_execution(node_id, class_type, input_data_all, None, return_types=return_types)
|
||||
except Exception as e:
|
||||
logger.error(f"Error collecting metadata (pre-execution): {str(e)}")
|
||||
|
||||
@@ -180,7 +186,8 @@ class MetadataHook:
|
||||
class_type = obj.__class__.__name__
|
||||
node_id = unique_id
|
||||
if node_id is not None:
|
||||
registry.update_node_execution(node_id, class_type, results)
|
||||
return_types = getattr(obj, 'RETURN_TYPES', None)
|
||||
registry.update_node_execution(node_id, class_type, results, return_types=return_types)
|
||||
except Exception as e:
|
||||
logger.error(f"Error collecting metadata (post-execution): {str(e)}")
|
||||
|
||||
@@ -202,6 +209,9 @@ class MetadataHook:
|
||||
if hasattr(prompt, 'original_prompt'):
|
||||
registry.set_current_prompt(prompt)
|
||||
|
||||
# Store extra_data for accessing full workflow node properties
|
||||
registry.set_extra_data(extra_data)
|
||||
|
||||
# Execute the original function
|
||||
return await original_execute(*args, **kwargs)
|
||||
|
||||
|
||||
@@ -1,15 +1,68 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from .constants import IMAGES
|
||||
|
||||
# Check if running in standalone mode
|
||||
standalone_mode = os.environ.get("LORA_MANAGER_STANDALONE", "0") == "1" or os.environ.get("HF_HUB_DISABLE_TELEMETRY", "0") == "0"
|
||||
|
||||
from .constants import MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IS_SAMPLER
|
||||
from .constants import MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IS_SAMPLER, OVERWRITE
|
||||
from .node_extractors import NODE_EXTRACTORS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Keys that identify metadata hint marks stored in node.properties.lm_marker_role
|
||||
_META_MARK_PREFIX = "meta_"
|
||||
_MARK_PRIMARY_MODEL = "primary_model"
|
||||
_MARK_PRIMARY_SAMPLER = "primary_sampler"
|
||||
_MARK_POSITIVE_PROMPT = "positive_prompt"
|
||||
_MARK_NEGATIVE_PROMPT = "negative_prompt"
|
||||
|
||||
class MetadataProcessor:
|
||||
"""Process and format collected metadata"""
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _get_user_marks(metadata):
|
||||
"""Scan workflow nodes (from extra_data.extra_pnginfo.workflow) for user-assigned
|
||||
metadata hint marks stored in node.properties.lm_marker_role.
|
||||
|
||||
Returns a dict mapping mark type keys to node IDs.
|
||||
Example: {'primary_model': '42', 'primary_sampler': '17'}
|
||||
"""
|
||||
marks: dict[str, str] = {}
|
||||
|
||||
# Primary source: extra_data.extra_pnginfo.workflow.nodes (has full properties)
|
||||
extra_data = metadata.get("extra_data")
|
||||
if extra_data and isinstance(extra_data, dict):
|
||||
extra_pnginfo = extra_data.get("extra_pnginfo", {})
|
||||
if isinstance(extra_pnginfo, dict):
|
||||
workflow = extra_pnginfo.get("workflow", {})
|
||||
nodes = workflow.get("nodes", [])
|
||||
for node in nodes:
|
||||
node_id = str(node.get("id", ""))
|
||||
role = node.get("properties", {}).get("lm_marker_role", "")
|
||||
if role.startswith(_META_MARK_PREFIX):
|
||||
mark_type = role[len(_META_MARK_PREFIX):]
|
||||
if mark_type in marks:
|
||||
logger.warning(
|
||||
"Duplicate meta hint '%s': node %s (previous: %s), "
|
||||
"last match wins",
|
||||
mark_type, node_id, marks[mark_type],
|
||||
)
|
||||
marks[mark_type] = node_id
|
||||
|
||||
# Fallback: try prompt.original_prompt (API-only submissions may not have workflow)
|
||||
if not marks:
|
||||
prompt = metadata.get("current_prompt")
|
||||
if prompt and getattr(prompt, "original_prompt", None):
|
||||
for node_id, node_data in prompt.original_prompt.items():
|
||||
role = node_data.get("properties", {}).get("lm_marker_role", "")
|
||||
if role.startswith(_META_MARK_PREFIX):
|
||||
mark_type = role[len(_META_MARK_PREFIX):]
|
||||
marks[mark_type] = node_id
|
||||
|
||||
return marks
|
||||
|
||||
@staticmethod
|
||||
def find_primary_sampler(metadata, downstream_id=None):
|
||||
"""
|
||||
@@ -471,20 +524,57 @@ class MetadataProcessor:
|
||||
"checkpoint": None,
|
||||
"loras": "",
|
||||
"size": None,
|
||||
"clip_skip": None
|
||||
"clip_skip": None,
|
||||
"additional_data": "",
|
||||
}
|
||||
|
||||
# Get the prompt object for node relationship tracing
|
||||
prompt = metadata.get("current_prompt")
|
||||
|
||||
# Find the primary KSampler node
|
||||
primary_sampler_id, primary_sampler = MetadataProcessor.find_primary_sampler(metadata, id)
|
||||
|
||||
# Directly get checkpoint from metadata instead of tracing
|
||||
# Pass primary_sampler_id to avoid redundant calculation
|
||||
checkpoint = MetadataProcessor.find_primary_checkpoint(metadata, id, primary_sampler_id)
|
||||
if checkpoint:
|
||||
params["checkpoint"] = checkpoint
|
||||
|
||||
# ---- User marks: override heuristic inference with user-assigned hints ----
|
||||
user_marks = MetadataProcessor._get_user_marks(metadata)
|
||||
|
||||
# Find the primary KSampler node (user mark takes priority)
|
||||
primary_sampler_id = None
|
||||
primary_sampler = None
|
||||
if _MARK_PRIMARY_SAMPLER in user_marks:
|
||||
marked_id = user_marks[_MARK_PRIMARY_SAMPLER]
|
||||
sampler_data = metadata.get(SAMPLING, {}).get(marked_id)
|
||||
if sampler_data and sampler_data.get(IS_SAMPLER):
|
||||
primary_sampler_id = marked_id
|
||||
primary_sampler = sampler_data
|
||||
else:
|
||||
logger.warning(
|
||||
"User-marked primary sampler %s has no runtime metadata, "
|
||||
"falling back to heuristic",
|
||||
marked_id,
|
||||
)
|
||||
if primary_sampler is None:
|
||||
primary_sampler_id, primary_sampler = MetadataProcessor.find_primary_sampler(metadata, id)
|
||||
|
||||
# Resolve checkpoint / model (user mark takes priority)
|
||||
if _MARK_PRIMARY_MODEL in user_marks:
|
||||
marked_id = user_marks[_MARK_PRIMARY_MODEL]
|
||||
if marked_id in metadata.get(MODELS, {}):
|
||||
params["checkpoint"] = metadata[MODELS][marked_id].get("name")
|
||||
else:
|
||||
extra_data = metadata.get("extra_data")
|
||||
extra_pnginfo = extra_data.get("extra_pnginfo", {}) if extra_data and isinstance(extra_data, dict) else {}
|
||||
workflow = extra_pnginfo.get("workflow", {}) if isinstance(extra_pnginfo, dict) else {}
|
||||
node_type = "unknown"
|
||||
for n in workflow.get("nodes", []):
|
||||
if str(n.get("id", "")) == marked_id:
|
||||
node_type = n.get("type", "unknown")
|
||||
break
|
||||
logger.warning(
|
||||
"User-marked primary model %s (type=%s, registered=%s) has no runtime metadata, "
|
||||
"falling back to heuristic",
|
||||
marked_id, node_type, node_type in NODE_EXTRACTORS,
|
||||
)
|
||||
if params["checkpoint"] is None:
|
||||
checkpoint = MetadataProcessor.find_primary_checkpoint(metadata, id, primary_sampler_id)
|
||||
if checkpoint:
|
||||
params["checkpoint"] = checkpoint
|
||||
|
||||
# Check if guidance parameter exists in any sampling node
|
||||
for node_id, sampler_info in metadata.get(SAMPLING, {}).items():
|
||||
@@ -539,7 +629,22 @@ class MetadataProcessor:
|
||||
|
||||
# For SamplerCustom, handle any additional parameters
|
||||
MetadataProcessor.handle_custom_advanced_sampler(metadata, prompt, primary_sampler_id, params)
|
||||
|
||||
|
||||
# ---- User marks: override prompts with explicitly tagged nodes ----
|
||||
prompts_data = metadata.get(PROMPTS, {})
|
||||
if _MARK_POSITIVE_PROMPT in user_marks:
|
||||
pos_id = user_marks[_MARK_POSITIVE_PROMPT]
|
||||
if pos_id in prompts_data:
|
||||
prompt_text = prompts_data[pos_id].get("text") or prompts_data[pos_id].get("positive_text")
|
||||
if prompt_text:
|
||||
params["prompt"] = prompt_text
|
||||
if _MARK_NEGATIVE_PROMPT in user_marks:
|
||||
neg_id = user_marks[_MARK_NEGATIVE_PROMPT]
|
||||
if neg_id in prompts_data:
|
||||
prompt_text = prompts_data[neg_id].get("text") or prompts_data[neg_id].get("negative_text")
|
||||
if prompt_text:
|
||||
params["negative_prompt"] = prompt_text
|
||||
|
||||
# Size extraction is same for all sampler types
|
||||
# Check if the sampler itself has size information (from latent_image)
|
||||
if primary_sampler_id in metadata.get(SIZE, {}):
|
||||
@@ -568,7 +673,26 @@ class MetadataProcessor:
|
||||
break
|
||||
if params["clip_skip"] is None:
|
||||
params["clip_skip"] = "1"
|
||||
|
||||
|
||||
# ---- Apply manual metadata overwrites ----
|
||||
for overwrite_info in metadata.get(OVERWRITE, {}).values():
|
||||
overwrite_params = overwrite_info.get("parameters", {})
|
||||
for key, value in overwrite_params.items():
|
||||
if key == "clip_skip":
|
||||
# Accept any value from overwrite node (sentinel -25 already
|
||||
# filtered upstream). Needed because falsy check treats 0
|
||||
# as "not set" even though 0 is a valid wired input here.
|
||||
params[key] = value
|
||||
elif value: # truthy check — only overwrite when user provided a real value
|
||||
params[key] = value
|
||||
|
||||
# Bridge: the overwrite node exposes the field as "model" (more accurate),
|
||||
# but the internal pipeline key remains "checkpoint" for backward compatibility
|
||||
# with A1111 metadata format and downstream consumers.
|
||||
if params.get("model"):
|
||||
params["checkpoint"] = params["model"]
|
||||
del params["model"]
|
||||
|
||||
return params
|
||||
|
||||
@staticmethod
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
import time
|
||||
from nodes import NODE_CLASS_MAPPINGS # type: ignore
|
||||
from typing import Any
|
||||
from nodes import NODE_CLASS_MAPPINGS # pyright: ignore[reportMissingImports, reportAttributeAccessIssue]
|
||||
from .node_extractors import NODE_EXTRACTORS, GenericNodeExtractor
|
||||
from .constants import METADATA_CATEGORIES, IMAGES
|
||||
from .constants import METADATA_CATEGORIES, IMAGES, OVERWRITE
|
||||
|
||||
|
||||
class MetadataRegistry:
|
||||
@@ -9,6 +10,15 @@ class MetadataRegistry:
|
||||
|
||||
_instance = None
|
||||
|
||||
current_prompt_id: Any = None
|
||||
current_prompt: Any = None
|
||||
metadata: dict[str, Any] = {}
|
||||
prompt_metadata: dict[str, Any] = {}
|
||||
executed_nodes: set[str] = set()
|
||||
node_cache: dict[str, Any] = {}
|
||||
max_prompt_history: int = 3
|
||||
metadata_categories: list[str] = METADATA_CATEGORIES
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
@@ -61,6 +71,7 @@ class MetadataRegistry:
|
||||
{
|
||||
"execution_order": [],
|
||||
"current_prompt": None, # Will store the prompt object
|
||||
"extra_data": None, # Will store the API extra_data for workflow metadata
|
||||
"timestamp": time.time(),
|
||||
}
|
||||
)
|
||||
@@ -75,6 +86,11 @@ class MetadataRegistry:
|
||||
# Store the prompt in the metadata for later relationship tracing
|
||||
self.prompt_metadata[self.current_prompt_id]["current_prompt"] = prompt
|
||||
|
||||
def set_extra_data(self, extra_data):
|
||||
"""Store the API extra_data (contains extra_pnginfo.workflow with node properties)"""
|
||||
if self.current_prompt_id and self.current_prompt_id in self.prompt_metadata:
|
||||
self.prompt_metadata[self.current_prompt_id]["extra_data"] = extra_data
|
||||
|
||||
def get_metadata(self, prompt_id=None):
|
||||
"""Get collected metadata for a prompt"""
|
||||
key = prompt_id if prompt_id is not None else self.current_prompt_id
|
||||
@@ -122,20 +138,28 @@ class MetadataRegistry:
|
||||
cache_key = f"{node_id}:{class_type}"
|
||||
|
||||
# Check if this node type is relevant for metadata collection
|
||||
if class_type in NODE_EXTRACTORS:
|
||||
if class_type in NODE_EXTRACTORS or cache_key in self.node_cache:
|
||||
# Check if we have cached metadata for this node
|
||||
if cache_key in self.node_cache:
|
||||
cached_data = self.node_cache[cache_key]
|
||||
|
||||
# Detect bypass (mode=4) / mute (mode=2) — these nodes
|
||||
# were intentionally disabled and should not contribute
|
||||
# overwrite values from a previous execution's cache.
|
||||
node_mode = node_data.get("mode", 0)
|
||||
node_is_disabled = node_mode in (2, 4)
|
||||
|
||||
# Apply cached metadata to the current metadata
|
||||
for category in self.metadata_categories:
|
||||
if category == OVERWRITE and node_is_disabled:
|
||||
continue
|
||||
if category in cached_data and node_id in cached_data[category]:
|
||||
if node_id not in metadata[category]:
|
||||
metadata[category][node_id] = cached_data[category][
|
||||
node_id
|
||||
]
|
||||
|
||||
def record_node_execution(self, node_id, class_type, inputs, outputs):
|
||||
def record_node_execution(self, node_id, class_type, inputs, outputs, return_types=None):
|
||||
"""Record information about a node's execution"""
|
||||
if not self.current_prompt_id:
|
||||
return
|
||||
@@ -158,17 +182,18 @@ class MetadataRegistry:
|
||||
|
||||
# Extract node-specific metadata
|
||||
extractor = NODE_EXTRACTORS.get(class_type, GenericNodeExtractor)
|
||||
extractor.extract(
|
||||
node_id,
|
||||
processed_inputs,
|
||||
outputs,
|
||||
self.prompt_metadata[self.current_prompt_id],
|
||||
)
|
||||
if extractor is GenericNodeExtractor:
|
||||
extractor.extract(node_id, processed_inputs, outputs,
|
||||
self.prompt_metadata[self.current_prompt_id],
|
||||
return_types=return_types)
|
||||
else:
|
||||
extractor.extract(node_id, processed_inputs, outputs,
|
||||
self.prompt_metadata[self.current_prompt_id])
|
||||
|
||||
# Cache this node's metadata
|
||||
self._cache_node_metadata(node_id, class_type)
|
||||
|
||||
def update_node_execution(self, node_id, class_type, outputs):
|
||||
def update_node_execution(self, node_id, class_type, outputs, return_types=None):
|
||||
"""Update node metadata with output information"""
|
||||
if not self.current_prompt_id:
|
||||
return
|
||||
@@ -179,9 +204,17 @@ class MetadataRegistry:
|
||||
# Use the same extractor to update with outputs
|
||||
extractor = NODE_EXTRACTORS.get(class_type, GenericNodeExtractor)
|
||||
if hasattr(extractor, "update"):
|
||||
extractor.update(
|
||||
node_id, processed_outputs, self.prompt_metadata[self.current_prompt_id]
|
||||
)
|
||||
if extractor is GenericNodeExtractor:
|
||||
extractor.update(
|
||||
node_id, processed_outputs,
|
||||
self.prompt_metadata[self.current_prompt_id],
|
||||
return_types=return_types,
|
||||
)
|
||||
else:
|
||||
extractor.update(
|
||||
node_id, processed_outputs,
|
||||
self.prompt_metadata[self.current_prompt_id],
|
||||
)
|
||||
|
||||
# Update the cached metadata for this node
|
||||
self._cache_node_metadata(node_id, class_type)
|
||||
|
||||
@@ -2,7 +2,8 @@ import json
|
||||
import os
|
||||
import re
|
||||
|
||||
from .constants import MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IMAGES, IS_SAMPLER
|
||||
from .constants import MODELS, PROMPTS, SAMPLING, LORAS, SIZE, IMAGES, IS_SAMPLER, OVERWRITE
|
||||
from .overwrite_utils import collect_overwrite_params
|
||||
|
||||
|
||||
def _store_checkpoint_metadata(metadata, node_id, model_name):
|
||||
@@ -31,11 +32,78 @@ class NodeMetadataExtractor:
|
||||
pass
|
||||
|
||||
class GenericNodeExtractor(NodeMetadataExtractor):
|
||||
"""Default extractor for nodes without specific handling"""
|
||||
"""Fallback extractor with type-signature-based detection.
|
||||
|
||||
When a node is not in the NODE_EXTRACTORS registry, the hook layer
|
||||
passes ``return_types`` from ``obj.RETURN_TYPES``:
|
||||
|
||||
* ``MODEL`` output: common input fields (ckpt_name, unet_name, etc.)
|
||||
are checked for a model file name and stored as checkpoint metadata.
|
||||
* ``CONDITIONING`` output: common text input fields are checked for
|
||||
prompt text and stored as prompt metadata.
|
||||
"""
|
||||
|
||||
# Input field names that carry a model path in loader-style nodes.
|
||||
_MODEL_NAME_FIELDS = (
|
||||
"ckpt_name", "unet_name", "model_path", "model_name", "gguf_name",
|
||||
)
|
||||
|
||||
# Extensions used by checkpoint_scanner.py — only record values that look
|
||||
# like real model filenames to avoid capturing unrelated string fields.
|
||||
_MODEL_EXTENSIONS = {
|
||||
".ckpt", ".pt", ".pt2", ".bin", ".pth", ".safetensors", ".pkl", ".sft", ".gguf",
|
||||
}
|
||||
|
||||
# Input field names that may carry prompt text in encoder-style nodes.
|
||||
_TEXT_FIELDS = ("text", "clip_l", "t5xxl", "prompt", "positive", "negative")
|
||||
|
||||
@staticmethod
|
||||
def extract(node_id, inputs, outputs, metadata):
|
||||
pass
|
||||
|
||||
def extract(node_id, inputs, outputs, metadata, return_types=None):
|
||||
if return_types is None:
|
||||
return
|
||||
|
||||
# — MODEL loader detection (checkpoint / UNET / GGUF) —
|
||||
if "MODEL" in return_types or any("MODEL" in str(t) for t in return_types):
|
||||
for field in GenericNodeExtractor._MODEL_NAME_FIELDS:
|
||||
val = inputs.get(field)
|
||||
if val and isinstance(val, str) and val.strip():
|
||||
name = val.strip()
|
||||
if not any(name.lower().endswith(ext) for ext in GenericNodeExtractor._MODEL_EXTENSIONS):
|
||||
continue
|
||||
_store_checkpoint_metadata(metadata, node_id, name)
|
||||
return
|
||||
|
||||
# — CONDITIONING encoder detection (CLIPTextEncode, Flux, custom) —
|
||||
if "CONDITIONING" in return_types or any("CONDITIONING" in str(t) for t in return_types):
|
||||
text = None
|
||||
for field in GenericNodeExtractor._TEXT_FIELDS:
|
||||
val = inputs.get(field)
|
||||
if val and isinstance(val, str) and val.strip():
|
||||
text = val.strip()
|
||||
break
|
||||
if text:
|
||||
prompt_data = metadata.setdefault(PROMPTS, {})
|
||||
prompt_data[node_id] = {
|
||||
"text": text,
|
||||
"node_id": node_id,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def update(node_id, outputs, metadata, return_types=None):
|
||||
if return_types is None:
|
||||
return
|
||||
if "CONDITIONING" not in return_types and not any(
|
||||
"CONDITIONING" in str(t) for t in return_types
|
||||
):
|
||||
return
|
||||
if node_id not in metadata.get(PROMPTS, {}):
|
||||
return
|
||||
if outputs and isinstance(outputs, list) and len(outputs) > 0:
|
||||
if isinstance(outputs[0], tuple) and len(outputs[0]) > 0:
|
||||
cond = outputs[0][0]
|
||||
if cond is not None:
|
||||
metadata[PROMPTS][node_id]["conditioning"] = cond
|
||||
|
||||
class CheckpointLoaderExtractor(NodeMetadataExtractor):
|
||||
@staticmethod
|
||||
def extract(node_id, inputs, outputs, metadata):
|
||||
@@ -1154,6 +1222,28 @@ class CR_ApplyControlNetStackExtractor(NodeMetadataExtractor):
|
||||
metadata[PROMPTS][node_id]["positive_encoded"] = transformed_positive
|
||||
metadata[PROMPTS][node_id]["negative_encoded"] = transformed_negative
|
||||
|
||||
class MetadataOverwriteExtractor(NodeMetadataExtractor):
|
||||
"""Extract manually specified metadata from MetadataOverwriteLM node.
|
||||
|
||||
Stores truthy input values under the OVERWRITE category so that
|
||||
extract_generation_params can merge them over the inferred params.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def extract(node_id, inputs, outputs, metadata):
|
||||
if not inputs:
|
||||
return
|
||||
|
||||
overwrite_params = collect_overwrite_params(inputs)
|
||||
|
||||
if overwrite_params:
|
||||
metadata.setdefault(OVERWRITE, {})
|
||||
metadata[OVERWRITE][node_id] = {
|
||||
"parameters": overwrite_params,
|
||||
"node_id": node_id,
|
||||
}
|
||||
|
||||
|
||||
# Registry of node-specific extractors
|
||||
# Keys are node class names
|
||||
NODE_EXTRACTORS = {
|
||||
@@ -1221,5 +1311,7 @@ NODE_EXTRACTORS = {
|
||||
"CFGGuider": CFGGuiderExtractor, # Add CFGGuider
|
||||
# Image
|
||||
"VAEDecode": VAEDecodeExtractor, # Added VAEDecode extractor
|
||||
# Metadata overwrite
|
||||
"MetadataOverwriteLM": MetadataOverwriteExtractor,
|
||||
# Add other nodes as needed
|
||||
}
|
||||
|
||||
@@ -0,0 +1,42 @@
|
||||
"""Shared helpers for Metadata Overwrite node metadata collection.
|
||||
|
||||
Used by both the MetadataOverwriteLM node (execution time) and the
|
||||
MetadataOverwriteExtractor (hook time) so the conversion/filtering logic
|
||||
cannot drift between the two paths.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from typing import Any, Dict
|
||||
|
||||
from ..utils.utils import model_patcher_to_name
|
||||
from .constants import CLIP_SKIP_SENTINEL, METADATA_OVERWRITE_FIELDS
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def collect_overwrite_params(values: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Convert node input values into non-default overwrite parameters.
|
||||
|
||||
For most fields, a falsy value (empty string, 0) means "not set" and is
|
||||
skipped. clip_skip uses a dedicated sentinel (-25) so that a wired value
|
||||
of 0 is preserved. The ``model`` field accepts either a manual string or
|
||||
a wired MODEL (ModelPatcher) connection; in the latter case the source
|
||||
model name is extracted from the patcher's ``cached_patcher_init`` and
|
||||
stored as a ComfyUI-style relative path.
|
||||
"""
|
||||
result: Dict[str, Any] = {}
|
||||
for key in METADATA_OVERWRITE_FIELDS:
|
||||
value = values.get(key)
|
||||
if key == "model" and not isinstance(value, str):
|
||||
value = model_patcher_to_name(value)
|
||||
if value is None:
|
||||
logger.warning(
|
||||
"Could not extract model name from wired MODEL input "
|
||||
"(no cached_patcher_init); model metadata overwrite skipped"
|
||||
)
|
||||
if key == "clip_skip":
|
||||
if value != CLIP_SKIP_SENTINEL:
|
||||
result[key] = value
|
||||
elif value:
|
||||
result[key] = value
|
||||
return result
|
||||
@@ -43,7 +43,7 @@ SCANNER_GETTER_NAMES = tuple(SCANNER_TYPE_MAP.keys())
|
||||
|
||||
async def _find_model_entry(
|
||||
model_path: str,
|
||||
) -> tuple[object, object, str | None] | tuple[None, None, None]:
|
||||
) -> tuple[Any, object, str | None] | tuple[None, None, None]:
|
||||
"""Iterate all scanners and return the first (scanner, entry, getter_name)
|
||||
that owns *model_path*. Returns ``(None, None, None)`` when no scanner
|
||||
claims it.
|
||||
@@ -73,7 +73,7 @@ async def _find_model_entry(
|
||||
|
||||
async def _find_scanner_for_model(
|
||||
model_path: str,
|
||||
) -> tuple[object, object] | tuple[None, None]:
|
||||
) -> tuple[Any, object] | tuple[None, None]:
|
||||
"""Find the (scanner, cache_entry) responsible for *model_path*."""
|
||||
scanner, entry, _ = await _find_model_entry(model_path)
|
||||
return scanner, entry
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import logging
|
||||
from typing import List, Tuple
|
||||
import comfy.sd # type: ignore
|
||||
import folder_paths # type: ignore
|
||||
from typing import Any, List, Tuple
|
||||
import comfy.sd # pyright: ignore[reportMissingImports]
|
||||
import folder_paths # pyright: ignore[reportMissingImports]
|
||||
from ..utils.utils import get_checkpoint_info_absolute, _format_model_name_for_comfyui
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -18,9 +18,9 @@ class CheckpointLoaderLM:
|
||||
CATEGORY = "Lora Manager/loaders"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
def INPUT_TYPES(cls):
|
||||
# Get list of checkpoint names from scanner (includes extra folder paths)
|
||||
checkpoint_names = s._get_checkpoint_names()
|
||||
checkpoint_names = cls._get_checkpoint_names()
|
||||
return {
|
||||
"required": {
|
||||
"ckpt_name": (
|
||||
@@ -89,7 +89,7 @@ class CheckpointLoaderLM:
|
||||
logger.error(f"Error getting checkpoint names: {e}")
|
||||
return []
|
||||
|
||||
def load_checkpoint(self, ckpt_name: str) -> Tuple:
|
||||
def load_checkpoint(self, ckpt_name: str) -> Tuple[Any, Any, Any]:
|
||||
"""Load a checkpoint by name, supporting extra folder paths
|
||||
|
||||
Args:
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
"""Create Hook LoRA (LoraManager) — multi-LoRA hook node compatible with ComfyUI's built-in hook pipeline.
|
||||
|
||||
Produces ``("HOOKS",)`` output that chains seamlessly with downstream hook consumers
|
||||
(ConditioningSetProperties, SetHookKeyframes, CombineHooks, SetClipHooks, etc.).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
|
||||
from ..utils.utils import get_lora_info_absolute
|
||||
from .utils import (
|
||||
FlexibleOptionalInputType,
|
||||
any_type,
|
||||
apply_lora_syntax_format,
|
||||
get_loras_list,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CreateHookLoraLM:
|
||||
NAME = "Create Hook LoRA (LoraManager)"
|
||||
CATEGORY = "Lora Manager/hooks"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"text": (
|
||||
"AUTOCOMPLETE_TEXT_LORAS",
|
||||
{
|
||||
"placeholder": "Search LoRAs to add...",
|
||||
"tooltip": (
|
||||
"Search and select LoRAs. Each LoRA gets its own "
|
||||
"model/clip strength. Hooks chain with prev_hooks."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
"optional": FlexibleOptionalInputType(any_type),
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("HOOKS", "STRING", "STRING")
|
||||
RETURN_NAMES = ("HOOKS", "trigger_words", "active_loras")
|
||||
FUNCTION = "create_hook"
|
||||
|
||||
def create_hook(self, text: str, **kwargs):
|
||||
"""Create a HookGroup from the selected LoRAs, chained with prev_hooks.
|
||||
|
||||
Each active LoRA from the widget is loaded and wrapped in a WeightHook
|
||||
via :func:`comfy.hooks.create_hook_lora`. All hooks are combined into a
|
||||
single group and returned alongside trigger words and a human-readable
|
||||
summary of the active LoRAs.
|
||||
"""
|
||||
del text # used by the frontend widget only
|
||||
|
||||
# Lazy imports: comfy is not available in CI/test environment at module level
|
||||
import comfy.hooks # pyright: ignore[reportMissingImports] # noqa: C0415
|
||||
import comfy.utils # pyright: ignore[reportMissingImports] # noqa: C0415
|
||||
|
||||
prev_hooks: comfy.hooks.HookGroup | None = kwargs.get("prev_hooks")
|
||||
|
||||
hook_group = prev_hooks.clone() if prev_hooks is not None else comfy.hooks.HookGroup()
|
||||
|
||||
all_trigger_words: list[str] = []
|
||||
active_loras: list[tuple[str, float, float]] = []
|
||||
|
||||
for lora in get_loras_list(kwargs):
|
||||
if not lora.get("active", False):
|
||||
continue
|
||||
|
||||
lora_name = apply_lora_syntax_format(lora["name"])
|
||||
model_strength = float(lora["strength"])
|
||||
clip_strength = float(lora.get("clipStrength", model_strength))
|
||||
|
||||
# Skip useless no-op entries (both strengths are zero)
|
||||
if model_strength == 0.0 and clip_strength == 0.0:
|
||||
continue
|
||||
|
||||
lora_path, trigger_words = get_lora_info_absolute(lora_name)
|
||||
if not lora_path or not os.path.isfile(lora_path):
|
||||
logger.warning("LoRA '%s' not found — skipping", lora_name)
|
||||
continue
|
||||
|
||||
try:
|
||||
lora_weights = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
||||
|
||||
lora_hooks = comfy.hooks.create_hook_lora(
|
||||
lora=lora_weights,
|
||||
strength_model=model_strength,
|
||||
strength_clip=clip_strength,
|
||||
)
|
||||
except Exception:
|
||||
logger.exception("Failed to load LoRA '%s' — skipping", lora_name)
|
||||
continue
|
||||
hook_group = hook_group.clone_and_combine(lora_hooks)
|
||||
|
||||
active_loras.append((lora_name, model_strength, clip_strength))
|
||||
all_trigger_words.extend(trigger_words)
|
||||
|
||||
# Format trigger words (group mode separator)
|
||||
trigger_words_text = ",, ".join(all_trigger_words) if all_trigger_words else ""
|
||||
|
||||
# Format active LoRAs summary
|
||||
formatted_loras = []
|
||||
for name, model_s, clip_s in active_loras:
|
||||
if abs(model_s - clip_s) > 0.001:
|
||||
formatted_loras.append(
|
||||
f"<lora:{name}:{model_s}:{clip_s}>"
|
||||
)
|
||||
else:
|
||||
formatted_loras.append(f"<lora:{name}:{model_s}>")
|
||||
active_loras_text = " ".join(formatted_loras)
|
||||
|
||||
return (hook_group, trigger_words_text, active_loras_text)
|
||||
@@ -1,8 +1,8 @@
|
||||
import importlib
|
||||
import logging
|
||||
|
||||
import comfy.sd # type: ignore
|
||||
import comfy.utils # type: ignore
|
||||
import comfy.sd # pyright: ignore[reportMissingImports]
|
||||
import comfy.utils # pyright: ignore[reportMissingImports]
|
||||
|
||||
from ..utils.utils import get_lora_info_absolute
|
||||
from .utils import (
|
||||
|
||||
@@ -1,26 +1,102 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
_STACK_INPUT_PATTERN = re.compile(r"^lora_stack(?:_([ab])|(\d+))$")
|
||||
|
||||
|
||||
def _is_stack_input(name: str) -> bool:
|
||||
return bool(_STACK_INPUT_PATTERN.match(name))
|
||||
|
||||
|
||||
def _stack_slot_number(name: str) -> int:
|
||||
"""Numeric slot used to order stack inputs; legacy a/b map to 1/2."""
|
||||
match = _STACK_INPUT_PATTERN.match(name)
|
||||
if not match:
|
||||
return -1
|
||||
letter, digits = match.group(1), match.group(2)
|
||||
if digits is not None:
|
||||
return int(digits)
|
||||
return 1 if letter == "a" else 2
|
||||
|
||||
|
||||
class _LoraStackOptionalInputs:
|
||||
"""Lookup that preserves explicit optional inputs and dynamic lora_stack slots."""
|
||||
|
||||
def __init__(self, explicit_inputs: dict[str, tuple[str, dict[str, Any]]]) -> None:
|
||||
self._explicit_inputs = explicit_inputs
|
||||
|
||||
def __contains__(self, item: object) -> bool:
|
||||
if not isinstance(item, str):
|
||||
return False
|
||||
return item in self._explicit_inputs or _is_stack_input(item)
|
||||
|
||||
def __getitem__(self, key: str) -> tuple[str, dict[str, Any]]:
|
||||
if key in self._explicit_inputs:
|
||||
return self._explicit_inputs[key]
|
||||
if _is_stack_input(key):
|
||||
return (
|
||||
"LORA_STACK",
|
||||
{
|
||||
"tooltip": "A LoRA stack to combine. Connect to add more inputs.",
|
||||
},
|
||||
)
|
||||
raise KeyError(key)
|
||||
|
||||
|
||||
class LoraStackCombinerLM:
|
||||
NAME = "Lora Stack Combiner (LoraManager)"
|
||||
CATEGORY = "Lora Manager/stackers"
|
||||
DESCRIPTION = (
|
||||
"Combines multiple LoRA stacks into a single stack. "
|
||||
"Supports dynamic inputs: connect a stack to add more inputs."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
optional_inputs: dict[str, tuple[str, dict[str, Any]]] = {
|
||||
"lora_stack1": (
|
||||
"LORA_STACK",
|
||||
{
|
||||
"tooltip": "A LoRA stack to combine. Connect to add more inputs.",
|
||||
},
|
||||
),
|
||||
"lora_stack2": (
|
||||
"LORA_STACK",
|
||||
{
|
||||
"tooltip": "A LoRA stack to combine. Connect to add more inputs.",
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
stack = inspect.stack()
|
||||
if len(stack) > 2 and stack[2].function == "get_input_info":
|
||||
optional_inputs = _LoraStackOptionalInputs(optional_inputs) # pyright: ignore[reportAssignmentType]
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"lora_stack_a": ("LORA_STACK",),
|
||||
"lora_stack_b": ("LORA_STACK",),
|
||||
},
|
||||
"required": {},
|
||||
"optional": optional_inputs,
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("LORA_STACK",)
|
||||
RETURN_NAMES = ("LORA_STACK",)
|
||||
FUNCTION = "combine_stacks"
|
||||
|
||||
def combine_stacks(self, lora_stack_a, lora_stack_b):
|
||||
combined_stack = []
|
||||
def combine_stacks(self, lora_stack1=None, lora_stack2=None, **kwargs):
|
||||
stacks = {
|
||||
"lora_stack1": lora_stack1,
|
||||
"lora_stack2": lora_stack2,
|
||||
}
|
||||
for key, value in kwargs.items():
|
||||
if _is_stack_input(key) and value is not None:
|
||||
stacks[key] = value
|
||||
|
||||
if lora_stack_a:
|
||||
combined_stack.extend(lora_stack_a)
|
||||
if lora_stack_b:
|
||||
combined_stack.extend(lora_stack_b)
|
||||
combined_stack = []
|
||||
for key in sorted(stacks, key=_stack_slot_number):
|
||||
stack = stacks[key]
|
||||
if stack:
|
||||
combined_stack.extend(stack)
|
||||
|
||||
return (combined_stack,)
|
||||
|
||||
@@ -0,0 +1,169 @@
|
||||
"""Metadata Overwrite node — allows users to manually specify generation parameters
|
||||
that override the automatically collected/inferred metadata.
|
||||
|
||||
Most inputs have falsy defaults (empty string / 0) which are skipped.
|
||||
clip_skip uses a sentinel default (-25) so that a wired value of 0 is
|
||||
preserved — both ComfyUI and A1111 conventions have no meaningful 0 value,
|
||||
but users may wire 0 to express "no clip skip / default".
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from ..metadata_collector.constants import CLIP_SKIP_SENTINEL as _CLIP_SKIP_SENTINEL
|
||||
from ..metadata_collector.overwrite_utils import collect_overwrite_params
|
||||
|
||||
|
||||
class MetadataOverwriteLM:
|
||||
NAME = "Metadata Overwrite (LoraManager)"
|
||||
CATEGORY = "Lora Manager/utils"
|
||||
DESCRIPTION = (
|
||||
"Manually specify generation parameters to override automatically collected "
|
||||
"metadata. Only filled/connected inputs will take effect — empty defaults "
|
||||
"are ignored."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> dict[str, Any]:
|
||||
return {
|
||||
"optional": {
|
||||
"prompt": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": "Positive prompt. Only overwrites when non-empty.",
|
||||
},
|
||||
),
|
||||
"negative_prompt": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": "Negative prompt. Only overwrites when non-empty.",
|
||||
},
|
||||
),
|
||||
"seed": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 0xFFFFFFFFFFFFFFFF,
|
||||
"control_after_generate": False,
|
||||
"tooltip": "Seed value. Only overwrites when > 0.",
|
||||
},
|
||||
),
|
||||
"steps": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 10000,
|
||||
"tooltip": "Number of steps. Only overwrites when > 0.",
|
||||
},
|
||||
),
|
||||
"cfg_scale": (
|
||||
"FLOAT",
|
||||
{
|
||||
"default": 0.0,
|
||||
"min": 0.0,
|
||||
"max": 100.0,
|
||||
"tooltip": "CFG scale. Only overwrites when > 0.",
|
||||
},
|
||||
),
|
||||
"sampler": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Sampler name. Only overwrites when non-empty.",
|
||||
},
|
||||
),
|
||||
"scheduler": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": "Scheduler name. Only overwrites when non-empty.",
|
||||
},
|
||||
),
|
||||
"model": (
|
||||
"STRING,MODEL",
|
||||
{
|
||||
"default": "",
|
||||
"widgetType": "STRING",
|
||||
"tooltip": (
|
||||
"The checkpoint or diffusion model (UNet) used "
|
||||
"for generation. Fill in the name manually or "
|
||||
"connect a MODEL output — the model name is then "
|
||||
"extracted automatically. Only overwrites when "
|
||||
"non-empty."
|
||||
),
|
||||
},
|
||||
),
|
||||
"loras": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": (
|
||||
"LoRA syntax, e.g. <lora:name:strength> "
|
||||
"or <lora:name:model_strength:clip_strength>, "
|
||||
"separated by spaces. Only overwrites when non-empty."
|
||||
),
|
||||
},
|
||||
),
|
||||
"size": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"tooltip": (
|
||||
"Image size in WIDTHxHEIGHT format (e.g. 512x768). "
|
||||
"Only overwrites when non-empty."
|
||||
),
|
||||
},
|
||||
),
|
||||
"clip_skip": (
|
||||
"INT",
|
||||
{
|
||||
"default": _CLIP_SKIP_SENTINEL,
|
||||
"min": -25,
|
||||
"max": 24,
|
||||
"tooltip": (
|
||||
"Clip skip (ComfyUI: -24..-1, A1111: 1+). "
|
||||
"Default -25 means not set — any other value "
|
||||
"overwrites."
|
||||
),
|
||||
},
|
||||
),
|
||||
"additional_data": (
|
||||
"STRING",
|
||||
{
|
||||
"default": "",
|
||||
"multiline": True,
|
||||
"tooltip": (
|
||||
"Additional data to embed in the image metadata. "
|
||||
"Inserted between Clip skip and Model hash in the "
|
||||
"A1111-compatible parameters string. "
|
||||
'Example: "Copyright": "Some license info"'
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("METADATA",)
|
||||
RETURN_NAMES = ("metadata",)
|
||||
FUNCTION = "collect_metadata"
|
||||
OUTPUT_NODE = True
|
||||
|
||||
def collect_metadata(self, **kwargs: Any) -> tuple[dict[str, Any]]:
|
||||
"""Collect non-default input values into a metadata dict.
|
||||
|
||||
For most fields, a falsy value (empty string, 0) means "not set"
|
||||
and is skipped. clip_skip uses a dedicated sentinel (-25) so that
|
||||
a wired value of 0 is preserved and reaches the metadata pipeline.
|
||||
|
||||
The ``model`` field accepts either a manual string or a wired MODEL
|
||||
(ModelPatcher) connection; in the latter case the underlying model
|
||||
name is extracted from the patcher's ``cached_patcher_init`` and
|
||||
stored as a ComfyUI-style relative path.
|
||||
"""
|
||||
return (collect_overwrite_params(kwargs),)
|
||||
+12
-13
@@ -15,15 +15,15 @@ import os
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple, Union
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union, cast
|
||||
|
||||
import comfy.utils # type: ignore
|
||||
import folder_paths # type: ignore
|
||||
import comfy.utils # pyright: ignore[reportMissingImports]
|
||||
import folder_paths # pyright: ignore[reportMissingImports]
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from safetensors import safe_open
|
||||
|
||||
from nunchaku.lora.flux.nunchaku_converter import (
|
||||
from nunchaku.lora.flux.nunchaku_converter import ( # pyright: ignore[reportMissingTypeStubs]
|
||||
pack_lowrank_weight,
|
||||
unpack_lowrank_weight,
|
||||
)
|
||||
@@ -87,10 +87,6 @@ def _rename_layer_underscore_layer_name(old_name: str) -> str:
|
||||
return new_name
|
||||
|
||||
|
||||
def _is_indexable_module(module):
|
||||
return isinstance(module, (nn.ModuleList, nn.Sequential, list, tuple))
|
||||
|
||||
|
||||
def _get_module_by_name(model: nn.Module, name: str) -> Optional[nn.Module]:
|
||||
if not name:
|
||||
return model
|
||||
@@ -100,7 +96,7 @@ def _get_module_by_name(model: nn.Module, name: str) -> Optional[nn.Module]:
|
||||
continue
|
||||
if hasattr(module, part):
|
||||
module = getattr(module, part)
|
||||
elif part.isdigit() and _is_indexable_module(module):
|
||||
elif part.isdigit() and isinstance(module, (nn.ModuleList, nn.Sequential, list, tuple)):
|
||||
try:
|
||||
module = module[int(part)]
|
||||
except (IndexError, TypeError):
|
||||
@@ -267,7 +263,9 @@ def _handle_proj_out_split(lora_dict: Dict[str, Dict[str, torch.Tensor]], base_k
|
||||
return result, consumed
|
||||
|
||||
|
||||
def _apply_lora_to_module(module: nn.Module, a_tensor: torch.Tensor, b_tensor: torch.Tensor, module_name: str, model: nn.Module) -> None:
|
||||
def _apply_lora_to_module(module: Any, a_tensor: torch.Tensor, b_tensor: torch.Tensor, module_name: str, model: Any) -> None:
|
||||
# These modules are dynamic torch containers; monkey-patched attributes
|
||||
# below are set at runtime, so the module/model types are deliberately Any.
|
||||
if not hasattr(module, "in_features") or not hasattr(module, "out_features"):
|
||||
raise ValueError(f"{module_name}: unsupported module without in/out features")
|
||||
if a_tensor.shape[1] != module.in_features or b_tensor.shape[0] != module.out_features:
|
||||
@@ -336,7 +334,7 @@ def _apply_lora_to_module(module: nn.Module, a_tensor: torch.Tensor, b_tensor: t
|
||||
raise ValueError(f"{module_name}: unsupported module type {type(module)}")
|
||||
|
||||
|
||||
def reset_lora_v2(model: nn.Module) -> None:
|
||||
def reset_lora_v2(model: Any) -> None:
|
||||
slots = getattr(model, "_lora_slots", None)
|
||||
if not slots:
|
||||
return
|
||||
@@ -344,6 +342,7 @@ def reset_lora_v2(model: nn.Module) -> None:
|
||||
module = _get_module_by_name(model, name)
|
||||
if module is None:
|
||||
continue
|
||||
module = cast(Any, module)
|
||||
module_type = info.get("type", "nunchaku")
|
||||
if module_type == "nunchaku":
|
||||
base_rank = info["base_rank"]
|
||||
@@ -371,7 +370,7 @@ def reset_lora_v2(model: nn.Module) -> None:
|
||||
def compose_loras_v2(model: nn.Module, lora_configs: List[Tuple[Union[str, Path, Dict[str, torch.Tensor]], float]], apply_awq_mod: bool = True) -> bool:
|
||||
del apply_awq_mod # retained for interface compatibility
|
||||
reset_lora_v2(model)
|
||||
aggregated_weights: Dict[str, List[Dict[str, object]]] = defaultdict(list)
|
||||
aggregated_weights: Dict[str, List[Dict[str, Any]]] = defaultdict(list)
|
||||
saw_supported_format = False
|
||||
unresolved_targets = 0
|
||||
|
||||
@@ -471,7 +470,7 @@ def compose_loras_v2(model: nn.Module, lora_configs: List[Tuple[Union[str, Path,
|
||||
class ComfyQwenImageWrapperLM(nn.Module):
|
||||
def __init__(self, model: nn.Module, config=None, apply_awq_mod: bool = True):
|
||||
super().__init__()
|
||||
self.model = model
|
||||
self.model: Any = model
|
||||
self.config = {} if config is None else config
|
||||
self.dtype = next(model.parameters()).dtype
|
||||
self.loras: List[Tuple[Union[str, Path, Dict[str, torch.Tensor]], float]] = []
|
||||
|
||||
+2
-2
@@ -67,7 +67,7 @@ class PromptLM:
|
||||
|
||||
stack = inspect.stack()
|
||||
if len(stack) > 2 and stack[2].function == "get_input_info":
|
||||
optional_inputs = _PromptOptionalInputs(optional_inputs) # type: ignore[assignment]
|
||||
optional_inputs = _PromptOptionalInputs(optional_inputs) # pyright: ignore[reportAssignmentType]
|
||||
|
||||
return {
|
||||
"required": {
|
||||
@@ -126,7 +126,7 @@ class PromptLM:
|
||||
else:
|
||||
prompt = expanded_text
|
||||
|
||||
from nodes import CLIPTextEncode # type: ignore
|
||||
from nodes import CLIPTextEncode # pyright: ignore[reportMissingImports, reportAttributeAccessIssue]
|
||||
|
||||
conditioning = CLIPTextEncode().encode(clip, prompt)[0]
|
||||
return (conditioning, prompt)
|
||||
|
||||
+362
-130
@@ -5,7 +5,7 @@ import time
|
||||
import uuid
|
||||
from typing import Any, Dict, Optional
|
||||
import numpy as np
|
||||
import folder_paths # type: ignore
|
||||
import folder_paths # pyright: ignore[reportMissingImports]
|
||||
from ..services.service_registry import ServiceRegistry
|
||||
from ..metadata_collector.metadata_processor import MetadataProcessor
|
||||
from ..metadata_collector import get_metadata
|
||||
@@ -13,9 +13,159 @@ from ..utils.constants import CARD_PREVIEW_WIDTH
|
||||
from ..utils.exif_utils import ExifUtils
|
||||
from ..utils.utils import calculate_recipe_fingerprint, sanitize_folder_name
|
||||
from PIL import Image, PngImagePlugin
|
||||
import piexif
|
||||
import piexif # pyright: ignore[reportMissingTypeStubs]
|
||||
import logging
|
||||
|
||||
# Civitai-compatible sampler name mapping: ComfyUI internal → A1111 display name
|
||||
CIVITAI_SAMPLER_MAP = {
|
||||
"euler": "Euler",
|
||||
"euler_ancestral": "Euler a",
|
||||
"lms": "LMS",
|
||||
"heun": "Heun",
|
||||
"dpm_2": "DPM2",
|
||||
"dpm_2_ancestral": "DPM2 a",
|
||||
"dpmpp_2s_ancestral": "DPM++ 2S a",
|
||||
"dpmpp_2m": "DPM++ 2M",
|
||||
"dpmpp_sde": "DPM++ SDE",
|
||||
"dpmpp_sde_gpu": "DPM++ SDE",
|
||||
"dpmpp_2m_sde": "DPM++ 2M SDE",
|
||||
"dpmpp_2m_sde_gpu": "DPM++ 2M SDE",
|
||||
"dpmpp_3m_sde": "DPM++ 3M SDE",
|
||||
"dpm_fast": "DPM fast",
|
||||
"dpm_adaptive": "DPM adaptive",
|
||||
"ddim": "DDIM",
|
||||
"plms": "PLMS",
|
||||
"uni_pc_bh2": "UniPC",
|
||||
"uni_pc": "UniPC",
|
||||
"lcm": "LCM",
|
||||
}
|
||||
|
||||
# Base model display name → AIR URN slug
|
||||
# Sourced from civitai source: src/shared/constants/basemodel.constants.ts
|
||||
BASE_MODEL_AIR_SLUG = {
|
||||
# Stable Diffusion family
|
||||
"SD 1.4": "sd1",
|
||||
"SD 1.5": "sd1",
|
||||
"SD 1.5 LCM": "sd1",
|
||||
"SD 1.5 Hyper": "sd1",
|
||||
"SD 2.0": "sd2",
|
||||
"SD 2.0 768": "sd2",
|
||||
"SD 2.1": "sd2",
|
||||
"SD 2.1 768": "sd2",
|
||||
"SD 2.1 Unclip": "sd2",
|
||||
"SD 3.0": "sd3",
|
||||
"SD 3.5": "sd35",
|
||||
"SD 3.5 Large": "sd35",
|
||||
"SD 3.5 Large Turbo": "sd35",
|
||||
"SD 3.5 Medium": "sd35",
|
||||
"SDXL 0.9": "sdxl",
|
||||
"SDXL 1.0": "sdxl",
|
||||
"SDXL 1.0 LCM": "sdxl",
|
||||
"SDXL Lightning": "sdxl",
|
||||
"SDXL Hyper": "sdxl",
|
||||
"SDXL Turbo": "sdxl",
|
||||
"SDXL Distilled": "sdxldistilled",
|
||||
"Stable Cascade": "scascade",
|
||||
"Stable Video Diffusion": "svd",
|
||||
"SVD": "svd",
|
||||
"SVD XT": "svdxt",
|
||||
|
||||
# SDXL community fine-tunes
|
||||
"Pony": "pony",
|
||||
"Pony Diffusion": "pony",
|
||||
"Illustrious": "illustrious",
|
||||
"NoobAI": "noobai",
|
||||
"Animagine": "illustrious",
|
||||
|
||||
# Flux family
|
||||
"Flux.1": "flux1",
|
||||
"Flux.1 D": "flux1",
|
||||
"Flux.1 S": "flux1",
|
||||
"Flux.1 Krea": "fluxkrea",
|
||||
"Flux.1 Kontext": "flux1kontext",
|
||||
"Flux.2": "flux2",
|
||||
"Flux.2 D": "flux2",
|
||||
"Flux.2 Klein 9B": "flux2klein_9b",
|
||||
"Flux.2 Klein 9B Base": "flux2klein_9b_base",
|
||||
"Flux.2 Klein 4B": "flux2klein_4b",
|
||||
"Flux.2 Klein 4B Base": "flux2klein_4b_base",
|
||||
|
||||
# Other image models (sorted alphabetically)
|
||||
"AuraFlow": "auraflow",
|
||||
"Chroma": "chroma",
|
||||
"HiDream": "hidream",
|
||||
"HiDream-O1": "hidream-o1",
|
||||
"Hunyuan DiT": "hydit1",
|
||||
"Hunyuan Video": "hyv1",
|
||||
"Kolors": "kolors",
|
||||
"Lumina": "lumina",
|
||||
"Mochi": "mochi",
|
||||
"ODOR": "odor",
|
||||
"PixArt Alpha": "pixarta",
|
||||
"PixArt Sigma": "pixarte",
|
||||
"Playground v2": "playgroundv2",
|
||||
"Playground v2.5": "playgroundv2",
|
||||
"Pony Diffusion V7": "ponyv7",
|
||||
|
||||
# Video models
|
||||
"CogVideoX": "cogvideox",
|
||||
"LTX Video": "ltxv",
|
||||
"LTX Video 2": "ltxv2",
|
||||
"LTX Video 2.3": "ltxv23",
|
||||
"Wan Video": "wanvideo",
|
||||
"Wan Video 1.3B T2V": "wanvideo_13b_t2v",
|
||||
"Wan Video 14B T2V": "wanvideo_14b_t2v",
|
||||
"Wan Video 14B I2V 480p": "wanvideo_14b_i2v_480p",
|
||||
"Wan Video 14B I2V 720p": "wanvideo_14b_i2v_720p",
|
||||
|
||||
# Third-party / proprietary image models
|
||||
"Boogu": "boogu",
|
||||
"Ernie": "ernie",
|
||||
"Grok": "grok",
|
||||
"HappyHorse": "happyhorse",
|
||||
"Ideogram": "ideogram",
|
||||
"Ideogram 4.0": "ideogram",
|
||||
"Imagen": "imagen4",
|
||||
"Imagen 4": "imagen4",
|
||||
"Krea": "krea2",
|
||||
"Krea 2": "krea2",
|
||||
"Lens": "lens",
|
||||
"MAI": "mai",
|
||||
"Nano Banana": "nanobanana",
|
||||
"OpenAI": "openai",
|
||||
"Reve": "reve",
|
||||
"Reve 2": "reve",
|
||||
"Reve 2.1": "reve",
|
||||
"Seedream": "seedream",
|
||||
"Sora": "sora2",
|
||||
"Sora 2": "sora2",
|
||||
"Veo": "veo3",
|
||||
"Veo 2": "veo3",
|
||||
"Veo 3": "veo3",
|
||||
"ZImageTurbo": "zimageturbo",
|
||||
"ZImageBase": "zimagebase",
|
||||
"ZImage": "zimagebase",
|
||||
|
||||
# Third-party video models
|
||||
"Hailuo by MiniMax": "minimax",
|
||||
"Haiper": "haiper",
|
||||
"Kling": "kling",
|
||||
"Lightricks": "lightricks",
|
||||
"Seedance": "seedance",
|
||||
"Vidu": "vidu",
|
||||
|
||||
# Qwen family
|
||||
"Qwen": "qwen",
|
||||
"Qwen 2": "qwen2",
|
||||
|
||||
# Anima
|
||||
"Anima": "anima",
|
||||
|
||||
# Special
|
||||
"Upscaler": "upscaler",
|
||||
"Other": "other",
|
||||
}
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@@ -70,11 +220,29 @@ class SaveImageLM:
|
||||
"tooltip": "Compression quality for JPEG and lossy WebP formats (1-100). Higher values mean better quality but larger files.",
|
||||
},
|
||||
),
|
||||
"webp_method": (
|
||||
"INT",
|
||||
{
|
||||
"default": 6,
|
||||
"min": 0,
|
||||
"max": 6,
|
||||
"tooltip": "WebP compression method (0-6). 0=fastest/largest, 6=slowest/smallest. Only applies when file_format is 'webp'.",
|
||||
},
|
||||
),
|
||||
"jpeg_subsampling": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 2,
|
||||
"tooltip": "JPEG chroma subsampling level. 0=4:4:4 (best quality), 1=4:2:2, 2=4:2:0 (smallest files). Only applies when file_format is 'jpeg'.",
|
||||
},
|
||||
),
|
||||
"embed_workflow": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": "Embeds the complete workflow data into the image metadata. Only works with PNG and WebP formats.",
|
||||
"tooltip": "When enabled, saved images store the complete workflow. Drag the image back into ComfyUI to restore the original node graph. PNG and WebP only.",
|
||||
},
|
||||
),
|
||||
"save_with_metadata": (
|
||||
@@ -84,6 +252,13 @@ class SaveImageLM:
|
||||
"tooltip": "When enabled, embeds generation parameters into the saved image metadata. Disable to skip writing generation metadata.",
|
||||
},
|
||||
),
|
||||
"add_loras_to_prompt": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
"default": False,
|
||||
"tooltip": "When enabled, appends the LoRA syntax line (e.g. <lora:name:strength>) after the positive prompt in the saved metadata.",
|
||||
},
|
||||
),
|
||||
"add_counter_to_filename": (
|
||||
"BOOLEAN",
|
||||
{
|
||||
@@ -142,148 +317,197 @@ class SaveImageLM:
|
||||
|
||||
return None
|
||||
|
||||
def format_metadata(self, metadata_dict):
|
||||
"""Format metadata in the requested format similar to userComment example"""
|
||||
if not metadata_dict:
|
||||
return ""
|
||||
def _resolve_model_cache_entry(self, scanner_type: str, name: str):
|
||||
"""Resolve model hash, civitai metadata, and base_model from scanner cache.
|
||||
Returns (hash_str, civitai_dict, base_model_str). All values are empty defaults when not found."""
|
||||
scanner = ServiceRegistry.get_service_sync(scanner_type)
|
||||
if scanner is None or not name:
|
||||
return "", {}, ""
|
||||
|
||||
# Helper function to only add parameter if value is not None
|
||||
def add_param_if_not_none(param_list, label, value):
|
||||
if value is not None:
|
||||
param_list.append(f"{label}: {value}")
|
||||
entry = self._get_cached_model_by_name(scanner, name)
|
||||
if entry is None:
|
||||
basename = os.path.splitext(os.path.basename(name))[0]
|
||||
hash_val = scanner.get_hash_by_filename(basename)
|
||||
return (hash_val or "").lower(), {}, ""
|
||||
|
||||
hash_val = (entry.get("sha256") or "").lower()
|
||||
civitai = entry.get("civitai") or {}
|
||||
base_model = entry.get("base_model") or ""
|
||||
return hash_val, civitai, base_model
|
||||
|
||||
@staticmethod
|
||||
def _get_civitai_sampler_name(sampler_name: str, scheduler: str) -> str:
|
||||
if sampler_name in CIVITAI_SAMPLER_MAP:
|
||||
civitai_name = CIVITAI_SAMPLER_MAP[sampler_name]
|
||||
if scheduler == "karras":
|
||||
civitai_name += " Karras"
|
||||
elif scheduler == "exponential":
|
||||
civitai_name += " Exponential"
|
||||
return civitai_name
|
||||
else:
|
||||
if scheduler and scheduler != "normal":
|
||||
return f"{sampler_name}_{scheduler}"
|
||||
return sampler_name
|
||||
|
||||
@staticmethod
|
||||
def _build_air_string(base_model: str, model_type: str, model_id: int, version_id: int) -> str:
|
||||
slug = BASE_MODEL_AIR_SLUG.get(base_model, "other")
|
||||
type_lower = model_type.lower() if model_type else "other"
|
||||
return f"urn:air:{slug}:{type_lower}:civitai:{model_id}@{version_id}"
|
||||
|
||||
def format_metadata(self, metadata_dict: dict[str, Any], add_loras_to_prompt: bool = False) -> str:
|
||||
"""Format metadata as A1111-compatible parameters string with Hashes JSON and Civitai resources."""
|
||||
if not metadata_dict: return ""
|
||||
|
||||
# Extract the prompt and negative prompt
|
||||
prompt = metadata_dict.get("prompt", "")
|
||||
negative_prompt = metadata_dict.get("negative_prompt", "")
|
||||
|
||||
# Extract loras from the prompt if present
|
||||
steps = metadata_dict.get("steps")
|
||||
cfg = metadata_dict.get("guidance")
|
||||
if cfg is None:
|
||||
cfg = metadata_dict.get("cfg_scale")
|
||||
if cfg is None:
|
||||
cfg = metadata_dict.get("cfg")
|
||||
seed = metadata_dict.get("seed")
|
||||
size = metadata_dict.get("size")
|
||||
sampler = metadata_dict.get("sampler") or ""
|
||||
scheduler = metadata_dict.get("scheduler") or "normal"
|
||||
checkpoint = metadata_dict.get("checkpoint") or ""
|
||||
loras_text = metadata_dict.get("loras", "")
|
||||
lora_hashes = {}
|
||||
clip_skip = metadata_dict.get("clip_skip")
|
||||
|
||||
# If loras are found, add them on a new line after the prompt
|
||||
# Parse LoRA entries from <lora:name:strength> format
|
||||
lora_entries: list[tuple[str, float]] = []
|
||||
if loras_text:
|
||||
prompt_with_loras = f"{prompt}\n{loras_text}"
|
||||
for match in re.findall(r"<lora:([^:]+):([^>]+)>", loras_text):
|
||||
lora_name, strength_str = match
|
||||
try:
|
||||
strength = float(strength_str)
|
||||
except (ValueError, TypeError):
|
||||
strength = 1.0
|
||||
lora_entries.append((lora_name, strength))
|
||||
|
||||
# Extract lora names from the format <lora:name:strength>
|
||||
lora_matches = re.findall(r"<lora:([^:]+):([^>]+)>", loras_text)
|
||||
# Resolve checkpoint hash and Civitai data from local cache
|
||||
ckpt_hash, ckpt_civitai, ckpt_base_model = "", {}, ""
|
||||
ckpt_display_name = ""
|
||||
if checkpoint:
|
||||
ckpt_hash, ckpt_civitai, ckpt_base_model = self._resolve_model_cache_entry(
|
||||
"checkpoint_scanner", checkpoint
|
||||
)
|
||||
ckpt_display_name = os.path.splitext(os.path.basename(checkpoint))[0]
|
||||
|
||||
# Get hash for each lora
|
||||
for lora_name, strength in lora_matches:
|
||||
hash_value = self.get_lora_hash(lora_name)
|
||||
if hash_value:
|
||||
lora_hashes[lora_name] = hash_value
|
||||
else:
|
||||
prompt_with_loras = prompt
|
||||
# Resolve LoRA hash and Civitai data from local cache
|
||||
loras_data: list[dict[str, Any]] = []
|
||||
for lora_name, strength in lora_entries:
|
||||
lora_hash, lora_civitai, lora_base_model = self._resolve_model_cache_entry(
|
||||
"lora_scanner", lora_name
|
||||
)
|
||||
loras_data.append({
|
||||
"name": lora_name,
|
||||
"strength": strength,
|
||||
"hash": lora_hash,
|
||||
"civitai": lora_civitai,
|
||||
"base_model": lora_base_model,
|
||||
})
|
||||
|
||||
# Format the first part (prompt and loras)
|
||||
metadata_parts = [prompt_with_loras]
|
||||
# Build Hashes JSON (A1111 / Civitai standard format)
|
||||
hashes: dict[str, str] = {}
|
||||
if ckpt_hash:
|
||||
hashes["model"] = ckpt_hash[:10].upper()
|
||||
for lora in loras_data:
|
||||
if lora["hash"]:
|
||||
hashes[f"LORA:{lora['name']}"] = lora["hash"][:10].upper()
|
||||
|
||||
# Add negative prompt
|
||||
# Build Civitai resources JSON array
|
||||
civitai_resources: list[dict[str, Any]] = []
|
||||
if ckpt_civitai.get("id", 0) > 0:
|
||||
ckpt_resource: dict[str, Any] = {}
|
||||
ckpt_type = (ckpt_civitai.get("model") or {}).get("type", "Checkpoint")
|
||||
model_id = ckpt_civitai.get("modelId", 0)
|
||||
version_id = ckpt_civitai.get("id", 0)
|
||||
if model_id and version_id:
|
||||
ckpt_resource["air"] = self._build_air_string(
|
||||
ckpt_base_model, ckpt_type, int(model_id), int(version_id)
|
||||
)
|
||||
elif version_id:
|
||||
ckpt_resource["modelVersionId"] = int(version_id)
|
||||
if ckpt_civitai.get("name"):
|
||||
ckpt_resource["versionName"] = ckpt_civitai["name"]
|
||||
if ckpt_resource:
|
||||
civitai_resources.append(ckpt_resource)
|
||||
|
||||
for lora in loras_data:
|
||||
lora_civitai = lora["civitai"]
|
||||
if not lora_civitai or lora_civitai.get("id", 0) <= 0:
|
||||
continue
|
||||
lora_resource: dict[str, Any] = {"weight": lora["strength"]}
|
||||
lora_type = (lora_civitai.get("model") or {}).get("type", "LORA")
|
||||
model_id = lora_civitai.get("modelId", 0)
|
||||
version_id = lora_civitai.get("id", 0)
|
||||
if model_id and version_id:
|
||||
lora_resource["air"] = self._build_air_string(
|
||||
lora["base_model"], lora_type, int(model_id), int(version_id)
|
||||
)
|
||||
elif version_id:
|
||||
lora_resource["modelVersionId"] = int(version_id)
|
||||
if lora_civitai.get("name"):
|
||||
lora_resource["versionName"] = lora_civitai["name"]
|
||||
civitai_resources.append(lora_resource)
|
||||
|
||||
sampler_name = CIVITAI_SAMPLER_MAP.get(sampler, sampler) if sampler else None
|
||||
|
||||
scheduler_mapping = {
|
||||
"normal": "Normal",
|
||||
"karras": "Karras",
|
||||
"exponential": "Exponential",
|
||||
"sgm_uniform": "SGM Uniform",
|
||||
"sgm_quadratic": "SGM Quadratic",
|
||||
}
|
||||
scheduler_name = scheduler_mapping.get(scheduler, scheduler) if scheduler else None
|
||||
|
||||
# Build output lines
|
||||
prompt_line = prompt if prompt else ""
|
||||
if add_loras_to_prompt and loras_text:
|
||||
prompt_line = f"{prompt_line}\n{loras_text}" if prompt_line else loras_text
|
||||
lines = [prompt_line] if prompt_line else [""]
|
||||
if negative_prompt:
|
||||
metadata_parts.append(f"Negative prompt: {negative_prompt}")
|
||||
lines.append(f"Negative prompt: {negative_prompt}")
|
||||
|
||||
# Format the second part (generation parameters)
|
||||
params = []
|
||||
|
||||
# Add standard parameters in the correct order
|
||||
if "steps" in metadata_dict:
|
||||
add_param_if_not_none(params, "Steps", metadata_dict.get("steps"))
|
||||
|
||||
# Combine sampler and scheduler information
|
||||
sampler_name = None
|
||||
scheduler_name = None
|
||||
|
||||
if "sampler" in metadata_dict:
|
||||
sampler = metadata_dict.get("sampler")
|
||||
# Convert ComfyUI sampler names to user-friendly names
|
||||
sampler_mapping = {
|
||||
"euler": "Euler",
|
||||
"euler_ancestral": "Euler a",
|
||||
"dpm_2": "DPM2",
|
||||
"dpm_2_ancestral": "DPM2 a",
|
||||
"heun": "Heun",
|
||||
"dpm_fast": "DPM fast",
|
||||
"dpm_adaptive": "DPM adaptive",
|
||||
"lms": "LMS",
|
||||
"dpmpp_2s_ancestral": "DPM++ 2S a",
|
||||
"dpmpp_sde": "DPM++ SDE",
|
||||
"dpmpp_sde_gpu": "DPM++ SDE",
|
||||
"dpmpp_2m": "DPM++ 2M",
|
||||
"dpmpp_2m_sde": "DPM++ 2M SDE",
|
||||
"dpmpp_2m_sde_gpu": "DPM++ 2M SDE",
|
||||
"ddim": "DDIM",
|
||||
}
|
||||
sampler_name = sampler_mapping.get(sampler, sampler)
|
||||
|
||||
if "scheduler" in metadata_dict:
|
||||
scheduler = metadata_dict.get("scheduler")
|
||||
scheduler_mapping = {
|
||||
"normal": "Simple",
|
||||
"karras": "Karras",
|
||||
"exponential": "Exponential",
|
||||
"sgm_uniform": "SGM Uniform",
|
||||
"sgm_quadratic": "SGM Quadratic",
|
||||
}
|
||||
scheduler_name = scheduler_mapping.get(scheduler, scheduler)
|
||||
|
||||
# Add combined sampler and scheduler information
|
||||
params: list[str] = []
|
||||
if steps is not None:
|
||||
params.append(f"Steps: {steps}")
|
||||
if sampler_name:
|
||||
if scheduler_name:
|
||||
params.append(f"Sampler: {sampler_name} {scheduler_name}")
|
||||
else:
|
||||
params.append(f"Sampler: {sampler_name}")
|
||||
if cfg is not None:
|
||||
params.append(f"CFG scale: {cfg}")
|
||||
if seed is not None:
|
||||
params.append(f"Seed: {seed}")
|
||||
if size:
|
||||
params.append(f"Size: {size}")
|
||||
if clip_skip is not None:
|
||||
try:
|
||||
params.append(f"Clip skip: {abs(int(clip_skip))}")
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
additional_data = metadata_dict.get("additional_data", "")
|
||||
if additional_data:
|
||||
params.append(additional_data)
|
||||
if ckpt_hash:
|
||||
params.append(f"Model hash: {ckpt_hash[:10].upper()}")
|
||||
if ckpt_display_name:
|
||||
params.append(f"Model: {ckpt_display_name}")
|
||||
if hashes:
|
||||
params.append(f"Hashes: {json.dumps(hashes, separators=(',', ':'))}")
|
||||
params.append("Version: ComfyUI")
|
||||
if civitai_resources:
|
||||
params.append(
|
||||
f"Civitai resources: {json.dumps(civitai_resources, separators=(',', ':'))}"
|
||||
)
|
||||
|
||||
# CFG scale (Use guidance if available, otherwise fall back to cfg_scale or cfg)
|
||||
if "guidance" in metadata_dict:
|
||||
add_param_if_not_none(params, "CFG scale", metadata_dict.get("guidance"))
|
||||
elif "cfg_scale" in metadata_dict:
|
||||
add_param_if_not_none(params, "CFG scale", metadata_dict.get("cfg_scale"))
|
||||
elif "cfg" in metadata_dict:
|
||||
add_param_if_not_none(params, "CFG scale", metadata_dict.get("cfg"))
|
||||
|
||||
# Seed
|
||||
if "seed" in metadata_dict:
|
||||
add_param_if_not_none(params, "Seed", metadata_dict.get("seed"))
|
||||
|
||||
# Size
|
||||
if "size" in metadata_dict:
|
||||
add_param_if_not_none(params, "Size", metadata_dict.get("size"))
|
||||
|
||||
# Model info
|
||||
if "checkpoint" in metadata_dict:
|
||||
# Ensure checkpoint is a string before processing
|
||||
checkpoint = metadata_dict.get("checkpoint")
|
||||
if checkpoint is not None:
|
||||
# Get model hash
|
||||
model_hash = self.get_checkpoint_hash(checkpoint)
|
||||
|
||||
# Extract basename without path
|
||||
checkpoint_name = os.path.basename(checkpoint)
|
||||
# Remove extension if present
|
||||
checkpoint_name = os.path.splitext(checkpoint_name)[0]
|
||||
|
||||
# Add model hash if available
|
||||
if model_hash:
|
||||
params.append(
|
||||
f"Model hash: {model_hash[:10]}, Model: {checkpoint_name}"
|
||||
)
|
||||
else:
|
||||
params.append(f"Model: {checkpoint_name}")
|
||||
|
||||
# Add LoRA hashes if available
|
||||
if lora_hashes:
|
||||
lora_hash_parts = []
|
||||
for lora_name, hash_value in lora_hashes.items():
|
||||
lora_hash_parts.append(f"{lora_name}: {hash_value[:10]}")
|
||||
|
||||
if lora_hash_parts:
|
||||
params.append(f'Lora hashes: "{", ".join(lora_hash_parts)}"')
|
||||
|
||||
# Combine all parameters with commas
|
||||
metadata_parts.append(", ".join(params))
|
||||
|
||||
# Join all parts with a new line
|
||||
return "\n".join(metadata_parts)
|
||||
lines.append(", ".join(params))
|
||||
return "\n".join(lines)
|
||||
|
||||
# credit to nkchocoai
|
||||
# Add format_filename method to handle pattern substitution
|
||||
@@ -573,10 +797,13 @@ class SaveImageLM:
|
||||
extra_pnginfo=None,
|
||||
lossless_webp=True,
|
||||
quality=100,
|
||||
webp_method=6,
|
||||
jpeg_subsampling=0,
|
||||
embed_workflow=False,
|
||||
save_with_metadata=True,
|
||||
add_counter_to_filename=True,
|
||||
save_as_recipe=False,
|
||||
add_loras_to_prompt=False,
|
||||
):
|
||||
"""Save images with metadata"""
|
||||
results = []
|
||||
@@ -585,7 +812,7 @@ class SaveImageLM:
|
||||
raw_metadata = get_metadata()
|
||||
metadata_dict = MetadataProcessor.to_dict(raw_metadata, id)
|
||||
|
||||
metadata = self.format_metadata(metadata_dict)
|
||||
metadata = self.format_metadata(metadata_dict, add_loras_to_prompt)
|
||||
|
||||
# Process filename_prefix with pattern substitution
|
||||
filename_prefix = self.format_filename(filename_prefix, metadata_dict)
|
||||
@@ -627,15 +854,14 @@ class SaveImageLM:
|
||||
elif file_format == "jpeg":
|
||||
file = base_filename + ".jpg"
|
||||
file_extension = ".jpg"
|
||||
save_kwargs = {"quality": quality, "optimize": True}
|
||||
save_kwargs = {"quality": quality, "optimize": True, "subsampling": jpeg_subsampling}
|
||||
elif file_format == "webp":
|
||||
file = base_filename + ".webp"
|
||||
file_extension = ".webp"
|
||||
# Add optimization param to control performance
|
||||
save_kwargs = {
|
||||
"quality": quality,
|
||||
"lossless": lossless_webp,
|
||||
"method": 0,
|
||||
"method": webp_method,
|
||||
}
|
||||
else:
|
||||
raise ValueError(f"Unsupported file format: {file_format}")
|
||||
@@ -722,10 +948,13 @@ class SaveImageLM:
|
||||
extra_pnginfo=None,
|
||||
lossless_webp=True,
|
||||
quality=100,
|
||||
webp_method=6,
|
||||
jpeg_subsampling=0,
|
||||
embed_workflow=False,
|
||||
save_with_metadata=True,
|
||||
add_counter_to_filename=True,
|
||||
save_as_recipe=False,
|
||||
add_loras_to_prompt=False,
|
||||
):
|
||||
"""Process and save image with metadata"""
|
||||
# Make sure the output directory exists
|
||||
@@ -751,10 +980,13 @@ class SaveImageLM:
|
||||
extra_pnginfo,
|
||||
lossless_webp,
|
||||
quality,
|
||||
webp_method,
|
||||
jpeg_subsampling,
|
||||
embed_workflow,
|
||||
save_with_metadata,
|
||||
add_counter_to_filename,
|
||||
save_as_recipe,
|
||||
add_loras_to_prompt,
|
||||
)
|
||||
|
||||
return {
|
||||
|
||||
+27
-6
@@ -1,12 +1,27 @@
|
||||
import logging
|
||||
import os
|
||||
from typing import List, Tuple
|
||||
import comfy.sd # type: ignore
|
||||
from typing import Any, List, Tuple
|
||||
import comfy.sd # pyright: ignore[reportMissingImports]
|
||||
from ..utils.utils import get_checkpoint_info_absolute, _format_model_name_for_comfyui
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _reload_gguf_unet(
|
||||
unet_path: str, weight_dtype: str, disable_dynamic: bool = False
|
||||
) -> object:
|
||||
"""Reload a GGUF diffusion model from disk (cached_patcher_init factory).
|
||||
|
||||
Mirrors the GGUF branch of UNETLoaderLM.load_unet so ModelPatcher
|
||||
deepclone/dynamic machinery can rebuild GGUF models with the correct
|
||||
GGMLOps. ``disable_dynamic`` is accepted for signature compatibility
|
||||
with core ComfyUI loaders.
|
||||
"""
|
||||
loader = UNETLoaderLM()
|
||||
model, = loader._load_gguf_unet(unet_path, unet_path, weight_dtype)
|
||||
return model
|
||||
|
||||
|
||||
class UNETLoaderLM:
|
||||
"""UNET Loader with support for extra folder paths
|
||||
|
||||
@@ -19,9 +34,9 @@ class UNETLoaderLM:
|
||||
CATEGORY = "Lora Manager/loaders"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
def INPUT_TYPES(cls):
|
||||
# Get list of unet names from scanner (includes extra folder paths)
|
||||
unet_names = s._get_unet_names()
|
||||
unet_names = cls._get_unet_names()
|
||||
return {
|
||||
"required": {
|
||||
"unet_name": (
|
||||
@@ -90,7 +105,7 @@ class UNETLoaderLM:
|
||||
logger.error(f"Error getting unet names: {e}")
|
||||
return []
|
||||
|
||||
def load_unet(self, unet_name: str, weight_dtype: str) -> Tuple:
|
||||
def load_unet(self, unet_name: str, weight_dtype: str) -> Tuple[Any, ...]:
|
||||
"""Load a diffusion model by name, supporting extra folder paths
|
||||
|
||||
Args:
|
||||
@@ -133,7 +148,7 @@ class UNETLoaderLM:
|
||||
|
||||
def _load_gguf_unet(
|
||||
self, unet_path: str, unet_name: str, weight_dtype: str
|
||||
) -> Tuple:
|
||||
) -> Tuple[Any, ...]:
|
||||
"""Load a GGUF format diffusion model
|
||||
|
||||
Args:
|
||||
@@ -196,6 +211,12 @@ class UNETLoaderLM:
|
||||
# Wrap with GGUFModelPatcher
|
||||
model = GGUFModelPatcher.clone(model)
|
||||
|
||||
# Register a reload factory so the MODEL carries its source path
|
||||
# (cached_patcher_init) like core ComfyUI loaders do — required
|
||||
# for model-name extraction downstream and for ModelPatcher
|
||||
# deepclone/dynamic machinery.
|
||||
model.cached_patcher_init = (_reload_gguf_unet, (unet_path, weight_dtype))
|
||||
|
||||
return (model,)
|
||||
|
||||
except Exception as e:
|
||||
|
||||
+7
-3
@@ -1,3 +1,6 @@
|
||||
from typing import Any
|
||||
|
||||
|
||||
class AnyType(str):
|
||||
"""A special class that is always equal in not equal comparisons. Credit to pythongosssss"""
|
||||
|
||||
@@ -6,7 +9,7 @@ class AnyType(str):
|
||||
|
||||
|
||||
# Credit to Regis Gaughan, III (rgthree)
|
||||
class FlexibleOptionalInputType(dict):
|
||||
class FlexibleOptionalInputType(dict[str, Any]):
|
||||
"""A special class to make flexible nodes that pass data to our python handlers.
|
||||
|
||||
Enables both flexible/dynamic input types (like for Any Switch) or a dynamic number of inputs
|
||||
@@ -23,6 +26,7 @@ class FlexibleOptionalInputType(dict):
|
||||
"""
|
||||
|
||||
def __init__(self, type):
|
||||
super().__init__()
|
||||
self.type = type
|
||||
|
||||
def __getitem__(self, key):
|
||||
@@ -40,7 +44,7 @@ import re
|
||||
import logging
|
||||
import copy
|
||||
import sys
|
||||
import folder_paths # type: ignore
|
||||
import folder_paths # pyright: ignore[reportMissingImports]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -70,7 +74,7 @@ def extract_lora_name(lora_path):
|
||||
return apply_lora_syntax_format(name_no_ext)
|
||||
|
||||
|
||||
def parse_lora_syntax(text: str) -> list[dict]:
|
||||
def parse_lora_syntax(text: str) -> list[dict[str, Any]]:
|
||||
"""Parse <lora:name:strength> syntax from text input into a list of dicts.
|
||||
|
||||
Each entry contains: name, model_strength, clip_strength.
|
||||
|
||||
+17
-5
@@ -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.
|
||||
"""Base classes for recipe parsers."""
|
||||
|
||||
import json
|
||||
@@ -38,7 +42,7 @@ class RecipeMetadataParser(ABC):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
async def populate_lora_from_civitai(lora_entry: Dict[str, Any], civitai_info_tuple: Tuple[Dict[str, Any], Optional[str]],
|
||||
async def populate_lora_from_civitai(lora_entry: Dict[str, Any], civitai_info_tuple: Tuple[Dict[str, Any] | None, str | None] | Dict[str, Any],
|
||||
recipe_scanner=None, base_model_counts=None, hash_value=None) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Populate a lora entry with information from Civitai API response
|
||||
@@ -175,10 +179,18 @@ class RecipeMetadataParser(ABC):
|
||||
lora_entry['localPath'] = local_path
|
||||
lora_entry['file_name'] = os.path.splitext(os.path.basename(local_path))[0]
|
||||
|
||||
# Get thumbnail from local preview if available
|
||||
# Get thumbnail from local preview if available.
|
||||
# Match the cache item by local path first (get_path_by_hash
|
||||
# cascade: 10-char autov2 / 12-char autov3), then by hash.
|
||||
lora_cache = await lora_scanner.get_cached_data()
|
||||
lora_item = next((item for item in lora_cache.raw_data
|
||||
if item['sha256'].lower() == lora_entry['hash'].lower()), None)
|
||||
h = (lora_entry.get("hash") or "").lower()
|
||||
lora_item = next((item for item in lora_cache.raw_data
|
||||
if (item.get("file_path") or "") == local_path), None)
|
||||
if lora_item is None:
|
||||
lora_item = next((item for item in lora_cache.raw_data
|
||||
if (item.get("sha256") or "").lower() == h
|
||||
or (item.get("autov3") or "").lower() == h
|
||||
or (item.get("sha256") or "")[:10].lower() == h), None)
|
||||
if lora_item and 'preview_url' in lora_item:
|
||||
lora_entry['thumbnailUrl'] = config.get_preview_static_url(lora_item['preview_url'])
|
||||
except Exception as e:
|
||||
@@ -194,7 +206,7 @@ class RecipeMetadataParser(ABC):
|
||||
return lora_entry
|
||||
|
||||
@staticmethod
|
||||
async def populate_checkpoint_from_civitai(checkpoint: Dict[str, Any], civitai_info: Dict[str, Any]) -> Dict[str, Any]:
|
||||
async def populate_checkpoint_from_civitai(checkpoint: Dict[str, Any], civitai_info: Dict[str, Any] | Tuple[Dict[str, Any] | None, str | None] | None) -> Dict[str, Any]:
|
||||
"""
|
||||
Populate checkpoint information from Civitai API response
|
||||
|
||||
|
||||
@@ -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 logging
|
||||
import json
|
||||
import os
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Factory for creating recipe metadata parsers."""
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
from .parsers import (
|
||||
RecipeFormatParser,
|
||||
ComfyMetadataParser,
|
||||
@@ -31,7 +32,8 @@ class RecipeParserFactory:
|
||||
# First, try CivitaiApiMetadataParser for dict input
|
||||
if isinstance(metadata, dict):
|
||||
try:
|
||||
if CivitaiApiMetadataParser().is_metadata_matching(metadata):
|
||||
user_comment: Any = metadata
|
||||
if CivitaiApiMetadataParser().is_metadata_matching(user_comment):
|
||||
return CivitaiApiMetadataParser()
|
||||
except Exception as e:
|
||||
logger.debug(f"CivitaiApiMetadataParser check failed: {e}")
|
||||
|
||||
@@ -52,7 +52,7 @@ class AutomaticMetadataParser(RecipeMetadataParser):
|
||||
negative_and_params = ""
|
||||
|
||||
# Initialize metadata
|
||||
metadata = {
|
||||
metadata: Dict[str, Any] = {
|
||||
"prompt": prompt,
|
||||
"loras": []
|
||||
}
|
||||
|
||||
@@ -4,7 +4,7 @@ import json
|
||||
import logging
|
||||
from typing import Dict, Any, Union
|
||||
from ..base import RecipeMetadataParser
|
||||
from ..constants import GEN_PARAM_KEYS
|
||||
from ..constants import GEN_PARAM_KEYS, VALID_LORA_TYPES
|
||||
from ...services.metadata_service import get_default_metadata_provider
|
||||
from ...config import config
|
||||
|
||||
@@ -14,15 +14,16 @@ logger = logging.getLogger(__name__)
|
||||
class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
"""Parser for Civitai image metadata format"""
|
||||
|
||||
def is_metadata_matching(self, metadata) -> bool:
|
||||
def is_metadata_matching(self, user_comment) -> bool:
|
||||
"""Check if the metadata matches the Civitai image metadata format
|
||||
|
||||
Args:
|
||||
metadata: The metadata from the image (dict)
|
||||
user_comment: The metadata from the image (dict)
|
||||
|
||||
Returns:
|
||||
bool: True if this parser can handle the metadata
|
||||
"""
|
||||
metadata = user_comment
|
||||
if not metadata or not isinstance(metadata, dict):
|
||||
return False
|
||||
|
||||
@@ -73,7 +74,7 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
|
||||
return False
|
||||
|
||||
async def parse_metadata( # type: ignore[override]
|
||||
async def parse_metadata( # pyright: ignore[reportIncompatibleMethodOverride]
|
||||
self, user_comment, recipe_scanner=None, civitai_client=None,
|
||||
local_cache: dict[str, Any] | None = None,
|
||||
) -> Dict[str, Any]:
|
||||
@@ -89,8 +90,7 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
Returns:
|
||||
Dict containing parsed recipe data
|
||||
"""
|
||||
metadata: Dict[str, Any] = user_comment # type: ignore[assignment]
|
||||
metadata = user_comment
|
||||
metadata: Dict[str, Any] = user_comment
|
||||
try:
|
||||
# Get metadata provider instead of using civitai_client directly
|
||||
metadata_provider = await get_default_metadata_provider()
|
||||
@@ -116,7 +116,7 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
metadata = inner_meta
|
||||
|
||||
# Initialize result structure
|
||||
result = {
|
||||
result: Dict[str, Any] = {
|
||||
"base_model": None,
|
||||
"loras": [],
|
||||
"model": None,
|
||||
@@ -125,10 +125,10 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
}
|
||||
|
||||
# Track already added LoRAs to prevent duplicates
|
||||
added_loras = {} # key: model_version_id or hash, value: index in result["loras"]
|
||||
added_loras: Dict[str, Any] = {} # key: model_version_id or hash, value: index in result["loras"]
|
||||
|
||||
# Extract hash information from hashes field for LoRA matching
|
||||
lora_hashes = {}
|
||||
lora_hashes: Dict[str, Any] = {}
|
||||
if "hashes" in metadata and isinstance(metadata["hashes"], dict):
|
||||
for key, hash_value in metadata["hashes"].items():
|
||||
key_str = str(key)
|
||||
@@ -184,7 +184,7 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
if model_info:
|
||||
result["base_model"] = model_info.get("baseModel", "")
|
||||
|
||||
base_model_counts = {}
|
||||
base_model_counts: Dict[str, int] = {}
|
||||
|
||||
# Process standard resources array
|
||||
if "resources" in metadata and isinstance(metadata["resources"], list):
|
||||
@@ -196,7 +196,7 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
# identification because it has an explicit type field and hash,
|
||||
# unlike modelVersionIds which is a flat list with no type info.
|
||||
if resource_type == "model":
|
||||
checkpoint_entry = {
|
||||
checkpoint_entry: Dict[str, Any] = {
|
||||
"id": 0,
|
||||
"modelId": 0,
|
||||
"name": resource.get("name", "Unknown Model"),
|
||||
@@ -216,7 +216,8 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
# Try to look up base model from the checkpoint hash
|
||||
cp_hash = checkpoint_entry.get("hash")
|
||||
if cp_hash and metadata_provider:
|
||||
local_cached = local_cache.get(cp_hash) if local_cache else None
|
||||
# local_cache keys are stored lowercase
|
||||
local_cached = local_cache.get(cp_hash.lower()) if local_cache else None
|
||||
if local_cached:
|
||||
self._populate_entry_from_cache(
|
||||
checkpoint_entry, local_cached
|
||||
@@ -294,8 +295,15 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
|
||||
# Try to get info from Civitai if hash is available
|
||||
if lora_hash and metadata_provider:
|
||||
local_cached = local_cache.get(lora_hash) if local_cache else None
|
||||
# local_cache keys are stored lowercase
|
||||
local_cached = local_cache.get(lora_hash.lower()) if local_cache else None
|
||||
if local_cached:
|
||||
cached_type = self._cache_item_model_type(local_cached)
|
||||
if cached_type and cached_type not in VALID_LORA_TYPES:
|
||||
logger.debug(
|
||||
f"Skipping non-LoRA cache item for hash {lora_hash}"
|
||||
)
|
||||
continue
|
||||
self._populate_entry_from_cache(
|
||||
lora_entry, local_cached
|
||||
)
|
||||
@@ -304,6 +312,12 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
added_loras[str(lora_entry["id"])] = len(
|
||||
result["loras"]
|
||||
)
|
||||
# Mirror base.py:150-151 counts for API-path loras
|
||||
bm = local_cached.get("base_model") or ""
|
||||
if bm:
|
||||
base_model_counts[bm] = base_model_counts.get(
|
||||
bm, 0
|
||||
) + 1
|
||||
else:
|
||||
try:
|
||||
civitai_info = (
|
||||
@@ -649,30 +663,47 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
}
|
||||
|
||||
if metadata_provider:
|
||||
try:
|
||||
civitai_info = await metadata_provider.get_model_by_hash(
|
||||
lora_hash
|
||||
)
|
||||
|
||||
populated_entry = await self.populate_lora_from_civitai(
|
||||
lora_entry,
|
||||
civitai_info,
|
||||
recipe_scanner,
|
||||
base_model_counts,
|
||||
lora_hash,
|
||||
)
|
||||
|
||||
if populated_entry is None:
|
||||
# local_cache keys are stored lowercase
|
||||
local_cached = local_cache.get(lora_hash.lower()) if local_cache else None
|
||||
if local_cached:
|
||||
cached_type = self._cache_item_model_type(local_cached)
|
||||
if cached_type and cached_type not in VALID_LORA_TYPES:
|
||||
logger.debug(
|
||||
f"Skipping non-LoRA cache item for hash {lora_hash}"
|
||||
)
|
||||
continue
|
||||
|
||||
lora_entry = populated_entry
|
||||
|
||||
self._populate_entry_from_cache(lora_entry, local_cached)
|
||||
# Mirror base.py:150-151 counts for API-path loras
|
||||
bm = local_cached.get("base_model") or ""
|
||||
if bm:
|
||||
base_model_counts[bm] = base_model_counts.get(bm, 0) + 1
|
||||
if "id" in lora_entry and lora_entry["id"]:
|
||||
added_loras[str(lora_entry["id"])] = len(result["loras"])
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error fetching Civitai info for LoRA hash {lora_hash}: {e}"
|
||||
)
|
||||
else:
|
||||
try:
|
||||
civitai_info = await metadata_provider.get_model_by_hash(
|
||||
lora_hash
|
||||
)
|
||||
|
||||
populated_entry = await self.populate_lora_from_civitai(
|
||||
lora_entry,
|
||||
civitai_info,
|
||||
recipe_scanner,
|
||||
base_model_counts,
|
||||
lora_hash,
|
||||
)
|
||||
|
||||
if populated_entry is None:
|
||||
continue
|
||||
|
||||
lora_entry = populated_entry
|
||||
|
||||
if "id" in lora_entry and lora_entry["id"]:
|
||||
added_loras[str(lora_entry["id"])] = len(result["loras"])
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error fetching Civitai info for LoRA hash {lora_hash}: {e}"
|
||||
)
|
||||
|
||||
added_loras[lora_hash] = len(result["loras"])
|
||||
result["loras"].append(lora_entry)
|
||||
@@ -711,32 +742,51 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
|
||||
# Try to get info from Civitai if hash is available
|
||||
if lora_entry["hash"] and metadata_provider:
|
||||
try:
|
||||
civitai_info = await metadata_provider.get_model_by_hash(
|
||||
lora_hash
|
||||
)
|
||||
|
||||
populated_entry = await self.populate_lora_from_civitai(
|
||||
lora_entry,
|
||||
civitai_info,
|
||||
recipe_scanner,
|
||||
base_model_counts,
|
||||
lora_hash,
|
||||
)
|
||||
|
||||
if populated_entry is None:
|
||||
# local_cache keys are stored lowercase
|
||||
local_cached = local_cache.get(lora_hash.lower()) if local_cache else None
|
||||
if local_cached:
|
||||
cached_type = self._cache_item_model_type(local_cached)
|
||||
if cached_type and cached_type not in VALID_LORA_TYPES:
|
||||
logger.debug(
|
||||
f"Skipping non-LoRA cache item for hash {lora_hash}"
|
||||
)
|
||||
lora_index += 1
|
||||
continue # Skip invalid LoRA types
|
||||
|
||||
lora_entry = populated_entry
|
||||
|
||||
continue # Skip non-LoRA cache items
|
||||
self._populate_entry_from_cache(lora_entry, local_cached)
|
||||
# Mirror base.py:150-151 counts for API-path loras
|
||||
bm = local_cached.get("base_model") or ""
|
||||
if bm:
|
||||
base_model_counts[bm] = base_model_counts.get(bm, 0) + 1
|
||||
# If we have a version ID from Civitai, track it for deduplication
|
||||
if "id" in lora_entry and lora_entry["id"]:
|
||||
added_loras[str(lora_entry["id"])] = len(result["loras"])
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error fetching Civitai info for LoRA hash {lora_entry['hash']}: {e}"
|
||||
)
|
||||
else:
|
||||
try:
|
||||
civitai_info = await metadata_provider.get_model_by_hash(
|
||||
lora_hash
|
||||
)
|
||||
|
||||
populated_entry = await self.populate_lora_from_civitai(
|
||||
lora_entry,
|
||||
civitai_info,
|
||||
recipe_scanner,
|
||||
base_model_counts,
|
||||
lora_hash,
|
||||
)
|
||||
|
||||
if populated_entry is None:
|
||||
lora_index += 1
|
||||
continue # Skip invalid LoRA types
|
||||
|
||||
lora_entry = populated_entry
|
||||
|
||||
# If we have a version ID from Civitai, track it for deduplication
|
||||
if "id" in lora_entry and lora_entry["id"]:
|
||||
added_loras[str(lora_entry["id"])] = len(result["loras"])
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error fetching Civitai info for LoRA hash {lora_entry['hash']}: {e}"
|
||||
)
|
||||
|
||||
# Track by hash if we have it
|
||||
if lora_hash:
|
||||
@@ -795,3 +845,14 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
base_model = cache_item.get("base_model", "")
|
||||
if base_model:
|
||||
entry["baseModel"] = base_model
|
||||
|
||||
@staticmethod
|
||||
def _cache_item_model_type(cache_item: dict[str, Any]) -> str:
|
||||
"""Lowercased civitai.model.type of a cache item, or '' when unknown."""
|
||||
civ = cache_item.get("civitai")
|
||||
if not isinstance(civ, dict):
|
||||
return ""
|
||||
model_info = civ.get("model")
|
||||
if not isinstance(model_info, dict):
|
||||
return ""
|
||||
return (model_info.get("type") or "").lower()
|
||||
|
||||
@@ -30,7 +30,7 @@ class MetaFormatParser(RecipeMetadataParser):
|
||||
prompt = parts[0].strip()
|
||||
|
||||
# Initialize metadata
|
||||
metadata = {"prompt": prompt, "loras": []}
|
||||
metadata: Dict[str, Any] = {"prompt": prompt, "loras": []}
|
||||
|
||||
# Extract negative prompt and parameters if available
|
||||
if len(parts) > 1:
|
||||
|
||||
@@ -91,7 +91,15 @@ class RecipeFormatParser(RecipeMetadataParser):
|
||||
exists_locally = lora_scanner.has_hash(lora['hash'])
|
||||
if exists_locally:
|
||||
lora_cache = await lora_scanner.get_cached_data()
|
||||
lora_item = next((item for item in lora_cache.raw_data if item['sha256'].lower() == lora['hash'].lower()), None)
|
||||
# Cascade match: full sha256, stored autov3, or autov2 (sha256[:10]).
|
||||
h = (lora.get('hash') or '').lower()
|
||||
lora_item = next(
|
||||
(item for item in lora_cache.raw_data
|
||||
if (item.get("sha256") or "").lower() == h
|
||||
or (item.get("autov3") or "").lower() == h
|
||||
or (item.get("sha256") or "")[:10].lower() == h),
|
||||
None
|
||||
)
|
||||
if lora_item:
|
||||
lora_entry['existsLocally'] = True
|
||||
lora_entry['inLibrary'] = True
|
||||
@@ -148,7 +156,7 @@ class RecipeFormatParser(RecipeMetadataParser):
|
||||
checkpoint_data = recipe_metadata.get('checkpoint') or {}
|
||||
if isinstance(checkpoint_data, dict) and checkpoint_data:
|
||||
version_id = checkpoint_data.get('modelVersionId') or checkpoint_data.get('id')
|
||||
checkpoint_entry = {
|
||||
checkpoint_entry: Dict[str, Any] = {
|
||||
'id': version_id or 0,
|
||||
'modelId': checkpoint_data.get('modelId', 0),
|
||||
'name': checkpoint_data.get('name', 'Unknown Checkpoint'),
|
||||
|
||||
@@ -2,7 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING, Callable, Dict, Mapping
|
||||
from typing import TYPE_CHECKING, Awaitable, Callable, Dict, Mapping
|
||||
|
||||
import jinja2
|
||||
from aiohttp import web
|
||||
@@ -84,7 +84,7 @@ class BaseModelRoutes(ABC):
|
||||
self.metadata_progress_callback = WebSocketBroadcastCallback()
|
||||
|
||||
self._handler_set: ModelHandlerSet | None = None
|
||||
self._handler_mapping: Dict[str, Callable[[web.Request], web.StreamResponse]] | None = None
|
||||
self._handler_mapping: Dict[str, Callable[[web.Request], Awaitable[web.Response]]] | None = None
|
||||
|
||||
self._preview_service = PreviewAssetService(
|
||||
metadata_manager=MetadataManager,
|
||||
@@ -131,7 +131,7 @@ class BaseModelRoutes(ABC):
|
||||
self._handler_set = None
|
||||
self._handler_mapping = None
|
||||
|
||||
def _ensure_handler_mapping(self) -> Mapping[str, Callable[[web.Request], web.StreamResponse]]:
|
||||
def _ensure_handler_mapping(self) -> Mapping[str, Callable[[web.Request], Awaitable[web.StreamResponse]]]:
|
||||
if self._handler_mapping is None:
|
||||
handler_set = self._create_handler_set()
|
||||
self._handler_set = handler_set
|
||||
@@ -220,7 +220,7 @@ class BaseModelRoutes(ABC):
|
||||
)
|
||||
|
||||
@property
|
||||
def route_handlers(self) -> Mapping[str, Callable[[web.Request], web.StreamResponse]]:
|
||||
def route_handlers(self) -> Mapping[str, Callable[[web.Request], Awaitable[web.StreamResponse]]]:
|
||||
return self._ensure_handler_mapping()
|
||||
|
||||
def setup_routes(self, app: web.Application, prefix: str) -> None:
|
||||
@@ -237,7 +237,7 @@ class BaseModelRoutes(ABC):
|
||||
"""Setup model-specific routes."""
|
||||
raise NotImplementedError
|
||||
|
||||
def _parse_specific_params(self, request: web.Request) -> Dict:
|
||||
def _parse_specific_params(self, request: web.Request) -> Dict[str, Any]:
|
||||
"""Parse model-specific parameters - to be overridden by subclasses."""
|
||||
return {}
|
||||
|
||||
@@ -253,7 +253,7 @@ class BaseModelRoutes(ABC):
|
||||
"""Find the appropriate model file from the files list - can be overridden by subclasses."""
|
||||
return next((file for file in files if file.get("type") in ("Model", "Diffusion Model") and file.get("primary") is True), None)
|
||||
|
||||
def get_handler(self, name: str) -> Callable[[web.Request], web.StreamResponse]:
|
||||
def get_handler(self, name: str) -> Callable[[web.Request], Awaitable[web.StreamResponse]]:
|
||||
"""Expose handlers for subclasses or tests."""
|
||||
return self._ensure_handler_mapping()[name]
|
||||
|
||||
@@ -285,7 +285,7 @@ class BaseModelRoutes(ABC):
|
||||
)
|
||||
return self.model_lifecycle_service
|
||||
|
||||
def _make_handler_proxy(self, name: str) -> Callable[[web.Request], web.StreamResponse]:
|
||||
def _make_handler_proxy(self, name: str) -> Callable[[web.Request], Awaitable[web.StreamResponse]]:
|
||||
async def proxy(request: web.Request) -> web.StreamResponse:
|
||||
try:
|
||||
handler = self.get_handler(name)
|
||||
|
||||
@@ -4,7 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import Callable, Mapping
|
||||
from typing import Awaitable, Callable, Mapping
|
||||
|
||||
import jinja2
|
||||
from aiohttp import web
|
||||
@@ -61,7 +61,9 @@ class BaseRecipeRoutes:
|
||||
self._i18n_registered = False
|
||||
self._startup_hooks_registered = False
|
||||
self._handler_set: RecipeHandlerSet | None = None
|
||||
self._handler_mapping: dict[str, Callable] | None = None
|
||||
self._handler_mapping: Mapping[
|
||||
str, Callable[[web.Request], Awaitable[web.StreamResponse]]
|
||||
] | None = None
|
||||
|
||||
async def attach_dependencies(self, app: web.Application | None = None) -> None:
|
||||
"""Resolve shared services from the registry."""
|
||||
@@ -84,7 +86,9 @@ class BaseRecipeRoutes:
|
||||
app.on_startup.append(self.attach_dependencies)
|
||||
self._startup_hooks_registered = True
|
||||
|
||||
def to_route_mapping(self) -> Mapping[str, Callable]:
|
||||
def to_route_mapping(
|
||||
self,
|
||||
) -> Mapping[str, Callable[[web.Request], Awaitable[web.StreamResponse]]]:
|
||||
"""Return a mapping of handler name to coroutine for registrar binding."""
|
||||
|
||||
if self._handler_mapping is None:
|
||||
@@ -124,17 +128,17 @@ class BaseRecipeRoutes:
|
||||
or os.environ.get("HF_HUB_DISABLE_TELEMETRY", "0") == "0"
|
||||
)
|
||||
if not standalone_mode:
|
||||
from ..metadata_collector import get_metadata # type: ignore[import-not-found]
|
||||
from ..metadata_collector.metadata_processor import ( # type: ignore[import-not-found]
|
||||
from ..metadata_collector import get_metadata # pyright: ignore[reportMissingImports]
|
||||
from ..metadata_collector.metadata_processor import ( # pyright: ignore[reportMissingImports]
|
||||
MetadataProcessor,
|
||||
)
|
||||
from ..metadata_collector.metadata_registry import ( # type: ignore[import-not-found]
|
||||
from ..metadata_collector.metadata_registry import ( # pyright: ignore[reportMissingImports]
|
||||
MetadataRegistry,
|
||||
)
|
||||
else: # pragma: no cover - optional dependency path
|
||||
get_metadata = None # type: ignore[assignment]
|
||||
MetadataProcessor = None # type: ignore[assignment]
|
||||
MetadataRegistry = None # type: ignore[assignment]
|
||||
get_metadata = None # pyright: ignore[reportAssignmentType]
|
||||
MetadataProcessor = None # pyright: ignore[reportAssignmentType]
|
||||
MetadataRegistry = None # pyright: ignore[reportAssignmentType]
|
||||
|
||||
analysis_service = RecipeAnalysisService(
|
||||
exif_utils=ExifUtils,
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import logging
|
||||
from typing import Dict, List, Set
|
||||
from typing import Any, Dict, List, Set
|
||||
from aiohttp import web
|
||||
|
||||
from .base_model_routes import BaseModelRoutes
|
||||
@@ -28,13 +28,13 @@ class CheckpointRoutes(BaseModelRoutes):
|
||||
# Attach service dependencies
|
||||
self.attach_service(self.service)
|
||||
|
||||
def setup_routes(self, app: web.Application):
|
||||
def setup_routes(self, app: web.Application, prefix: str = "checkpoints"):
|
||||
"""Setup Checkpoint routes"""
|
||||
# Schedule service initialization on app startup
|
||||
app.on_startup.append(lambda _: self.initialize_services())
|
||||
|
||||
|
||||
# Setup common routes with 'checkpoints' prefix (includes page route)
|
||||
super().setup_routes(app, 'checkpoints')
|
||||
super().setup_routes(app, prefix)
|
||||
|
||||
def setup_specific_routes(self, registrar: ModelRouteRegistrar, prefix: str):
|
||||
"""Setup Checkpoint-specific routes"""
|
||||
@@ -53,9 +53,9 @@ class CheckpointRoutes(BaseModelRoutes):
|
||||
"""Get expected model types string for error messages"""
|
||||
return "Checkpoint"
|
||||
|
||||
def _parse_specific_params(self, request: web.Request) -> Dict:
|
||||
def _parse_specific_params(self, request: web.Request) -> Dict[str, Any]:
|
||||
"""Parse Checkpoint-specific parameters"""
|
||||
params: Dict = {}
|
||||
params: Dict[str, Any] = {}
|
||||
|
||||
if 'checkpoint_hash' in request.query:
|
||||
params['hash_filters'] = {'single_hash': request.query['checkpoint_hash'].lower()}
|
||||
@@ -70,7 +70,7 @@ class CheckpointRoutes(BaseModelRoutes):
|
||||
"""Get detailed information for a specific checkpoint by name"""
|
||||
try:
|
||||
name = request.match_info.get('name', '')
|
||||
checkpoint_info = await self.service.get_model_info_by_name(name)
|
||||
checkpoint_info = await self.service.get_model_info_by_name(name) # pyright: ignore[reportAttributeAccessIssue]
|
||||
|
||||
if checkpoint_info:
|
||||
return web.json_response(checkpoint_info)
|
||||
@@ -89,7 +89,7 @@ class CheckpointRoutes(BaseModelRoutes):
|
||||
roots.extend(config.checkpoints_roots or [])
|
||||
roots.extend(config.extra_checkpoints_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:
|
||||
@@ -114,7 +114,7 @@ class CheckpointRoutes(BaseModelRoutes):
|
||||
roots.extend(config.unet_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 and root not in seen:
|
||||
|
||||
@@ -26,13 +26,13 @@ class EmbeddingRoutes(BaseModelRoutes):
|
||||
# Attach service dependencies
|
||||
self.attach_service(self.service)
|
||||
|
||||
def setup_routes(self, app: web.Application):
|
||||
def setup_routes(self, app: web.Application, prefix: str = "embeddings"):
|
||||
"""Setup Embedding routes"""
|
||||
# Schedule service initialization on app startup
|
||||
app.on_startup.append(lambda _: self.initialize_services())
|
||||
|
||||
|
||||
# Setup common routes with 'embeddings' prefix (includes page route)
|
||||
super().setup_routes(app, 'embeddings')
|
||||
super().setup_routes(app, prefix)
|
||||
|
||||
def setup_specific_routes(self, registrar: ModelRouteRegistrar, prefix: str):
|
||||
"""Setup Embedding-specific routes"""
|
||||
@@ -51,7 +51,7 @@ class EmbeddingRoutes(BaseModelRoutes):
|
||||
"""Get detailed information for a specific embedding by name"""
|
||||
try:
|
||||
name = request.match_info.get('name', '')
|
||||
embedding_info = await self.service.get_model_info_by_name(name)
|
||||
embedding_info = await self.service.get_model_info_by_name(name) # pyright: ignore[reportAttributeAccessIssue]
|
||||
|
||||
if embedding_info:
|
||||
return web.json_response(embedding_info)
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Callable, Mapping
|
||||
from typing import Any, Awaitable, Callable, Mapping
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
@@ -35,7 +35,7 @@ class ExampleImagesRoutes:
|
||||
*,
|
||||
ws_manager,
|
||||
download_manager: DownloadManager | None = None,
|
||||
processor=ExampleImagesProcessor,
|
||||
processor: Any = ExampleImagesProcessor,
|
||||
file_manager=ExampleImagesFileManager,
|
||||
cleanup_service: ExampleImagesCleanupService | None = None,
|
||||
) -> None:
|
||||
@@ -46,7 +46,9 @@ class ExampleImagesRoutes:
|
||||
self._file_manager = file_manager
|
||||
self._cleanup_service = cleanup_service or ExampleImagesCleanupService()
|
||||
self._handler_set: ExampleImagesHandlerSet | None = None
|
||||
self._handler_mapping: Mapping[str, Callable[[web.Request], web.StreamResponse]] | None = None
|
||||
self._handler_mapping: Mapping[
|
||||
str, Callable[[web.Request], Awaitable[web.StreamResponse]]
|
||||
] | None = None
|
||||
|
||||
@classmethod
|
||||
def setup_routes(cls, app: web.Application, *, ws_manager) -> None:
|
||||
@@ -61,7 +63,9 @@ class ExampleImagesRoutes:
|
||||
registrar = ExampleImagesRouteRegistrar(app)
|
||||
registrar.register_routes(self.to_route_mapping())
|
||||
|
||||
def to_route_mapping(self) -> Mapping[str, Callable[[web.Request], web.StreamResponse]]:
|
||||
def to_route_mapping(
|
||||
self,
|
||||
) -> Mapping[str, Callable[[web.Request], Awaitable[web.StreamResponse]]]:
|
||||
"""Return the registrar-compatible mapping of handler names to callables."""
|
||||
|
||||
if self._handler_mapping is None:
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Mapping
|
||||
from typing import Awaitable, Callable, Mapping
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
@@ -170,7 +170,7 @@ class ExampleImagesHandlerSet:
|
||||
management: ExampleImagesManagementHandler
|
||||
files: ExampleImagesFileHandler
|
||||
|
||||
def to_route_mapping(self) -> Mapping[str, Callable[[web.Request], web.StreamResponse]]:
|
||||
def to_route_mapping(self) -> Mapping[str, Callable[[web.Request], Awaitable[web.StreamResponse]]]:
|
||||
"""Flatten handler methods into the registrar mapping."""
|
||||
|
||||
return {
|
||||
|
||||
@@ -276,7 +276,7 @@ def _collect_comfyui_session_logs(
|
||||
) -> dict[str, Any]:
|
||||
if log_entries is None:
|
||||
try:
|
||||
import app.logger as comfy_logger
|
||||
import app.logger as comfy_logger # pyright: ignore[reportMissingImports]
|
||||
|
||||
log_entries = list(comfy_logger.get_logs() or [])
|
||||
except Exception as exc: # pragma: no cover - environment dependent
|
||||
@@ -422,10 +422,10 @@ class PromptServerProtocol(Protocol):
|
||||
"""Subset of PromptServer used by the handlers."""
|
||||
|
||||
instance: "PromptServerProtocol"
|
||||
sockets: dict # maps clientId (sid) → WebSocketResponse
|
||||
sockets: dict[str, Any] # maps clientId (sid) → WebSocketResponse
|
||||
|
||||
def send_sync(
|
||||
self, event: str, payload: dict | None = None, sid: str | None = None
|
||||
self, event: str, payload: dict[str, Any] | None = None, sid: str | None = None
|
||||
) -> None: # pragma: no cover - protocol
|
||||
...
|
||||
|
||||
@@ -443,7 +443,12 @@ class UsageStatsFactory(Protocol):
|
||||
class MetadataProviderProtocol(Protocol):
|
||||
async def get_model_versions(
|
||||
self, model_id: int
|
||||
) -> dict | None: # pragma: no cover - protocol
|
||||
) -> dict[str, Any] | None: # pragma: no cover - protocol
|
||||
...
|
||||
|
||||
async def get_user_models(
|
||||
self, username: str, cursor: str | None = None
|
||||
) -> Any: # pragma: no cover - protocol
|
||||
...
|
||||
|
||||
|
||||
@@ -466,16 +471,16 @@ class MetadataArchiveManagerProtocol(Protocol):
|
||||
class BackupServiceProtocol(Protocol):
|
||||
async def create_snapshot(
|
||||
self, *, snapshot_type: str = "manual", persist: bool = False
|
||||
) -> dict: # pragma: no cover - protocol
|
||||
) -> dict[str, Any]: # pragma: no cover - protocol
|
||||
...
|
||||
|
||||
async def restore_snapshot(self, archive_path: str) -> dict: # pragma: no cover - protocol
|
||||
async def restore_snapshot(self, archive_path: str) -> dict[str, Any]: # pragma: no cover - protocol
|
||||
...
|
||||
|
||||
def get_status(self) -> dict: # pragma: no cover - protocol
|
||||
def get_status(self) -> dict[str, Any]: # pragma: no cover - protocol
|
||||
...
|
||||
|
||||
def get_available_snapshots(self) -> list[dict]: # pragma: no cover - protocol
|
||||
def get_available_snapshots(self) -> list[dict[str, Any]]: # pragma: no cover - protocol
|
||||
...
|
||||
|
||||
|
||||
@@ -491,7 +496,7 @@ class NodeRegistry:
|
||||
def __init__(self) -> None:
|
||||
self._lock = asyncio.Lock()
|
||||
# sid → {unique_id → node_info}
|
||||
self._tab_nodes: Dict[str, Dict[str, dict]] = {}
|
||||
self._tab_nodes: Dict[str, Dict[str, dict[str, Any]]] = {}
|
||||
self._ready = asyncio.Event()
|
||||
self._waiting_clients: set[str] = set()
|
||||
|
||||
@@ -504,7 +509,7 @@ class NodeRegistry:
|
||||
# Helpers to build one node dict (extracted so it's reused for each tab)
|
||||
# ------------------------------------------------------------------
|
||||
@staticmethod
|
||||
def _build_node_dict(node: dict) -> dict:
|
||||
def _build_node_dict(node: dict[str, Any]) -> dict[str, Any]:
|
||||
node_id = node["node_id"]
|
||||
graph_id = str(node["graph_id"])
|
||||
unique_id = f"{graph_id}:{node_id}"
|
||||
@@ -513,11 +518,11 @@ class NodeRegistry:
|
||||
bgcolor = node.get("bgcolor") or DEFAULT_NODE_COLOR
|
||||
|
||||
raw_capabilities = node.get("capabilities")
|
||||
capabilities: dict = {}
|
||||
capabilities: dict[str, Any] = {}
|
||||
if isinstance(raw_capabilities, dict):
|
||||
capabilities = dict(raw_capabilities)
|
||||
|
||||
raw_widget_names: list | None = node.get("widget_names")
|
||||
raw_widget_names: list[Any] | None = node.get("widget_names")
|
||||
if not isinstance(raw_widget_names, list):
|
||||
capability_widget_names = capabilities.get("widget_names")
|
||||
raw_widget_names = (
|
||||
@@ -565,9 +570,9 @@ class NodeRegistry:
|
||||
# ------------------------------------------------------------------
|
||||
# Public API
|
||||
# ------------------------------------------------------------------
|
||||
async def register_nodes(self, sid: str, nodes: list[dict]) -> None:
|
||||
async def register_nodes(self, sid: str, nodes: list[dict[str, Any]]) -> None:
|
||||
"""Register/replace the node list for a single ComfyUI tab (identified by *sid*)."""
|
||||
tab_nodes: dict[str, dict] = {}
|
||||
tab_nodes: dict[str, dict[str, Any]] = {}
|
||||
for node in nodes:
|
||||
nd = self._build_node_dict(node)
|
||||
tab_nodes[nd["unique_id"]] = nd
|
||||
@@ -602,7 +607,7 @@ class NodeRegistry:
|
||||
except asyncio.TimeoutError:
|
||||
return False
|
||||
|
||||
async def get_merged_registry(self, active_sids: set[str] | None = None) -> dict:
|
||||
async def get_merged_registry(self, active_sids: set[str] | None = None) -> dict[str, Any]:
|
||||
"""Return the union of all known tab nodes, pruning any tab that is no
|
||||
longer connected."""
|
||||
async with self._lock:
|
||||
@@ -619,8 +624,8 @@ class NodeRegistry:
|
||||
len(stale_sids), stale_sids,
|
||||
)
|
||||
|
||||
merged: dict[str, dict] = {}
|
||||
tab_info: dict[str, dict] = {}
|
||||
merged: dict[str, dict[str, Any]] = {}
|
||||
tab_info: dict[str, dict[str, Any]] = {}
|
||||
for sid, nodes in self._tab_nodes.items():
|
||||
tab_info[sid] = {
|
||||
"node_count": len(nodes),
|
||||
@@ -653,7 +658,7 @@ class SupportersHandler:
|
||||
def __init__(self, logger: logging.Logger | None = None) -> None:
|
||||
self._logger = logger or logging.getLogger(__name__)
|
||||
|
||||
def _load_supporters(self) -> dict:
|
||||
def _load_supporters(self) -> dict[str, Any]:
|
||||
"""Load supporters data from JSON file."""
|
||||
try:
|
||||
current_file = os.path.abspath(__file__)
|
||||
@@ -1229,10 +1234,8 @@ class DoctorHandler:
|
||||
settings_snapshot = _sanitize_sensitive_data(
|
||||
getattr(self._settings, "settings", {}) or {}
|
||||
)
|
||||
startup_messages_getter = getattr(self._settings, "get_startup_messages", None)
|
||||
startup_messages = (
|
||||
list(startup_messages_getter()) if callable(startup_messages_getter) else []
|
||||
)
|
||||
startup_messages_getter: Any = getattr(self._settings, "get_startup_messages", None)
|
||||
startup_messages = list(startup_messages_getter()) if startup_messages_getter else []
|
||||
|
||||
environment = {
|
||||
"app_version": app_version,
|
||||
@@ -1439,7 +1442,7 @@ class SettingsHandler:
|
||||
*,
|
||||
settings_service=None,
|
||||
metadata_provider_updater: Callable[
|
||||
[], Awaitable[None]
|
||||
[], Awaitable[Any]
|
||||
] = update_metadata_providers,
|
||||
downloader_factory: Callable[
|
||||
[], Awaitable[DownloaderProtocol]
|
||||
@@ -1484,8 +1487,8 @@ class SettingsHandler:
|
||||
settings_file = getattr(self._settings, "settings_file", None)
|
||||
if settings_file:
|
||||
response_data["settings_file"] = settings_file
|
||||
messages_getter = getattr(self._settings, "get_startup_messages", None)
|
||||
messages = list(messages_getter()) if callable(messages_getter) else []
|
||||
messages_getter: Any = getattr(self._settings, "get_startup_messages", None)
|
||||
messages = list(messages_getter()) if messages_getter else []
|
||||
return web.json_response(
|
||||
{
|
||||
"success": True,
|
||||
@@ -1562,6 +1565,11 @@ class SettingsHandler:
|
||||
{"success": False, "error": validation_error}
|
||||
)
|
||||
|
||||
if key == "update_channel" and value not in ("release", "nightly"):
|
||||
return web.json_response(
|
||||
{"success": False, "error": "update_channel must be 'release' or 'nightly'"}
|
||||
)
|
||||
|
||||
if value == "__DELETE__" and key in (
|
||||
"proxy_username",
|
||||
"proxy_password",
|
||||
@@ -2000,11 +2008,11 @@ async def _noop_backup_service() -> None:
|
||||
|
||||
@dataclass
|
||||
class ServiceRegistryAdapter:
|
||||
get_lora_scanner: Callable[[], Awaitable]
|
||||
get_checkpoint_scanner: Callable[[], Awaitable]
|
||||
get_embedding_scanner: Callable[[], Awaitable]
|
||||
get_downloaded_version_history_service: Callable[[], Awaitable]
|
||||
get_backup_service: Callable[[], Awaitable] = _noop_backup_service
|
||||
get_lora_scanner: Callable[[], Awaitable[Any]]
|
||||
get_checkpoint_scanner: Callable[[], Awaitable[Any]]
|
||||
get_embedding_scanner: Callable[[], Awaitable[Any]]
|
||||
get_downloaded_version_history_service: Callable[[], Awaitable[Any]]
|
||||
get_backup_service: Callable[[], Awaitable[Any]] = _noop_backup_service
|
||||
|
||||
|
||||
class ModelLibraryHandler:
|
||||
@@ -2045,8 +2053,8 @@ class ModelLibraryHandler:
|
||||
return await self._service_registry.get_downloaded_version_history_service()
|
||||
|
||||
@staticmethod
|
||||
def _with_downloaded_flag(versions: list[dict]) -> list[dict]:
|
||||
enriched: list[dict] = []
|
||||
def _with_downloaded_flag(versions: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
enriched: list[dict[str, Any]] = []
|
||||
for version in versions:
|
||||
entry = dict(version)
|
||||
entry.setdefault("hasBeenDownloaded", True)
|
||||
@@ -2239,7 +2247,7 @@ class ModelLibraryHandler:
|
||||
checkpoint_scanner = await self._service_registry.get_checkpoint_scanner()
|
||||
embedding_scanner = await self._service_registry.get_embedding_scanner()
|
||||
|
||||
results: list[dict] = []
|
||||
results: list[dict[str, Any]] = []
|
||||
for model_id in model_ids:
|
||||
lora_versions = await lora_scanner.get_model_versions_by_id(model_id)
|
||||
if lora_versions:
|
||||
@@ -2348,7 +2356,7 @@ class ModelLibraryHandler:
|
||||
)
|
||||
|
||||
try:
|
||||
model_version_id = int(data.get("modelVersionId"))
|
||||
model_version_id = int(data.get("modelVersionId")) # pyright: ignore[reportArgumentType]
|
||||
except (TypeError, ValueError):
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Parameter modelVersionId must be an integer"},
|
||||
@@ -2460,10 +2468,11 @@ class ModelLibraryHandler:
|
||||
"checkpoint": checkpoint_scanner,
|
||||
"embedding": embedding_scanner,
|
||||
}
|
||||
scanner = scanner_map.get(found_type)
|
||||
scanner = scanner_map.get(found_type or "")
|
||||
if scanner:
|
||||
persist = getattr(scanner, "_persist_current_cache", None)
|
||||
if callable(persist):
|
||||
scanner.bump_cache_version()
|
||||
persist: Any = getattr(scanner, "_persist_current_cache", None)
|
||||
if persist:
|
||||
await persist()
|
||||
|
||||
history_service = await self._get_download_history_service()
|
||||
@@ -2585,6 +2594,8 @@ class ModelLibraryHandler:
|
||||
status=400,
|
||||
)
|
||||
|
||||
cursor = request.query.get("cursor")
|
||||
|
||||
metadata_provider = await self._metadata_provider_factory()
|
||||
if not metadata_provider:
|
||||
return web.json_response(
|
||||
@@ -2593,7 +2604,7 @@ class ModelLibraryHandler:
|
||||
)
|
||||
|
||||
try:
|
||||
models = await metadata_provider.get_user_models(username)
|
||||
result = await metadata_provider.get_user_models(username, cursor)
|
||||
except NotImplementedError:
|
||||
return web.json_response(
|
||||
{
|
||||
@@ -2603,14 +2614,35 @@ class ModelLibraryHandler:
|
||||
status=501,
|
||||
)
|
||||
|
||||
if models is None:
|
||||
if result is None:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Failed to fetch user models"},
|
||||
status=502,
|
||||
)
|
||||
|
||||
if isinstance(result, dict):
|
||||
models = result.get("items")
|
||||
next_cursor = result.get("nextCursor")
|
||||
else:
|
||||
# Defensive: tolerate providers that still return a raw list
|
||||
models = result
|
||||
next_cursor = None
|
||||
|
||||
if not isinstance(models, list):
|
||||
models = []
|
||||
if next_cursor is not None and not isinstance(next_cursor, str):
|
||||
next_cursor = str(next_cursor)
|
||||
|
||||
estimated_total = None
|
||||
if cursor is None:
|
||||
get_count = getattr(metadata_provider, "get_creator_model_count", None)
|
||||
if get_count is not None:
|
||||
try:
|
||||
estimated_total = await get_count(username)
|
||||
except Exception: # best-effort only
|
||||
estimated_total = None
|
||||
if not isinstance(estimated_total, int):
|
||||
estimated_total = None
|
||||
|
||||
lora_scanner = await self._service_registry.get_lora_scanner()
|
||||
checkpoint_scanner = await self._service_registry.get_checkpoint_scanner()
|
||||
@@ -2621,15 +2653,16 @@ class ModelLibraryHandler:
|
||||
}
|
||||
lora_type_aliases = {model_type.lower() for model_type in VALID_LORA_TYPES}
|
||||
|
||||
type_scanner_map: Dict[str, object | None] = {
|
||||
type_scanner_map: Dict[str, Any] = {
|
||||
**{alias: lora_scanner for alias in lora_type_aliases},
|
||||
"checkpoint": checkpoint_scanner,
|
||||
"textualinversion": embedding_scanner,
|
||||
}
|
||||
|
||||
versions: list[dict] = []
|
||||
versions: list[dict[str, Any]] = []
|
||||
history_service = await self._get_download_history_service()
|
||||
model_ids: list[int] = []
|
||||
model_count = 0
|
||||
for model in models:
|
||||
try:
|
||||
model_ids.append(int(model.get("id")))
|
||||
@@ -2663,6 +2696,8 @@ class ModelLibraryHandler:
|
||||
if model_type not in normalized_allowed_types:
|
||||
continue
|
||||
|
||||
model_count += 1
|
||||
|
||||
scanner = type_scanner_map.get(model_type)
|
||||
if scanner is None:
|
||||
return web.json_response(
|
||||
@@ -2676,6 +2711,8 @@ class ModelLibraryHandler:
|
||||
tags_value = model.get("tags")
|
||||
tags = tags_value if isinstance(tags_value, list) else []
|
||||
model_id = model.get("id")
|
||||
if model_id is None:
|
||||
continue
|
||||
try:
|
||||
model_id_int = int(model_id)
|
||||
except (TypeError, ValueError):
|
||||
@@ -2691,6 +2728,8 @@ class ModelLibraryHandler:
|
||||
continue
|
||||
|
||||
version_id = version.get("id")
|
||||
if version_id is None:
|
||||
continue
|
||||
try:
|
||||
version_id_int = int(version_id)
|
||||
except (TypeError, ValueError):
|
||||
@@ -2728,7 +2767,15 @@ class ModelLibraryHandler:
|
||||
)
|
||||
|
||||
return web.json_response(
|
||||
{"success": True, "username": username, "versions": versions}
|
||||
{
|
||||
"success": True,
|
||||
"username": username,
|
||||
"versions": versions,
|
||||
"modelCount": model_count,
|
||||
"nextCursor": next_cursor,
|
||||
"hasMore": next_cursor is not None,
|
||||
"estimatedTotal": estimated_total,
|
||||
}
|
||||
)
|
||||
except Exception as exc: # pragma: no cover - defensive logging
|
||||
logger.error("Failed to get Civitai user models: %s", exc, exc_info=True)
|
||||
@@ -2744,7 +2791,7 @@ class MetadataArchiveHandler:
|
||||
] = get_metadata_archive_manager,
|
||||
settings_service=None,
|
||||
metadata_provider_updater: Callable[
|
||||
[], Awaitable[None]
|
||||
[], Awaitable[Any]
|
||||
] = update_metadata_providers,
|
||||
) -> None:
|
||||
self._metadata_archive_manager_factory = metadata_archive_manager_factory
|
||||
@@ -2891,7 +2938,7 @@ class BackupHandler:
|
||||
|
||||
if request.content_type.startswith("multipart/"):
|
||||
reader = await request.multipart()
|
||||
field = await reader.next()
|
||||
field: Any = await reader.next()
|
||||
uploaded = False
|
||||
while field is not None:
|
||||
if getattr(field, "filename", None):
|
||||
@@ -3510,7 +3557,7 @@ class NodeRegistryHandler:
|
||||
except (TypeError, ValueError):
|
||||
parsed_node_id = node_identifier
|
||||
|
||||
payload: dict = {
|
||||
payload: dict[str, Any] = {
|
||||
"id": parsed_node_id,
|
||||
"value": value,
|
||||
"mode": mode,
|
||||
@@ -3634,7 +3681,7 @@ class NodeRegistryHandler:
|
||||
except (TypeError, ValueError):
|
||||
parsed_node_id = node_identifier
|
||||
|
||||
payload: dict = {
|
||||
payload: dict[str, Any] = {
|
||||
"id": parsed_node_id,
|
||||
"value": value,
|
||||
"mode": mode,
|
||||
@@ -3701,8 +3748,8 @@ class MiscHandlerSet:
|
||||
doctor: DoctorHandler,
|
||||
example_workflows: ExampleWorkflowsHandler,
|
||||
base_model: BaseModelHandlerSet,
|
||||
hf_handler: HfHandler | None = None,
|
||||
agent_handler: AgentHandler | None = None,
|
||||
hf_handler: Any = None,
|
||||
agent_handler: Any = None,
|
||||
) -> None:
|
||||
self.health = health
|
||||
self.settings = settings
|
||||
|
||||
@@ -71,7 +71,7 @@ class ModelPageView:
|
||||
self._server_i18n = server_i18n
|
||||
self._logger = logger
|
||||
|
||||
def _load_supporters(self) -> dict:
|
||||
def _load_supporters(self) -> dict[str, Any]:
|
||||
"""Load supporters data from JSON file."""
|
||||
try:
|
||||
current_file = os.path.abspath(__file__)
|
||||
@@ -152,7 +152,7 @@ class ModelPageView:
|
||||
self._template_env.filters["t"] = (
|
||||
self._server_i18n.create_template_filter()
|
||||
)
|
||||
self._template_env._i18n_filter_added = True # type: ignore[attr-defined]
|
||||
self._template_env._i18n_filter_added = True # pyright: ignore[reportAttributeAccessIssue]
|
||||
|
||||
from ...services.llm_service import PROVIDER_PRESETS
|
||||
|
||||
@@ -199,7 +199,7 @@ class ModelListingHandler:
|
||||
self,
|
||||
*,
|
||||
service,
|
||||
parse_specific_params: Callable[[web.Request], Dict],
|
||||
parse_specific_params: Callable[[web.Request], Dict[str, Any]],
|
||||
logger: logging.Logger,
|
||||
) -> None:
|
||||
self._service = service
|
||||
@@ -287,7 +287,7 @@ class ModelListingHandler:
|
||||
)
|
||||
return web.json_response({"error": str(exc)}, status=500)
|
||||
|
||||
def _parse_common_params(self, request: web.Request) -> Dict:
|
||||
def _parse_common_params(self, request: web.Request) -> Dict[str, Any]:
|
||||
page = int(request.query.get("page", "1"))
|
||||
page_size = min(int(request.query.get("page_size", "20")), 100)
|
||||
sort_by = request.query.get("sort_by", "name")
|
||||
@@ -394,12 +394,14 @@ class ModelListingHandler:
|
||||
)
|
||||
|
||||
# View-local-versions filter: show all local versions of a specific model
|
||||
# Accepts either a CivitAI modelId (int) or a HF group key like "hf:user/repo"
|
||||
civitai_model_id = request.query.get("civitai_model_id")
|
||||
if civitai_model_id is not None:
|
||||
try:
|
||||
civitai_model_id = int(civitai_model_id)
|
||||
except (TypeError, ValueError):
|
||||
civitai_model_id = None
|
||||
# Keep as string — could be an HF group key (e.g. "hf:user/repo")
|
||||
pass
|
||||
|
||||
return {
|
||||
"page": page,
|
||||
@@ -537,6 +539,7 @@ class ModelManagementHandler:
|
||||
# Update model_data with new hash
|
||||
model_data["sha256"] = sha256
|
||||
model_data["hash_status"] = "completed"
|
||||
hash_status = "completed"
|
||||
else:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "No SHA256 hash found"}, status=400
|
||||
@@ -544,6 +547,32 @@ class ModelManagementHandler:
|
||||
|
||||
await MetadataManager.hydrate_model_data(model_data)
|
||||
|
||||
# hydrate_model_data replaces model_data with .metadata.json content,
|
||||
# which may lack sha256. Restore from cache and persist the fix.
|
||||
if not model_data.get("sha256"):
|
||||
if sha256:
|
||||
model_data["sha256"] = sha256
|
||||
model_data["hash_status"] = model_data.get("hash_status", hash_status)
|
||||
data_to_save = model_data.copy()
|
||||
data_to_save.pop("folder", None)
|
||||
await MetadataManager.save_metadata(file_path, data_to_save)
|
||||
else:
|
||||
sha256 = await calculate_sha256(file_path)
|
||||
if sha256:
|
||||
model_data["sha256"] = sha256.lower()
|
||||
model_data["hash_status"] = "completed"
|
||||
data_to_save = model_data.copy()
|
||||
data_to_save.pop("folder", None)
|
||||
await MetadataManager.save_metadata(file_path, data_to_save)
|
||||
else:
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": "Failed to compute SHA256 hash for model",
|
||||
},
|
||||
status=500,
|
||||
)
|
||||
|
||||
success, error = await self._metadata_sync.fetch_and_update_model(
|
||||
sha256=model_data["sha256"],
|
||||
file_path=file_path,
|
||||
@@ -566,7 +595,12 @@ class ModelManagementHandler:
|
||||
{"success": False, "error": OFFLINE_FRIENDLY_MESSAGE},
|
||||
status=503,
|
||||
)
|
||||
self._logger.error("Error fetching from CivitAI: %s", exc, exc_info=True)
|
||||
self._logger.error(
|
||||
"Error fetching from CivitAI for %s: %s",
|
||||
locals().get("file_path", "unknown"),
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
||||
|
||||
async def relink_civitai(self, request: web.Request) -> web.Response:
|
||||
@@ -624,7 +658,7 @@ class ModelManagementHandler:
|
||||
try:
|
||||
reader = await request.multipart()
|
||||
|
||||
field = await reader.next()
|
||||
field: Any = await reader.next()
|
||||
if field is None or field.name != "preview_file":
|
||||
raise ValueError("Expected 'preview_file' field")
|
||||
content_type = field.headers.get("Content-Type", "image/png")
|
||||
@@ -666,7 +700,7 @@ class ModelManagementHandler:
|
||||
{
|
||||
"success": True,
|
||||
"preview_url": config.get_preview_static_url(
|
||||
result["preview_path"]
|
||||
str(result["preview_path"])
|
||||
),
|
||||
"preview_nsfw_level": result["preview_nsfw_level"],
|
||||
}
|
||||
@@ -747,7 +781,7 @@ class ModelManagementHandler:
|
||||
|
||||
result = await self._preview_service.replace_preview(
|
||||
model_path=model_path,
|
||||
preview_data=preview_data,
|
||||
preview_data=preview_bytes,
|
||||
content_type=content_type,
|
||||
original_filename=original_filename,
|
||||
nsfw_level=nsfw_level,
|
||||
@@ -759,7 +793,7 @@ class ModelManagementHandler:
|
||||
{
|
||||
"success": True,
|
||||
"preview_url": config.get_preview_static_url(
|
||||
result["preview_path"]
|
||||
str(result["preview_path"])
|
||||
),
|
||||
"preview_nsfw_level": result["preview_nsfw_level"],
|
||||
}
|
||||
@@ -1454,8 +1488,73 @@ class ModelQueryHandler:
|
||||
search = request.query.get("search", "").strip()
|
||||
limit = min(int(request.query.get("limit", "15")), 100)
|
||||
offset = max(0, int(request.query.get("offset", "0")))
|
||||
|
||||
folder = request.query.get("folder")
|
||||
recursive = request.query.get("recursive", "true").lower() == "true"
|
||||
base_models = list(request.query.getall("base_model", []))
|
||||
model_types = list(request.query.getall("model_type", []))
|
||||
|
||||
tag_filters: Dict[str, str] = {}
|
||||
for tag in request.query.getall("tag_include", []):
|
||||
if tag:
|
||||
tag_filters[tag] = "include"
|
||||
for tag in request.query.getall("tag_exclude", []):
|
||||
if tag:
|
||||
tag_filters[tag] = "exclude"
|
||||
|
||||
auto_tag_filters: Dict[str, str] = {}
|
||||
for tag in request.query.getall("auto_tag_include", []):
|
||||
if tag:
|
||||
auto_tag_filters[tag] = "include"
|
||||
for tag in request.query.getall("auto_tag_exclude", []):
|
||||
if tag:
|
||||
auto_tag_filters[tag] = "exclude"
|
||||
|
||||
tag_logic = request.query.get("tag_logic", "any").lower()
|
||||
if tag_logic not in ("any", "all"):
|
||||
tag_logic = "any"
|
||||
|
||||
credit_required = request.query.get("credit_required")
|
||||
if credit_required is not None:
|
||||
credit_required = credit_required.lower() not in ("false", "0", "")
|
||||
|
||||
allow_selling_generated_content = request.query.get(
|
||||
"allow_selling_generated_content"
|
||||
)
|
||||
if allow_selling_generated_content is not None:
|
||||
allow_selling_generated_content = (
|
||||
allow_selling_generated_content.lower() not in ("false", "0", "")
|
||||
)
|
||||
|
||||
# The presence of the recursive param (always sent by the loras
|
||||
# widget when filter mode is on) signals that the filter pipeline
|
||||
# must run even when no concrete filter is set, so global settings
|
||||
# like show_only_sfw stay consistent with the list endpoint.
|
||||
apply_filters = (
|
||||
"recursive" in request.query
|
||||
or folder is not None
|
||||
or bool(base_models)
|
||||
or bool(model_types)
|
||||
or bool(tag_filters)
|
||||
or bool(auto_tag_filters)
|
||||
or credit_required is not None
|
||||
or allow_selling_generated_content is not None
|
||||
)
|
||||
|
||||
matching_paths = await self._service.search_relative_paths(
|
||||
search, limit, offset
|
||||
search,
|
||||
limit,
|
||||
offset,
|
||||
folder=folder,
|
||||
recursive=recursive,
|
||||
base_models=base_models,
|
||||
model_types=model_types,
|
||||
tags=tag_filters,
|
||||
auto_tags=auto_tag_filters,
|
||||
tag_logic=tag_logic,
|
||||
credit_required=credit_required,
|
||||
allow_selling_generated_content=allow_selling_generated_content,
|
||||
apply_filters=apply_filters,
|
||||
)
|
||||
return web.json_response(
|
||||
{"success": True, "relative_paths": matching_paths}
|
||||
@@ -1961,7 +2060,7 @@ class ModelCivitaiHandler:
|
||||
settings_service: SettingsManager,
|
||||
ws_manager: WebSocketManager,
|
||||
logger: logging.Logger,
|
||||
metadata_provider_factory: Callable[[], Awaitable],
|
||||
metadata_provider_factory: Callable[[], Awaitable[Any]],
|
||||
validate_model_type: Callable[[str], bool],
|
||||
expected_model_types: Callable[[], str],
|
||||
find_model_file: Callable[
|
||||
@@ -2026,7 +2125,7 @@ class ModelCivitaiHandler:
|
||||
downloaded_version_ids = set(
|
||||
await history_service.get_downloaded_version_ids(
|
||||
self._service.model_type,
|
||||
model_id,
|
||||
int(model_id),
|
||||
)
|
||||
)
|
||||
except Exception as exc: # pragma: no cover - defensive logging
|
||||
@@ -2303,8 +2402,8 @@ class ModelUpdateHandler:
|
||||
self._logger.error("Failed to fetch license info: %s", exc, exc_info=True)
|
||||
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
||||
|
||||
updated: List[Dict[str, str]] = []
|
||||
errors: List[Dict[str, str]] = []
|
||||
updated: List[Dict[str, Any]] = []
|
||||
errors: List[Dict[str, Any]] = []
|
||||
for model_id in model_ids:
|
||||
license_payload = license_map.get(model_id)
|
||||
if not license_payload:
|
||||
@@ -2317,6 +2416,7 @@ class ModelUpdateHandler:
|
||||
model_section = civitai_section.get("model")
|
||||
if not isinstance(model_section, Mapping):
|
||||
model_section = {}
|
||||
model_section = dict(model_section)
|
||||
model_section.update(resolved_payload)
|
||||
civitai_section["model"] = model_section
|
||||
metadata_payload["civitai"] = civitai_section
|
||||
@@ -2332,7 +2432,7 @@ class ModelUpdateHandler:
|
||||
)
|
||||
errors.append({"filePath": metadata_path, "error": str(exc)})
|
||||
|
||||
response_payload = {"success": True, "updated": updated}
|
||||
response_payload: Dict[str, Any] = {"success": True, "updated": updated}
|
||||
missing_model_ids = [mid for mid in model_ids if mid not in license_map]
|
||||
if missing_model_ids:
|
||||
response_payload["missingModelIds"] = missing_model_ids
|
||||
@@ -2681,6 +2781,7 @@ class ModelUpdateHandler:
|
||||
civitai_payload = metadata_payload.get("civitai")
|
||||
if not isinstance(civitai_payload, Mapping):
|
||||
civitai_payload = {}
|
||||
civitai_payload = dict(civitai_payload)
|
||||
|
||||
model_payload = civitai_payload.get("model")
|
||||
if not isinstance(model_payload, Mapping):
|
||||
@@ -2725,7 +2826,7 @@ class ModelUpdateHandler:
|
||||
|
||||
return aggregated
|
||||
|
||||
def _extract_target_model_ids(self, payload: Dict) -> Optional[List[int]]:
|
||||
def _extract_target_model_ids(self, payload: Dict[str, Any]) -> Optional[List[int]]:
|
||||
if not isinstance(payload, Mapping):
|
||||
return None
|
||||
|
||||
@@ -2753,7 +2854,7 @@ class ModelUpdateHandler:
|
||||
return {}
|
||||
|
||||
to_dict = getattr(metadata, "to_dict", None)
|
||||
if callable(to_dict):
|
||||
if to_dict:
|
||||
try:
|
||||
return to_dict()
|
||||
except Exception:
|
||||
@@ -2764,7 +2865,7 @@ class ModelUpdateHandler:
|
||||
|
||||
return {}
|
||||
|
||||
async def _read_json(self, request: web.Request) -> Dict:
|
||||
async def _read_json(self, request: web.Request) -> Dict[str, Any]:
|
||||
if not request.can_read_body:
|
||||
return {}
|
||||
try:
|
||||
@@ -2796,7 +2897,7 @@ class ModelUpdateHandler:
|
||||
record,
|
||||
*,
|
||||
version_context: Optional[Dict[int, Dict[str, Any]]] = None,
|
||||
) -> Dict:
|
||||
) -> Dict[str, Any]:
|
||||
context = version_context or {}
|
||||
# Check user setting for hiding early access versions
|
||||
hide_early_access = False
|
||||
@@ -2825,7 +2926,7 @@ class ModelUpdateHandler:
|
||||
@staticmethod
|
||||
def _serialize_version(
|
||||
version, context: Optional[Dict[str, Any]]
|
||||
) -> Dict:
|
||||
) -> Dict[str, Any]:
|
||||
context = context or {}
|
||||
preview_override = context.get("preview_override")
|
||||
preview_url = (
|
||||
|
||||
@@ -10,7 +10,7 @@ import asyncio
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional
|
||||
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Tuple
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
@@ -44,6 +44,22 @@ EnsureDependenciesCallable = Callable[[], Awaitable[None]]
|
||||
RecipeScannerGetter = Callable[[], Any]
|
||||
CivitaiClientGetter = Callable[[], Any]
|
||||
|
||||
# Cap concurrent preview-dimension reads across requests. With a cold LRU
|
||||
# cache one page can touch up to page_size image files; 16 balances SSD and
|
||||
# HDD throughput without starving the event loop.
|
||||
_DIMS_READ_SEMAPHORE = asyncio.Semaphore(16)
|
||||
|
||||
|
||||
async def _read_preview_dims(path: str) -> Optional[Tuple[int, int]]:
|
||||
"""Read preview dimensions off the event loop under the concurrency cap.
|
||||
|
||||
PIL I/O runs in a worker thread so it never blocks the event loop, and the
|
||||
semaphore bounds how many files are opened at once even when many list
|
||||
requests land together.
|
||||
"""
|
||||
async with _DIMS_READ_SEMAPHORE:
|
||||
return await asyncio.to_thread(ExifUtils.get_image_dimensions, path)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RecipeHandlerSet:
|
||||
@@ -246,7 +262,8 @@ class RecipeListingHandler:
|
||||
recursive=recursive,
|
||||
)
|
||||
|
||||
for item in result.get("items", []):
|
||||
items = result.get("items", [])
|
||||
for item in items:
|
||||
file_path = item.get("file_path")
|
||||
if file_path:
|
||||
item["file_url"] = self.format_recipe_file_url(file_path)
|
||||
@@ -255,6 +272,26 @@ class RecipeListingHandler:
|
||||
item.setdefault("loras", [])
|
||||
item.setdefault("base_model", "")
|
||||
|
||||
# Batch preview dimension reads with asyncio.gather. The previous
|
||||
# loop awaited asyncio.to_thread once per item, so a page_size=100
|
||||
# request submitted 100 sequential thread calls (50-300ms cold-page
|
||||
# latency). gather runs them concurrently while the semaphore caps
|
||||
# disk opens; dimensions stay omitted (not null) when a preview has
|
||||
# no readable size (video, missing file).
|
||||
to_read = [
|
||||
(i, item.get("file_path"))
|
||||
for i, item in enumerate(items)
|
||||
if item.get("file_path")
|
||||
]
|
||||
if to_read:
|
||||
dims_list = await asyncio.gather(
|
||||
*(_read_preview_dims(path) for _, path in to_read)
|
||||
)
|
||||
for (idx, _), dims in zip(to_read, dims_list):
|
||||
if dims:
|
||||
item = items[idx]
|
||||
item["width"], item["height"] = dims
|
||||
|
||||
return web.json_response(result)
|
||||
except Exception as exc:
|
||||
self._logger.error("Error retrieving recipes: %s", exc, exc_info=True)
|
||||
@@ -1045,10 +1082,10 @@ class RecipeManagementHandler:
|
||||
*,
|
||||
image_url: str,
|
||||
name: str,
|
||||
lora_entries: list,
|
||||
checkpoint_entry: dict,
|
||||
gen_params_request: dict,
|
||||
tags: list,
|
||||
lora_entries: list[Any],
|
||||
checkpoint_entry: Dict[str, Any] | None,
|
||||
gen_params_request: Dict[str, Any] | None,
|
||||
tags: list[Any],
|
||||
base_model: str,
|
||||
source_path: str,
|
||||
) -> web.Response:
|
||||
@@ -1081,6 +1118,12 @@ class RecipeManagementHandler:
|
||||
_original_image_url,
|
||||
) = await self._download_remote_media(image_url)
|
||||
|
||||
# Build a version-cached map of local model hashes to cache items so
|
||||
# CivitaiApiMetadataParser can skip CivitAI API calls for models that
|
||||
# exist on disk. Built once and shared by every parse pass below.
|
||||
local_cache = await recipe_scanner.build_local_hash_cache()
|
||||
from ...recipes.parsers.civitai_image import CivitaiApiMetadataParser
|
||||
|
||||
# Extract embedded EXIF metadata (offloaded to thread pool in this call)
|
||||
embedded_gen_params = {}
|
||||
parsed_embedded = None
|
||||
@@ -1102,9 +1145,16 @@ class RecipeManagementHandler:
|
||||
)
|
||||
)
|
||||
if parser:
|
||||
parsed_embedded = await parser.parse_metadata(
|
||||
raw_embedded, recipe_scanner=recipe_scanner
|
||||
)
|
||||
if isinstance(parser, CivitaiApiMetadataParser):
|
||||
parsed_embedded = await parser.parse_metadata(
|
||||
raw_embedded,
|
||||
recipe_scanner=recipe_scanner,
|
||||
local_cache=local_cache,
|
||||
)
|
||||
else:
|
||||
parsed_embedded = await parser.parse_metadata(
|
||||
raw_embedded, recipe_scanner=recipe_scanner
|
||||
)
|
||||
if parsed_embedded and "gen_params" in parsed_embedded:
|
||||
embedded_gen_params = parsed_embedded["gen_params"]
|
||||
else:
|
||||
@@ -1135,9 +1185,16 @@ class RecipeManagementHandler:
|
||||
civitai_inner_meta
|
||||
)
|
||||
if parser:
|
||||
civitai_parsed = await parser.parse_metadata(
|
||||
civitai_inner_meta, recipe_scanner=recipe_scanner
|
||||
)
|
||||
if isinstance(parser, CivitaiApiMetadataParser):
|
||||
civitai_parsed = await parser.parse_metadata(
|
||||
civitai_inner_meta,
|
||||
recipe_scanner=recipe_scanner,
|
||||
local_cache=local_cache,
|
||||
)
|
||||
else:
|
||||
civitai_parsed = await parser.parse_metadata(
|
||||
civitai_inner_meta, recipe_scanner=recipe_scanner
|
||||
)
|
||||
if civitai_parsed and "gen_params" in civitai_parsed:
|
||||
# Merge: API gen_params override EXIF at field level,
|
||||
# EXIF fills in fields the API doesn't have.
|
||||
@@ -1641,7 +1698,7 @@ class RecipeManagementHandler:
|
||||
if not provider:
|
||||
return ""
|
||||
|
||||
version_info = await provider.get_model_version_info(version_id)
|
||||
version_info = await provider.get_model_version_info(str(version_id))
|
||||
if isinstance(version_info, tuple):
|
||||
version_info = version_info[0]
|
||||
|
||||
@@ -1761,6 +1818,12 @@ class RecipeManagementHandler:
|
||||
await self._download_remote_media(image_url)
|
||||
)
|
||||
|
||||
# Build a version-cached map of local model hashes to cache items so
|
||||
# CivitaiApiMetadataParser can skip CivitAI API calls for models that
|
||||
# exist on disk. Built once and shared by every parse pass below.
|
||||
local_cache = await recipe_scanner.build_local_hash_cache()
|
||||
from ...recipes.parsers.civitai_image import CivitaiApiMetadataParser
|
||||
|
||||
# Extract embedded EXIF metadata
|
||||
embedded_gen_params = {}
|
||||
parsed_embedded = None
|
||||
@@ -1782,9 +1845,16 @@ class RecipeManagementHandler:
|
||||
)
|
||||
)
|
||||
if parser:
|
||||
parsed_embedded = await parser.parse_metadata(
|
||||
raw_embedded, recipe_scanner=recipe_scanner
|
||||
)
|
||||
if isinstance(parser, CivitaiApiMetadataParser):
|
||||
parsed_embedded = await parser.parse_metadata(
|
||||
raw_embedded,
|
||||
recipe_scanner=recipe_scanner,
|
||||
local_cache=local_cache,
|
||||
)
|
||||
else:
|
||||
parsed_embedded = await parser.parse_metadata(
|
||||
raw_embedded, recipe_scanner=recipe_scanner
|
||||
)
|
||||
if parsed_embedded and "gen_params" in parsed_embedded:
|
||||
embedded_gen_params = parsed_embedded["gen_params"]
|
||||
finally:
|
||||
@@ -1822,9 +1892,16 @@ class RecipeManagementHandler:
|
||||
)
|
||||
)
|
||||
if parser:
|
||||
parsed_embedded = await parser.parse_metadata(
|
||||
raw_orig, recipe_scanner=recipe_scanner
|
||||
)
|
||||
if isinstance(parser, CivitaiApiMetadataParser):
|
||||
parsed_embedded = await parser.parse_metadata(
|
||||
raw_orig,
|
||||
recipe_scanner=recipe_scanner,
|
||||
local_cache=local_cache,
|
||||
)
|
||||
else:
|
||||
parsed_embedded = await parser.parse_metadata(
|
||||
raw_orig, recipe_scanner=recipe_scanner
|
||||
)
|
||||
if (
|
||||
parsed_embedded
|
||||
and "gen_params" in parsed_embedded
|
||||
@@ -1858,9 +1935,16 @@ class RecipeManagementHandler:
|
||||
civitai_inner_meta
|
||||
)
|
||||
if parser:
|
||||
civitai_parsed = await parser.parse_metadata(
|
||||
civitai_inner_meta, recipe_scanner=recipe_scanner
|
||||
)
|
||||
if isinstance(parser, CivitaiApiMetadataParser):
|
||||
civitai_parsed = await parser.parse_metadata(
|
||||
civitai_inner_meta,
|
||||
recipe_scanner=recipe_scanner,
|
||||
local_cache=local_cache,
|
||||
)
|
||||
else:
|
||||
civitai_parsed = await parser.parse_metadata(
|
||||
civitai_inner_meta, recipe_scanner=recipe_scanner
|
||||
)
|
||||
if civitai_parsed and "gen_params" in civitai_parsed:
|
||||
# Merge: API gen_params override EXIF at field level,
|
||||
# EXIF fills in fields the API doesn't have.
|
||||
@@ -2072,33 +2156,44 @@ class RecipeManagementHandler:
|
||||
parsed_input = {**image_data, **inner_meta}
|
||||
parsed_input.pop("meta", None)
|
||||
|
||||
# Build a local cache of {hash → cache_item} so the parser can
|
||||
# skip CivitAI API calls for models that exist on disk.
|
||||
local_cache: Dict[str, Dict[str, Any]] = {}
|
||||
lora_scanner = getattr(recipe_scanner, "_lora_scanner", None)
|
||||
if lora_scanner and model_hash:
|
||||
try:
|
||||
parent_cache_data = await lora_scanner.get_cached_data()
|
||||
for item in getattr(parent_cache_data, "raw_data", []):
|
||||
if item.get("sha256", "").lower() == model_hash.lower():
|
||||
local_cache[model_hash.lower()] = item
|
||||
# Compute AutoV3 so the parser can also match on
|
||||
# that hash type (CivitAI metadata resources use
|
||||
# AutoV3).
|
||||
file_path = item.get("file_path")
|
||||
if file_path and os.path.exists(file_path):
|
||||
try:
|
||||
from ...utils.file_utils import (
|
||||
calculate_autov3,
|
||||
)
|
||||
autov3 = calculate_autov3(file_path)
|
||||
if autov3:
|
||||
local_cache[autov3.lower()] = item
|
||||
except Exception:
|
||||
pass
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
# Build the shared local hash cache so the parser can skip CivitAI
|
||||
# API calls for models that exist on disk.
|
||||
local_cache: Dict[str, Dict[str, Any]] = (
|
||||
await recipe_scanner.build_local_hash_cache()
|
||||
)
|
||||
|
||||
# Bounded supplement for un-backfilled parents. The shared builder
|
||||
# never computes autov3; when the parent model exists on disk but
|
||||
# its cached entry has no stored AutoV3, compute it for that single
|
||||
# file and register the AutoV3 key so the parser can also match on
|
||||
# that hash type (CivitAI metadata resources use AutoV3). This runs
|
||||
# whenever the parent is found with an empty autov3, independent of
|
||||
# whether the sha256 key is already present in the shared cache.
|
||||
if model_hash:
|
||||
lora_scanner = getattr(recipe_scanner, "_lora_scanner", None)
|
||||
if lora_scanner:
|
||||
try:
|
||||
parent_cache_data = await lora_scanner.get_cached_data()
|
||||
for item in getattr(parent_cache_data, "raw_data", []):
|
||||
if item.get("sha256", "").lower() == model_hash.lower():
|
||||
autov3 = (item.get("autov3") or "").lower()
|
||||
if not autov3:
|
||||
file_path = item.get("file_path")
|
||||
if file_path and os.path.exists(file_path):
|
||||
try:
|
||||
from ...utils.file_utils import (
|
||||
calculate_autov3,
|
||||
)
|
||||
autov3 = (
|
||||
calculate_autov3(file_path) or ""
|
||||
).lower()
|
||||
except Exception:
|
||||
pass
|
||||
if autov3:
|
||||
local_cache[autov3] = item
|
||||
break
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
parser = self._analysis_service._recipe_parser_factory.create_parser(
|
||||
parsed_input
|
||||
@@ -2130,10 +2225,10 @@ class RecipeManagementHandler:
|
||||
parent_model_id: int | None = None
|
||||
parent_version_name: str | None = None
|
||||
parent_model_name: str | None = None
|
||||
# Prefer sha256 key; fall back to any cached entry.
|
||||
# Resolve the parent strictly by its sha256 key. There is no
|
||||
# arbitrary fallback: with a full-library cache, picking any entry
|
||||
# would corrupt the isDeleted reconciliation below.
|
||||
parent_item = local_cache.get(model_hash.lower()) if model_hash else None
|
||||
if parent_item is None and local_cache:
|
||||
parent_item = next(iter(local_cache.values()))
|
||||
if parent_item:
|
||||
civ = parent_item.get("civitai") or {}
|
||||
if isinstance(civ, dict):
|
||||
@@ -2349,7 +2444,7 @@ class RecipeAnalysisHandler:
|
||||
content_type = request.headers.get("Content-Type", "")
|
||||
if "multipart/form-data" in content_type:
|
||||
reader = await request.multipart()
|
||||
field = await reader.next()
|
||||
field: Any = await reader.next()
|
||||
if field is None or field.name != "image":
|
||||
raise RecipeValidationError("No image field found")
|
||||
image_chunks = bytearray()
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from aiohttp import web
|
||||
from typing import Dict
|
||||
from server import PromptServer # type: ignore
|
||||
from typing import Any, Dict
|
||||
from server import PromptServer # pyright: ignore[reportMissingImports]
|
||||
|
||||
from .base_model_routes import BaseModelRoutes
|
||||
from .model_route_registrar import ModelRouteRegistrar
|
||||
@@ -31,13 +31,13 @@ class LoraRoutes(BaseModelRoutes):
|
||||
# Attach service dependencies
|
||||
self.attach_service(self.service)
|
||||
|
||||
def setup_routes(self, app: web.Application):
|
||||
def setup_routes(self, app: web.Application, prefix: str = "loras"):
|
||||
"""Setup LoRA routes"""
|
||||
# Schedule service initialization on app startup
|
||||
app.on_startup.append(lambda _: self.initialize_services())
|
||||
|
||||
# Setup common routes with 'loras' prefix (includes page route)
|
||||
super().setup_routes(app, "loras")
|
||||
super().setup_routes(app, prefix)
|
||||
|
||||
def setup_specific_routes(self, registrar: ModelRouteRegistrar, prefix: str):
|
||||
"""Setup LoRA-specific routes"""
|
||||
@@ -73,7 +73,7 @@ class LoraRoutes(BaseModelRoutes):
|
||||
"POST", "/api/lm/{prefix}/get_trigger_words", prefix, self.get_trigger_words
|
||||
)
|
||||
|
||||
def _parse_specific_params(self, request: web.Request) -> Dict:
|
||||
def _parse_specific_params(self, request: web.Request) -> Dict[str, Any]:
|
||||
"""Parse LoRA-specific parameters"""
|
||||
params = {}
|
||||
|
||||
@@ -119,25 +119,6 @@ class LoraRoutes(BaseModelRoutes):
|
||||
logger.error(f"Error getting letter counts: {e}")
|
||||
return web.json_response({"success": False, "error": str(e)}, status=500)
|
||||
|
||||
async def get_lora_notes(self, request: web.Request) -> web.Response:
|
||||
"""Get notes for a specific LoRA file"""
|
||||
try:
|
||||
lora_name = request.query.get("name")
|
||||
if not lora_name:
|
||||
return web.Response(text="Lora file name is required", status=400)
|
||||
|
||||
notes = await self.service.get_lora_notes(lora_name)
|
||||
if notes is not None:
|
||||
return web.json_response({"success": True, "notes": notes})
|
||||
else:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "LoRA not found in cache"}, status=404
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting lora notes: {e}", exc_info=True)
|
||||
return web.json_response({"success": False, "error": str(e)}, status=500)
|
||||
|
||||
async def get_lora_trigger_words(self, request: web.Request) -> web.Response:
|
||||
"""Get trigger words for a specific LoRA file"""
|
||||
try:
|
||||
@@ -168,52 +149,6 @@ class LoraRoutes(BaseModelRoutes):
|
||||
logger.error(f"Error getting lora usage tips by path: {e}", exc_info=True)
|
||||
return web.json_response({"success": False, "error": str(e)}, status=500)
|
||||
|
||||
async def get_lora_preview_url(self, request: web.Request) -> web.Response:
|
||||
"""Get the static preview URL for a LoRA file"""
|
||||
try:
|
||||
lora_name = request.query.get("name")
|
||||
if not lora_name:
|
||||
return web.Response(text="Lora file name is required", status=400)
|
||||
|
||||
preview_url = await self.service.get_lora_preview_url(lora_name)
|
||||
if preview_url:
|
||||
return web.json_response({"success": True, "preview_url": preview_url})
|
||||
else:
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": "No preview URL found for the specified lora",
|
||||
},
|
||||
status=404,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting lora preview URL: {e}", exc_info=True)
|
||||
return web.json_response({"success": False, "error": str(e)}, status=500)
|
||||
|
||||
async def get_lora_civitai_url(self, request: web.Request) -> web.Response:
|
||||
"""Get the Civitai URL for a LoRA file"""
|
||||
try:
|
||||
lora_name = request.query.get("name")
|
||||
if not lora_name:
|
||||
return web.Response(text="Lora file name is required", status=400)
|
||||
|
||||
result = await self.service.get_lora_civitai_url(lora_name)
|
||||
if result["civitai_url"]:
|
||||
return web.json_response({"success": True, **result})
|
||||
else:
|
||||
return web.json_response(
|
||||
{
|
||||
"success": False,
|
||||
"error": "No Civitai data found for the specified lora",
|
||||
},
|
||||
status=404,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error getting lora Civitai URL: {e}", exc_info=True)
|
||||
return web.json_response({"success": False, "error": str(e)}, status=500)
|
||||
|
||||
async def get_random_loras(self, request: web.Request) -> web.Response:
|
||||
"""Get random LoRAs based on filters and strength ranges"""
|
||||
try:
|
||||
@@ -337,7 +272,7 @@ class LoraRoutes(BaseModelRoutes):
|
||||
graph_identifier = entry.get("graph_id")
|
||||
|
||||
try:
|
||||
parsed_node_id = int(node_identifier)
|
||||
parsed_node_id = int(node_identifier) # pyright: ignore[reportArgumentType]
|
||||
except (TypeError, ValueError):
|
||||
parsed_node_id = node_identifier
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@ miscellaneous endpoints share a consistent registration flow.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Iterable, Mapping
|
||||
from typing import Any, Callable, Iterable, Mapping
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
@@ -147,7 +147,7 @@ class MiscRouteRegistrar:
|
||||
handler_lookup[definition.handler_name],
|
||||
)
|
||||
|
||||
def _bind(self, method: str, path: str, handler: Callable) -> None:
|
||||
def _bind(self, method: str, path: str, handler: Callable[..., Any]) -> None:
|
||||
add_method_name = self._METHOD_MAP[method.upper()]
|
||||
add_method = getattr(self._app.router, add_method_name)
|
||||
add_method(path, handler)
|
||||
|
||||
@@ -7,7 +7,7 @@ import os
|
||||
from typing import Awaitable, Callable, Mapping
|
||||
|
||||
from aiohttp import web
|
||||
from server import PromptServer # type: ignore
|
||||
from server import PromptServer # pyright: ignore[reportMissingImports]
|
||||
|
||||
from ..services.metadata_service import (
|
||||
get_metadata_archive_manager,
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Iterable, Mapping
|
||||
from typing import Any, Callable, Iterable, Mapping
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
@@ -174,15 +174,15 @@ class ModelRouteRegistrar:
|
||||
handler_lookup[definition.handler_name],
|
||||
)
|
||||
|
||||
def add_route(self, method: str, path: str, handler: Callable) -> None:
|
||||
def add_route(self, method: str, path: str, handler: Callable[..., Any]) -> None:
|
||||
self._bind_route(method, path, handler)
|
||||
|
||||
def add_prefixed_route(
|
||||
self, method: str, path_template: str, prefix: str, handler: Callable
|
||||
self, method: str, path_template: str, prefix: str, handler: Callable[..., Any]
|
||||
) -> None:
|
||||
self._bind_route(method, path_template.replace("{prefix}", prefix), handler)
|
||||
|
||||
def _bind_route(self, method: str, path: str, handler: Callable) -> None:
|
||||
def _bind_route(self, method: str, path: str, handler: Callable[..., Any]) -> None:
|
||||
add_method_name = self._METHOD_MAP[method.upper()]
|
||||
add_method = getattr(self._app.router, add_method_name)
|
||||
add_method(path, handler)
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Mapping
|
||||
from typing import Any, Callable, Mapping
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
@@ -105,7 +105,7 @@ class RecipeRouteRegistrar:
|
||||
handler = handler_lookup[definition.handler_name]
|
||||
self._bind_route(definition.method, definition.path, handler)
|
||||
|
||||
def _bind_route(self, method: str, path: str, handler: Callable) -> None:
|
||||
def _bind_route(self, method: str, path: str, handler: Callable[..., Any]) -> None:
|
||||
add_method_name = self._METHOD_MAP[method.upper()]
|
||||
add_method = getattr(self._app.router, add_method_name)
|
||||
add_method(path, handler)
|
||||
|
||||
+11
-10
@@ -40,10 +40,11 @@ class StatsRoutes:
|
||||
"""Route handlers for Statistics page and API endpoints"""
|
||||
|
||||
def __init__(self):
|
||||
self.lora_scanner = None
|
||||
self.checkpoint_scanner = None
|
||||
self.embedding_scanner = None
|
||||
self.usage_stats = None
|
||||
self.lora_scanner: Any = None
|
||||
self.checkpoint_scanner: Any = None
|
||||
self.embedding_scanner: Any = None
|
||||
self.usage_stats: Any = None
|
||||
self._i18n_filter_added = False
|
||||
self.template_env = jinja2.Environment(
|
||||
loader=jinja2.FileSystemLoader(config.templates_path),
|
||||
autoescape=True
|
||||
@@ -95,9 +96,9 @@ class StatsRoutes:
|
||||
server_i18n.set_locale(user_language)
|
||||
|
||||
# 为模板环境添加i18n过滤器
|
||||
if not hasattr(self.template_env, '_i18n_filter_added'):
|
||||
if not self._i18n_filter_added:
|
||||
self.template_env.filters['t'] = server_i18n.create_template_filter()
|
||||
self.template_env._i18n_filter_added = True
|
||||
self._i18n_filter_added = True
|
||||
|
||||
template = self.template_env.get_template('statistics.html')
|
||||
rendered = template.render(
|
||||
@@ -549,7 +550,7 @@ class StatsRoutes:
|
||||
'error': str(e)
|
||||
}, status=500)
|
||||
|
||||
def _count_unused_models(self, models: List[Dict], usage_data: Dict) -> int:
|
||||
def _count_unused_models(self, models: List[Dict[str, Any]], usage_data: Dict[str, Any]) -> int:
|
||||
"""Count models that have never been used"""
|
||||
used_hashes = set(usage_data.keys())
|
||||
unused_count = 0
|
||||
@@ -560,7 +561,7 @@ class StatsRoutes:
|
||||
|
||||
return unused_count
|
||||
|
||||
def _get_top_used_models(self, usage_data: Dict, model_map: Dict, limit: int) -> List[Dict]:
|
||||
def _get_top_used_models(self, usage_data: Dict[str, Any], model_map: Dict[str, Any], limit: int) -> List[Dict[str, Any]]:
|
||||
"""Get top used models with their metadata"""
|
||||
sorted_usage = sorted(usage_data.items(), key=lambda x: x[1].get('total', 0), reverse=True)
|
||||
|
||||
@@ -578,7 +579,7 @@ class StatsRoutes:
|
||||
|
||||
return top_models
|
||||
|
||||
def _get_usage_timeline(self, usage_data: Dict, days: int) -> List[Dict]:
|
||||
def _get_usage_timeline(self, usage_data: Dict[str, Any], days: int) -> List[Dict[str, Any]]:
|
||||
"""Get usage timeline for the past N days"""
|
||||
timeline = []
|
||||
today = datetime.now()
|
||||
@@ -614,7 +615,7 @@ class StatsRoutes:
|
||||
|
||||
return list(reversed(timeline)) # Oldest to newest
|
||||
|
||||
def _format_size(self, size_bytes: int) -> str:
|
||||
def _format_size(self, size_bytes: float) -> str:
|
||||
"""Format file size in human readable format"""
|
||||
for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
|
||||
if size_bytes < 1024.0:
|
||||
|
||||
+324
-53
@@ -6,7 +6,7 @@ import shutil
|
||||
import tempfile
|
||||
import asyncio
|
||||
from aiohttp import web, ClientError
|
||||
from typing import Dict, List
|
||||
from typing import Any, Dict, List, cast
|
||||
|
||||
from ..utils.settings_paths import ensure_settings_file
|
||||
from ..services.downloader import get_downloader
|
||||
@@ -38,6 +38,84 @@ def _clean_excludes() -> List[str]:
|
||||
return excludes
|
||||
|
||||
|
||||
def _stage_preserved_items(plugin_root: str) -> tuple[str, list[str]]:
|
||||
"""Move preserved user-data items to a temp directory outside *plugin_root*.
|
||||
|
||||
This ensures that ``git reset --hard``, ``git clean -fd``, and ZIP-based
|
||||
replacement cannot touch these files even when ``-e`` exclusion patterns
|
||||
are mishandled (e.g. on Windows where forward-slash patterns may not
|
||||
match backslash-prefixed paths in some Git builds, or where file locks
|
||||
prevent deletion/recreation).
|
||||
|
||||
Returns:
|
||||
``(backup_root, staged_names)``: the temp directory path and the
|
||||
list of item names that were successfully moved.
|
||||
"""
|
||||
backup_root = tempfile.mkdtemp(prefix='lora_manager_update_')
|
||||
staged: list[str] = []
|
||||
for name in _PRESERVE_DIRS:
|
||||
src = os.path.join(plugin_root, name)
|
||||
if not os.path.lexists(src):
|
||||
continue
|
||||
dst = os.path.join(backup_root, name)
|
||||
try:
|
||||
shutil.move(src, dst)
|
||||
staged.append(name)
|
||||
logger.debug("Staged '%s' for update safety", name)
|
||||
except OSError:
|
||||
# ``shutil.move`` may fail on Windows if a file handle inside
|
||||
# the directory is still open (e.g. a SQLite WAL file). Fall
|
||||
# back to copy-then-remove.
|
||||
logger.debug("Move failed for '%s', falling back to copy", name)
|
||||
try:
|
||||
if os.path.isdir(src) and not os.path.islink(src):
|
||||
shutil.copytree(src, dst, symlinks=True)
|
||||
shutil.rmtree(src, ignore_errors=True)
|
||||
else:
|
||||
shutil.copy2(src, dst)
|
||||
os.remove(src)
|
||||
staged.append(name)
|
||||
logger.info("Copied (then removed) '%s' for update safety", name)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Could not stage '%s': %s (will rely on git -e / skip lists)", name, exc
|
||||
)
|
||||
return backup_root, staged
|
||||
|
||||
|
||||
def _restore_preserved_items(plugin_root: str, backup_root: str, staged: list[str]) -> None:
|
||||
"""Move staged items back from *backup_root* into *plugin_root*.
|
||||
|
||||
Any leftover placeholder at the destination (created by git checkout or
|
||||
ZIP extraction) is removed before the move.
|
||||
"""
|
||||
for name in staged:
|
||||
src = os.path.join(backup_root, name)
|
||||
dst = os.path.join(plugin_root, name)
|
||||
try:
|
||||
if os.path.lexists(dst):
|
||||
if os.path.isdir(dst) and not os.path.islink(dst):
|
||||
shutil.rmtree(dst, ignore_errors=True)
|
||||
else:
|
||||
os.remove(dst)
|
||||
shutil.move(src, dst)
|
||||
logger.debug("Restored '%s' after update", name)
|
||||
except OSError:
|
||||
logger.debug("Move failed restoring '%s', falling back to copy", name)
|
||||
try:
|
||||
if os.path.isdir(src) and not os.path.islink(src):
|
||||
shutil.copytree(src, dst, symlinks=True, dirs_exist_ok=True)
|
||||
shutil.rmtree(src, ignore_errors=True)
|
||||
else:
|
||||
shutil.copy2(src, dst)
|
||||
os.remove(src)
|
||||
logger.info("Copied '%s' back after update", name)
|
||||
except Exception as exc:
|
||||
logger.error("Failed to restore '%s': %s", name, exc)
|
||||
shutil.rmtree(backup_root, ignore_errors=True)
|
||||
|
||||
|
||||
|
||||
class UpdateRoutes:
|
||||
"""Routes for handling plugin update checks"""
|
||||
|
||||
@@ -47,6 +125,7 @@ class UpdateRoutes:
|
||||
app.router.add_get('/api/lm/check-updates', UpdateRoutes.check_updates)
|
||||
app.router.add_get('/api/lm/version-info', UpdateRoutes.get_version_info)
|
||||
app.router.add_post('/api/lm/perform-update', UpdateRoutes.perform_update)
|
||||
app.router.add_post('/api/lm/switch-channel', UpdateRoutes.switch_channel)
|
||||
|
||||
@staticmethod
|
||||
async def check_updates(request):
|
||||
@@ -65,10 +144,17 @@ class UpdateRoutes:
|
||||
|
||||
# Fetch remote version from GitHub
|
||||
if nightly:
|
||||
remote_version, changelog = await UpdateRoutes._get_nightly_version()
|
||||
releases = None
|
||||
local_hash = git_info.get('short_hash', '')
|
||||
nightly_version, releases_result = await asyncio.gather(
|
||||
UpdateRoutes._get_nightly_version(local_hash),
|
||||
UpdateRoutes._get_remote_version()
|
||||
)
|
||||
remote_version, _, behind_by, commit_date = nightly_version
|
||||
_, changelog, releases = releases_result
|
||||
else:
|
||||
remote_version, changelog, releases = await UpdateRoutes._get_remote_version()
|
||||
behind_by = 0
|
||||
commit_date = ''
|
||||
|
||||
# Compare versions
|
||||
if nightly:
|
||||
@@ -81,6 +167,10 @@ class UpdateRoutes:
|
||||
remote_version.replace('v', '')
|
||||
)
|
||||
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
plugin_root = os.path.dirname(os.path.dirname(current_dir))
|
||||
has_git = os.path.exists(os.path.join(plugin_root, '.git'))
|
||||
|
||||
response_data = {
|
||||
'success': True,
|
||||
'current_version': local_version,
|
||||
@@ -88,13 +178,13 @@ class UpdateRoutes:
|
||||
'update_available': update_available,
|
||||
'changelog': changelog,
|
||||
'git_info': git_info,
|
||||
'nightly': nightly
|
||||
'nightly': nightly,
|
||||
'has_git': has_git,
|
||||
'releases': releases,
|
||||
'behind_by': behind_by,
|
||||
'commit_date': commit_date
|
||||
}
|
||||
|
||||
# Include releases list for stable mode
|
||||
if releases is not None:
|
||||
response_data['releases'] = releases
|
||||
|
||||
return web.json_response(response_data)
|
||||
|
||||
except NETWORK_EXCEPTIONS as e:
|
||||
@@ -126,9 +216,14 @@ class UpdateRoutes:
|
||||
# Format: version-short_hash
|
||||
version_string = f"{local_version}-{short_hash}"
|
||||
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
plugin_root = os.path.dirname(os.path.dirname(current_dir))
|
||||
has_git = os.path.exists(os.path.join(plugin_root, '.git'))
|
||||
|
||||
return web.json_response({
|
||||
'success': True,
|
||||
'version': version_string
|
||||
'version': version_string,
|
||||
'has_git': has_git
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
@@ -156,20 +251,22 @@ class UpdateRoutes:
|
||||
if os.path.exists(settings_path):
|
||||
with open(settings_path, 'r', encoding='utf-8') as f:
|
||||
settings_backup = f.read()
|
||||
logger.info("Backed up settings.json")
|
||||
logger.debug("Backed up settings.json (%d bytes)", len(settings_backup))
|
||||
|
||||
git_folder = os.path.join(plugin_root, '.git')
|
||||
if os.path.exists(git_folder):
|
||||
# Git update
|
||||
success, new_version = await UpdateRoutes._perform_git_update(plugin_root, nightly)
|
||||
else:
|
||||
# Fallback: Download ZIP and replace files
|
||||
success, new_version = await UpdateRoutes._download_and_replace_zip(plugin_root)
|
||||
staged_backup_dir, staged_items = _stage_preserved_items(plugin_root)
|
||||
try:
|
||||
git_folder = os.path.join(plugin_root, '.git')
|
||||
if os.path.exists(git_folder):
|
||||
success, new_version = await UpdateRoutes._perform_git_update(plugin_root, nightly)
|
||||
else:
|
||||
success, new_version = await UpdateRoutes._download_and_replace_zip(plugin_root)
|
||||
finally:
|
||||
_restore_preserved_items(plugin_root, staged_backup_dir, staged_items)
|
||||
|
||||
if settings_backup and success:
|
||||
with open(settings_path, 'w', encoding='utf-8') as f:
|
||||
f.write(settings_backup)
|
||||
logger.info("Restored settings.json")
|
||||
logger.debug("Restored settings.json content (%d bytes)", len(settings_backup))
|
||||
|
||||
if success:
|
||||
return web.json_response({
|
||||
@@ -190,6 +287,164 @@ class UpdateRoutes:
|
||||
'error': str(e)
|
||||
})
|
||||
|
||||
@staticmethod
|
||||
async def switch_channel(request):
|
||||
"""
|
||||
Switch between release and nightly update channels.
|
||||
|
||||
ZIP/CNR install → Nightly: git init + checkout main (one-way upgrade)
|
||||
Git install → Release: git checkout latest tag (.git preserved)
|
||||
ZIP/CNR install → Release: ZIP download (no .git, stays in ZIP mode)
|
||||
Git install → Nightly: git checkout main + pull
|
||||
"""
|
||||
try:
|
||||
body = await request.json() if request.has_body else {}
|
||||
channel = body.get('channel', '')
|
||||
|
||||
if channel not in ('release', 'nightly'):
|
||||
return web.json_response({
|
||||
'success': False,
|
||||
'error': f'Invalid channel: {channel}. Must be "release" or "nightly".'
|
||||
})
|
||||
|
||||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
plugin_root = os.path.dirname(os.path.dirname(current_dir))
|
||||
|
||||
settings_path = ensure_settings_file(logger)
|
||||
settings_backup = None
|
||||
if os.path.exists(settings_path):
|
||||
with open(settings_path, 'r', encoding='utf-8') as f:
|
||||
settings_backup = f.read()
|
||||
logger.debug("Backed up settings.json before channel switch (%d bytes)", len(settings_backup))
|
||||
|
||||
staged_backup_dir, staged_items = _stage_preserved_items(plugin_root)
|
||||
try:
|
||||
git_folder = os.path.join(plugin_root, '.git')
|
||||
|
||||
if channel == 'nightly':
|
||||
git_backup = None
|
||||
if os.path.exists(git_folder):
|
||||
git_backup = UpdateRoutes._backup_git(git_folder, 'nightly')
|
||||
|
||||
success = False
|
||||
new_version = ''
|
||||
try:
|
||||
if os.path.exists(git_folder):
|
||||
success, new_version = await UpdateRoutes._perform_git_update(
|
||||
plugin_root, nightly=True
|
||||
)
|
||||
else:
|
||||
success, new_version = UpdateRoutes._init_git_repo(plugin_root)
|
||||
finally:
|
||||
UpdateRoutes._restore_git(git_backup, git_folder, success, 'nightly')
|
||||
else:
|
||||
success = False
|
||||
new_version = ''
|
||||
if os.path.exists(git_folder):
|
||||
success, new_version = await UpdateRoutes._perform_git_update(
|
||||
plugin_root, nightly=False
|
||||
)
|
||||
else:
|
||||
tracking_file = os.path.join(plugin_root, '.tracking')
|
||||
if os.path.exists(tracking_file):
|
||||
os.remove(tracking_file)
|
||||
success, new_version = await UpdateRoutes._download_and_replace_zip(plugin_root)
|
||||
finally:
|
||||
_restore_preserved_items(plugin_root, staged_backup_dir, staged_items)
|
||||
|
||||
if settings_backup and success:
|
||||
with open(settings_path, 'w', encoding='utf-8') as f:
|
||||
f.write(settings_backup)
|
||||
logger.debug("Restored settings.json content after channel switch (%d bytes)", len(settings_backup))
|
||||
|
||||
if success:
|
||||
return web.json_response({
|
||||
'success': True,
|
||||
'channel': channel,
|
||||
'new_version': new_version,
|
||||
'message': f'Switched to {channel} channel'
|
||||
})
|
||||
else:
|
||||
return web.json_response({
|
||||
'success': False,
|
||||
'error': f'Failed to switch to {channel} channel'
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to switch channel: %s", e, exc_info=True)
|
||||
return web.json_response({
|
||||
'success': False,
|
||||
'error': str(e)
|
||||
})
|
||||
|
||||
@staticmethod
|
||||
def _init_git_repo(plugin_root: str) -> tuple[bool, str]:
|
||||
"""
|
||||
Initialize a Git repository in a ZIP-installed plugin folder.
|
||||
Clones the remote history and checks out main branch.
|
||||
"""
|
||||
try:
|
||||
import git
|
||||
except ImportError:
|
||||
logger.error(
|
||||
"GitPython is not available: cannot initialize git repo. "
|
||||
"Install git or set $GIT_PYTHON_GIT_EXECUTABLE to the git binary path."
|
||||
)
|
||||
return False, ""
|
||||
|
||||
clean_excludes = _clean_excludes()
|
||||
|
||||
try:
|
||||
repo = git.Repo.init(plugin_root)
|
||||
origin = repo.create_remote(
|
||||
'origin',
|
||||
'https://github.com/willmiao/ComfyUI-Lora-Manager.git'
|
||||
)
|
||||
origin.fetch()
|
||||
|
||||
repo.create_head('main', origin.refs.main)
|
||||
repo.git.checkout('main', '--force')
|
||||
repo.git.reset('--hard')
|
||||
repo.git.clean('-fd', *clean_excludes)
|
||||
|
||||
tracking_file = os.path.join(plugin_root, '.tracking')
|
||||
if os.path.exists(tracking_file):
|
||||
os.remove(tracking_file)
|
||||
logger.info("Removed .tracking file (now in git mode)")
|
||||
|
||||
new_version = f"main-{repo.head.commit.hexsha[:7]}"
|
||||
logger.info("Initialized git repo on main branch: %s", new_version)
|
||||
return True, new_version
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to initialize git repo: %s", e, exc_info=True)
|
||||
return False, ""
|
||||
|
||||
@staticmethod
|
||||
def _backup_git(git_folder, label):
|
||||
try:
|
||||
backup_dir = tempfile.mkdtemp()
|
||||
backup = os.path.join(backup_dir, '.git')
|
||||
shutil.copytree(git_folder, backup)
|
||||
logger.info("Backed up .git before switching to %s", label)
|
||||
return backup
|
||||
except Exception as e:
|
||||
logger.error("Failed to backup .git before %s switch: %s", label, e)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _restore_git(git_backup, git_folder, success, label):
|
||||
if git_backup and not success:
|
||||
try:
|
||||
if os.path.exists(git_folder):
|
||||
shutil.rmtree(git_folder)
|
||||
shutil.copytree(git_backup, git_folder)
|
||||
logger.info("Restored .git after failed %s switch", label)
|
||||
except Exception as e:
|
||||
logger.error("Failed to restore .git after %s switch: %s", label, e)
|
||||
if git_backup:
|
||||
shutil.rmtree(os.path.dirname(git_backup), ignore_errors=True)
|
||||
|
||||
@staticmethod
|
||||
async def _download_and_replace_zip(plugin_root: str) -> tuple[bool, str]:
|
||||
"""
|
||||
@@ -212,9 +467,10 @@ class UpdateRoutes:
|
||||
if not success:
|
||||
logger.error(f"Failed to fetch release info: {data}")
|
||||
return False, ""
|
||||
|
||||
zip_url = data.get("zipball_url")
|
||||
version = data.get("tag_name", "unknown")
|
||||
|
||||
release_payload = cast(dict[str, Any], data)
|
||||
zip_url = release_payload.get("zipball_url", "")
|
||||
version = release_payload.get("tag_name", "unknown")
|
||||
|
||||
# Download ZIP to temporary file
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".zip") as tmp_zip:
|
||||
@@ -244,8 +500,7 @@ class UpdateRoutes:
|
||||
except Exception:
|
||||
logger.debug("Could not close downloaded-version history database", exc_info=True)
|
||||
|
||||
# Skip settings.json, civitai, model cache and runtime cache folders
|
||||
UpdateRoutes._clean_plugin_folder(plugin_root, skip_files=['settings.json', 'civitai', 'model_cache', 'cache', 'wildcards', 'backups', 'stats'])
|
||||
UpdateRoutes._clean_plugin_folder(plugin_root, skip_files=list(_PRESERVE_DIRS))
|
||||
|
||||
# Extract ZIP to temp dir
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
@@ -255,7 +510,7 @@ class UpdateRoutes:
|
||||
extracted_root = next(os.scandir(tmp_dir)).path
|
||||
|
||||
# Copy files, skipping user data that should be preserved
|
||||
skip_items = {'settings.json', 'civitai', 'wildcards', 'backups', 'stats'}
|
||||
skip_items = set(_PRESERVE_DIRS)
|
||||
for item in os.listdir(extracted_root):
|
||||
if item in skip_items:
|
||||
continue
|
||||
@@ -272,7 +527,7 @@ class UpdateRoutes:
|
||||
# for ComfyUI Manager to work properly
|
||||
tracking_info_file = os.path.join(plugin_root, '.tracking')
|
||||
tracking_files = []
|
||||
skip_tracked = {'civitai', 'wildcards', 'backups', 'stats'}
|
||||
skip_tracked = set(_PRESERVE_DIRS) - {'settings.json'}
|
||||
for root, dirs, files in os.walk(extracted_root):
|
||||
# Skip user data directories and their contents
|
||||
rel_root = os.path.relpath(root, extracted_root)
|
||||
@@ -295,7 +550,8 @@ class UpdateRoutes:
|
||||
except Exception as e:
|
||||
logger.error(f"ZIP update failed: {e}", exc_info=True)
|
||||
return False, ""
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _clean_plugin_folder(plugin_root, skip_files=None):
|
||||
skip_files = skip_files or []
|
||||
for item in os.listdir(plugin_root):
|
||||
@@ -308,41 +564,56 @@ class UpdateRoutes:
|
||||
os.remove(path)
|
||||
|
||||
@staticmethod
|
||||
async def _get_nightly_version() -> tuple[str, List[str]]:
|
||||
"""
|
||||
Fetch latest commit from main branch
|
||||
"""
|
||||
async def _get_nightly_version(local_hash: str = "") -> tuple[str, List[str], int, str]:
|
||||
repo_owner = "willmiao"
|
||||
repo_name = "ComfyUI-Lora-Manager"
|
||||
|
||||
# Use GitHub API to fetch the latest commit from main branch
|
||||
|
||||
github_url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/commits/main"
|
||||
|
||||
|
||||
try:
|
||||
downloader = await get_downloader()
|
||||
success, data = await downloader.make_request('GET', github_url, custom_headers={'Accept': 'application/vnd.github+json'})
|
||||
|
||||
success, data = await downloader.make_request(
|
||||
'GET', github_url,
|
||||
custom_headers={'Accept': 'application/vnd.github+json'}
|
||||
)
|
||||
|
||||
if not success:
|
||||
logger.warning(f"Failed to fetch GitHub commit: {data}")
|
||||
return "main", []
|
||||
|
||||
commit_sha = data.get('sha', '')[:7] # Short hash
|
||||
commit_message = data.get('commit', {}).get('message', '')
|
||||
|
||||
# Format as "main-{short_hash}"
|
||||
logger.warning("Failed to fetch GitHub commit: %s", data)
|
||||
return "main", [], 0, ""
|
||||
|
||||
commit_payload = cast(dict[str, Any], data)
|
||||
commit_sha = commit_payload.get('sha', '')[:7]
|
||||
commit_message = commit_payload.get('commit', {}).get('message', '')
|
||||
commit_date = commit_payload.get('commit', {}).get('committer', {}).get('date', '')[:10]
|
||||
|
||||
version = f"main-{commit_sha}"
|
||||
|
||||
# Use commit message as changelog
|
||||
changelog = [commit_message] if commit_message else []
|
||||
|
||||
return version, changelog
|
||||
|
||||
|
||||
behind_by = 0
|
||||
if local_hash and local_hash not in ('unknown', 'stable'):
|
||||
compare_url = (
|
||||
f"https://api.github.com/repos/{repo_owner}/{repo_name}"
|
||||
f"/compare/{local_hash}...main"
|
||||
)
|
||||
c_ok, c_data = await downloader.make_request(
|
||||
'GET', compare_url,
|
||||
custom_headers={'Accept': 'application/vnd.github+json'}
|
||||
)
|
||||
if c_ok:
|
||||
compare_payload = cast(dict[str, Any], c_data)
|
||||
if compare_payload.get('status') in ('ahead', 'diverged'):
|
||||
behind_by = compare_payload.get('ahead_by', 0)
|
||||
else:
|
||||
behind_by = compare_payload.get('behind_by', 0)
|
||||
|
||||
return version, changelog, behind_by, commit_date
|
||||
|
||||
except NETWORK_EXCEPTIONS as e:
|
||||
logger.warning("Unable to reach GitHub for nightly version: %s", e)
|
||||
return "main", []
|
||||
return "main", [], 0, ""
|
||||
except Exception as e:
|
||||
logger.error(f"Error fetching nightly version: {e}", exc_info=True)
|
||||
return "main", []
|
||||
logger.error("Error fetching nightly version: %s", e, exc_info=True)
|
||||
return "main", [], 0, ""
|
||||
|
||||
@staticmethod
|
||||
def _compare_nightly_versions(local_git_info: Dict[str, str], remote_version: str) -> bool:
|
||||
@@ -438,7 +709,7 @@ class UpdateRoutes:
|
||||
logger.info(f"Successfully updated to {new_version}")
|
||||
return True, new_version
|
||||
|
||||
except git.exc.GitError as e:
|
||||
except git.exc.GitError as e: # pyright: ignore[reportAttributeAccessIssue]
|
||||
logger.error(f"Git error during update: {e}")
|
||||
return False, ""
|
||||
except Exception as e:
|
||||
@@ -499,7 +770,7 @@ class UpdateRoutes:
|
||||
return git_info
|
||||
|
||||
@staticmethod
|
||||
async def _get_remote_version() -> tuple[str, List[str], List[Dict]]:
|
||||
async def _get_remote_version() -> tuple[str, List[str], List[Dict[str, Any]]]:
|
||||
"""
|
||||
Fetch remote version from GitHub
|
||||
Returns:
|
||||
@@ -521,7 +792,7 @@ class UpdateRoutes:
|
||||
|
||||
# Parse releases
|
||||
releases = []
|
||||
for i, release in enumerate(data):
|
||||
for i, release in enumerate(cast(list[dict[str, Any]], data)):
|
||||
version = release.get('tag_name', '')
|
||||
if not version.startswith('v'):
|
||||
version = f"v{version}"
|
||||
|
||||
@@ -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)):
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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).
|
||||
|
||||
|
||||
@@ -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 ```` 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 ````."""
|
||||
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:
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -0,0 +1,144 @@
|
||||
# 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
|
||||
from it have a NULL ``autov3`` column (the "not checked yet" state). This
|
||||
service computes the embedded AutoV3 hash for each such model — once per
|
||||
process — and persists it through the scanner's single write path
|
||||
(:meth:`ModelScanner.update_autov3_for_model`), marking every visited row so a
|
||||
subsequent run finds nothing left to do.
|
||||
|
||||
Three-state contract honored here:
|
||||
|
||||
- ``NULL`` (sqlite) / absent (dict) = not checked yet → backfill computes it
|
||||
- ``''`` (sqlite/dict) / JSON null = checked, no value available → never recompute
|
||||
- 12-char lowercase hex = value → never recompute
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
if TYPE_CHECKING: # pragma: no cover - type-check only; runtime imports are local
|
||||
from .model_scanner import ModelScanner
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _resolve_autov3(file_path: str) -> str:
|
||||
"""Resolve the AutoV3 hash for a model file.
|
||||
|
||||
Prefers the Civitai AutoV3 reported for the file whose SHA256 matches
|
||||
(the authoritative value for recipe matching); falls back to the embedded
|
||||
safetensors header hash. Returns ``''`` when neither is available.
|
||||
"""
|
||||
try:
|
||||
metadata_path = f"{os.path.splitext(file_path)[0]}.metadata.json"
|
||||
if os.path.exists(metadata_path):
|
||||
with open(metadata_path, "r", encoding="utf-8") as handle:
|
||||
payload = json.load(handle)
|
||||
if isinstance(payload, dict):
|
||||
from ..utils.models import autov3_from_civitai_files # local import avoids cycles
|
||||
|
||||
sha256 = (payload.get("sha256") or "").lower()
|
||||
civitai_autov3 = autov3_from_civitai_files(payload.get("civitai"), sha256)
|
||||
if civitai_autov3:
|
||||
return civitai_autov3
|
||||
except Exception:
|
||||
pass
|
||||
from ..utils.file_utils import calculate_autov3 # local import avoids cycles
|
||||
|
||||
return calculate_autov3(file_path) or ""
|
||||
|
||||
|
||||
class Autov3BackfillService:
|
||||
"""Compute and persist AutoV3 hashes for models missing a checked state."""
|
||||
|
||||
_instance: Optional["Autov3BackfillService"] = None
|
||||
_instance_lock = threading.Lock()
|
||||
|
||||
def __init__(self) -> None:
|
||||
# Re-entrancy guard per model type: scanners for different model types
|
||||
# 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[str] = set()
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls) -> "Autov3BackfillService":
|
||||
"""Return the process-wide singleton instance."""
|
||||
if cls._instance is None:
|
||||
with cls._instance_lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = cls()
|
||||
return cls._instance
|
||||
|
||||
async def backfill(self, scanner: "ModelScanner") -> int:
|
||||
"""Compute AutoV3 for every un-checked model of ``scanner.model_type``.
|
||||
|
||||
Each candidate file is read once via :func:`~py.utils.file_utils.calculate_autov3`
|
||||
(cheap: safetensors header only) and the result is persisted through
|
||||
``scanner.update_autov3_for_model``. Files that no longer exist on
|
||||
disk are skipped — they are intentionally NOT marked, because scanner
|
||||
cleanup removes the stale row later.
|
||||
|
||||
Returns:
|
||||
The number of models successfully updated. Never raises; on any
|
||||
failure a warning is logged and ``0`` is returned. A duplicate
|
||||
trigger for a model type that is already being backfilled returns
|
||||
``0`` immediately; different model types run concurrently.
|
||||
"""
|
||||
model_type = scanner.model_type
|
||||
if model_type in self._running_types:
|
||||
return 0
|
||||
self._running_types.add(model_type)
|
||||
try:
|
||||
# Local imports avoid import cycles at module load time.
|
||||
from .persistent_model_cache import get_persistent_cache
|
||||
from ..utils.file_utils import calculate_autov3
|
||||
|
||||
persistent = getattr(scanner, "_persistent_cache", None) or get_persistent_cache()
|
||||
paths = persistent.get_models_missing_autov3(model_type)
|
||||
|
||||
loop = asyncio.get_running_loop()
|
||||
count = 0
|
||||
for path in paths:
|
||||
# A file that no longer exists must not be marked; scanner
|
||||
# cleanup removes the stale row later. The existence check and
|
||||
# hash resolution run in the executor so the loop stays
|
||||
# responsive to API requests while the backfill iterates a
|
||||
# large library.
|
||||
if not await loop.run_in_executor(None, os.path.exists, path):
|
||||
continue
|
||||
autov3 = await loop.run_in_executor(None, _resolve_autov3, path)
|
||||
if await scanner.update_autov3_for_model(model_type, path, autov3):
|
||||
count += 1
|
||||
|
||||
if paths:
|
||||
logger.info(
|
||||
"AutoV3 backfill: updated %d/%d models for %s",
|
||||
count,
|
||||
len(paths),
|
||||
model_type,
|
||||
)
|
||||
else:
|
||||
# Steady state after the first run: nothing left to backfill.
|
||||
logger.debug("AutoV3 backfill: nothing to process for %s", model_type)
|
||||
return count
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"AutoV3 backfill failed for %s: %s",
|
||||
getattr(scanner, "model_type", "?"),
|
||||
exc,
|
||||
)
|
||||
return 0
|
||||
finally:
|
||||
self._running_types.discard(model_type)
|
||||
@@ -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
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
from abc import ABC, abstractmethod
|
||||
import asyncio
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional, Type, TYPE_CHECKING
|
||||
import random
|
||||
from typing import Any, Awaitable, Dict, List, Optional, Type, Union, TYPE_CHECKING, cast
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
@@ -69,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()
|
||||
|
||||
@@ -109,12 +110,15 @@ class BaseModelService(ABC):
|
||||
if civitai_model_id is not None:
|
||||
sorted_data = [
|
||||
item for item in sorted_data
|
||||
if self._extract_model_id(item) == civitai_model_id
|
||||
if self._extract_group_key(item) == civitai_model_id
|
||||
]
|
||||
# VLM mode: always sort by version ID descending (newest version first),
|
||||
# regardless of the current sort_by preference.
|
||||
# Fall back to modified timestamp for non-CivitAI sources.
|
||||
sorted_data.sort(
|
||||
key=lambda x: self._extract_version_id(x) or 0,
|
||||
key=lambda x: self._extract_version_id(x)
|
||||
or x.get("modified", 0)
|
||||
or 0,
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
@@ -129,18 +133,21 @@ class BaseModelService(ABC):
|
||||
ufs = self.settings.get("version_grouping", "same_base")
|
||||
group_by_base = ufs == "same_base"
|
||||
|
||||
dedup_map = {} # (modelId [,base_model]) -> (item, version_id)
|
||||
dedup_map = {} # (modelId [,base_model]) -> (item, version_or_modified)
|
||||
version_counter = {} # same-key -> count
|
||||
standalone = []
|
||||
for item in sorted_data:
|
||||
mid = self._extract_model_id(item)
|
||||
mid = self._extract_group_key(item)
|
||||
if mid is None:
|
||||
standalone.append(item)
|
||||
continue
|
||||
key = (mid, item.get("base_model") or "") if group_by_base else mid
|
||||
# Count all versions per key
|
||||
version_counter[key] = version_counter.get(key, 0) + 1
|
||||
vid = self._extract_version_id(item) or 0
|
||||
# Prefer CivitAI version_id; fall back to modified timestamp
|
||||
vid = self._extract_version_id(item)
|
||||
if vid is None:
|
||||
vid = item.get("modified", 0) or 0
|
||||
if key not in dedup_map or vid > dedup_map[key][1]:
|
||||
dedup_map[key] = (item, vid)
|
||||
# Attach version_count to each surviving grouped item (shallow copy
|
||||
@@ -171,19 +178,22 @@ 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_model_id(item)
|
||||
mid = self._extract_group_key(item)
|
||||
if mid is None:
|
||||
ungrouped_standalone.append(item)
|
||||
continue
|
||||
key = (mid, item.get("base_model") or "") if group_by_base else mid
|
||||
model_groups.setdefault(key, []).append(item)
|
||||
# Sort versions within each group by version id descending
|
||||
# Sort versions within each group by version id (descending);
|
||||
# fall back to modified timestamp for non-CivitAI sources.
|
||||
for items in model_groups.values():
|
||||
items.sort(
|
||||
key=lambda x: self._extract_version_id(x) or 0,
|
||||
key=lambda x: self._extract_version_id(x)
|
||||
or x.get("modified", 0)
|
||||
or 0,
|
||||
reverse=True,
|
||||
)
|
||||
# Sort groups by version count
|
||||
@@ -239,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)
|
||||
@@ -286,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]] = []
|
||||
@@ -316,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)
|
||||
|
||||
@@ -381,6 +391,12 @@ class BaseModelService(ABC):
|
||||
(item.get("model_name") or item.get("file_name") or "").lower(),
|
||||
item.get("file_path", "").lower(),
|
||||
)
|
||||
elif key_name == "random":
|
||||
# Seeded random shuffle: same seed -> same order (stable pagination)
|
||||
rng = random.Random(sort_params.seed or "random")
|
||||
result = list(data)
|
||||
rng.shuffle(result)
|
||||
return result
|
||||
elif key_name == "size":
|
||||
key_fn = lambda item: (
|
||||
int(item.get("size", 0) or 0),
|
||||
@@ -428,39 +444,50 @@ class BaseModelService(ABC):
|
||||
return entry
|
||||
|
||||
async def _apply_hash_filters(
|
||||
self, data: List[Dict], hash_filters: Dict
|
||||
) -> List[Dict]:
|
||||
"""Apply hash-based filtering"""
|
||||
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[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``
|
||||
field, both case-insensitively.
|
||||
"""
|
||||
if item.get("sha256", "").lower() in hash_set:
|
||||
return True
|
||||
autov3 = item.get("autov3", "")
|
||||
return bool(autov3) and autov3.lower() in hash_set
|
||||
|
||||
single_hash = hash_filters.get("single_hash")
|
||||
multiple_hashes = hash_filters.get("multiple_hashes")
|
||||
|
||||
if single_hash:
|
||||
# Filter by single hash
|
||||
single_hash = single_hash.lower()
|
||||
# Filter by single hash (SHA256 or AutoV3)
|
||||
return [
|
||||
item for item in data if item.get("sha256", "").lower() == single_hash
|
||||
item for item in data if matches_hash_set(item, {single_hash.lower()})
|
||||
]
|
||||
elif multiple_hashes:
|
||||
# Filter by multiple hashes
|
||||
hash_set = set(hash.lower() for hash in multiple_hashes)
|
||||
return [item for item in data if item.get("sha256", "").lower() in hash_set]
|
||||
# Filter by multiple hashes (SHA256 or AutoV3)
|
||||
hash_set = {hash.lower() for hash in multiple_hashes}
|
||||
return [item for item in data if matches_hash_set(item, hash_set)]
|
||||
|
||||
return data
|
||||
|
||||
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(
|
||||
@@ -479,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:
|
||||
@@ -526,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:
|
||||
@@ -559,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.
|
||||
@@ -575,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)
|
||||
@@ -612,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()
|
||||
@@ -632,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",
|
||||
@@ -698,7 +725,34 @@ class BaseModelService(ABC):
|
||||
return annotated
|
||||
|
||||
@staticmethod
|
||||
def _extract_model_id(item: Dict) -> Optional[int]:
|
||||
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):
|
||||
return None
|
||||
m = re.match(
|
||||
r"https?://huggingface\.co/([^/]+/[^/]+)", hf_url.strip()
|
||||
)
|
||||
if not m:
|
||||
return None
|
||||
return f"hf:{m.group(1)}"
|
||||
|
||||
@staticmethod
|
||||
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:
|
||||
1. CivitAI ``modelId`` (int)
|
||||
2. HF repo identity ``hf:{owner}/{repo}`` (str)
|
||||
3. ``None`` (no known grouping source)
|
||||
"""
|
||||
mid = BaseModelService._extract_model_id(item)
|
||||
if mid is not None:
|
||||
return mid
|
||||
return BaseModelService._extract_hf_group_key(item)
|
||||
|
||||
@staticmethod
|
||||
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
|
||||
@@ -711,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
|
||||
@@ -724,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
|
||||
@@ -776,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
|
||||
@@ -791,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
|
||||
@@ -800,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)
|
||||
|
||||
@@ -877,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 {}
|
||||
@@ -903,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()
|
||||
|
||||
@@ -932,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()
|
||||
|
||||
@@ -961,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
|
||||
@@ -1093,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
|
||||
@@ -1189,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.
|
||||
@@ -1216,19 +1270,87 @@ class BaseModelService(ABC):
|
||||
)
|
||||
|
||||
async def search_relative_paths(
|
||||
self, search_term: str, limit: int = 15, offset: int = 0
|
||||
self,
|
||||
search_term: str,
|
||||
limit: int = 15,
|
||||
offset: int = 0,
|
||||
*,
|
||||
folder: Optional[str] = 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,
|
||||
recursive: bool = True,
|
||||
apply_filters: bool = False,
|
||||
) -> List[str]:
|
||||
"""Search model relative file paths for autocomplete functionality"""
|
||||
"""Search model relative file paths for autocomplete functionality.
|
||||
|
||||
Optional filter kwargs mirror the filters used by the list endpoint
|
||||
(/api/lm/{prefix}/list). When no filter kwargs are provided the
|
||||
behavior is identical to plain token-based path matching.
|
||||
"""
|
||||
cache = await self.scanner.get_cached_data()
|
||||
include_terms, exclude_terms = self._parse_search_tokens(search_term)
|
||||
|
||||
data = cache.raw_data
|
||||
has_filters = any(
|
||||
[
|
||||
apply_filters,
|
||||
folder is not None,
|
||||
folder_include,
|
||||
folder_exclude,
|
||||
base_models,
|
||||
model_types,
|
||||
tags,
|
||||
auto_tags,
|
||||
credit_required is not None,
|
||||
allow_selling_generated_content is not None,
|
||||
]
|
||||
)
|
||||
if has_filters:
|
||||
# Auto-tags are not stored in the scanner cache — they are computed
|
||||
# on the fly. Pre-compute them only when an auto-tag filter is
|
||||
# active to avoid mutating cache entries unnecessarily.
|
||||
if auto_tags:
|
||||
from .auto_tag_service import extract_auto_tags
|
||||
|
||||
for item in data:
|
||||
if not item.get("auto_tags"):
|
||||
item["auto_tags"] = extract_auto_tags(item)
|
||||
|
||||
criteria = FilterCriteria(
|
||||
folder=folder,
|
||||
folder_include=folder_include,
|
||||
folder_exclude=folder_exclude,
|
||||
base_models=base_models,
|
||||
model_types=model_types,
|
||||
tags=tags,
|
||||
auto_tags=auto_tags,
|
||||
search_options={"recursive": recursive},
|
||||
tag_logic=tag_logic,
|
||||
)
|
||||
data = self.filter_set.apply(data, criteria)
|
||||
if credit_required is not None:
|
||||
data = await self._apply_credit_required_filter(
|
||||
data, credit_required
|
||||
)
|
||||
if allow_selling_generated_content is not None:
|
||||
data = await self._apply_allow_selling_filter(
|
||||
data, allow_selling_generated_content
|
||||
)
|
||||
|
||||
matching_paths = []
|
||||
|
||||
# Get model roots for path calculation
|
||||
model_roots = self.scanner.get_model_roots()
|
||||
|
||||
# Collect all matching paths first (needed for proper sorting and offset)
|
||||
for model in cache.raw_data:
|
||||
for model in data:
|
||||
file_path = model.get("file_path", "")
|
||||
if not file_path:
|
||||
continue
|
||||
|
||||
@@ -59,6 +59,7 @@ class CacheEntryValidator:
|
||||
'notes': ('', False),
|
||||
'usage_tips': ('', False),
|
||||
'hash_status': ('completed', False),
|
||||
'autov3': (None, False),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
@@ -119,8 +120,13 @@ class CacheEntryValidator:
|
||||
if is_required:
|
||||
errors.append(f"Required field '{field_name}' is missing or None")
|
||||
if auto_repair:
|
||||
working_entry[field_name] = cls._get_default_copy(default_value)
|
||||
repaired = True
|
||||
# A missing optional field whose default is None is already
|
||||
# semantically equal to its default (e.g. autov3: absent
|
||||
# means "not checked") — writing None back is a no-op, not
|
||||
# a repair.
|
||||
if default_value is not None:
|
||||
working_entry[field_name] = cls._get_default_copy(default_value)
|
||||
repaired = True
|
||||
continue
|
||||
|
||||
# Validate field type and value
|
||||
@@ -175,6 +181,15 @@ class CacheEntryValidator:
|
||||
# that invalidates the entry, but we also don't mark it repaired.
|
||||
pass
|
||||
|
||||
# Normalize autov3 to lowercase if needed (optional field, never stripped).
|
||||
autov3 = working_entry.get('autov3')
|
||||
if isinstance(autov3, str) and autov3:
|
||||
normalized_autov3 = autov3.lower()
|
||||
if normalized_autov3 != autov3:
|
||||
if auto_repair:
|
||||
working_entry['autov3'] = normalized_autov3
|
||||
repaired = True
|
||||
|
||||
# Determine if entry is valid
|
||||
# Entry is valid if no critical required field errors remain after repair
|
||||
# Critical fields are file_path and sha256
|
||||
@@ -242,6 +257,19 @@ class CacheEntryValidator:
|
||||
"""
|
||||
expected_type = type(default_value)
|
||||
|
||||
# Special case: autov3 is optional with a three-state contract.
|
||||
# None = not checked, "" = checked but unavailable, otherwise a
|
||||
# 12-character hex string (case-insensitive here; normalized to
|
||||
# lowercase separately).
|
||||
if field_name == 'autov3':
|
||||
if value is None or value == "":
|
||||
return None
|
||||
if not isinstance(value, str):
|
||||
return f"Field 'autov3' should be string or None, got {type(value).__name__}"
|
||||
if len(value) != 12 or any(c not in '0123456789abcdefABCDEF' for c in value):
|
||||
return "Field 'autov3' should be a 12-character hex string"
|
||||
return None
|
||||
|
||||
# Special handling for numeric types
|
||||
if expected_type == int:
|
||||
if not isinstance(value, (int, float)):
|
||||
|
||||
@@ -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
|
||||
@@ -6,7 +10,7 @@ from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from ..utils.models import CheckpointMetadata
|
||||
from ..utils.file_utils import find_preview_file, normalize_path
|
||||
from ..utils.file_utils import find_preview_file, normalize_path, calculate_autov3
|
||||
from ..utils.metadata_manager import MetadataManager
|
||||
from ..config import config
|
||||
from .model_scanner import ModelScanner
|
||||
@@ -62,6 +66,11 @@ class CheckpointScanner(ModelScanner):
|
||||
# Find preview image
|
||||
preview_url = find_preview_file(base_name, dir_path)
|
||||
|
||||
# AutoV3 reads only the safetensors header, so it is cheap even for
|
||||
# large checkpoints; record the checked state at creation time ("" =
|
||||
# checked but unavailable).
|
||||
autov3 = calculate_autov3(real_path)
|
||||
|
||||
# Create metadata WITHOUT calculating hash
|
||||
metadata = CheckpointMetadata(
|
||||
file_name=base_name,
|
||||
@@ -77,6 +86,7 @@ class CheckpointScanner(ModelScanner):
|
||||
sub_type="checkpoint",
|
||||
from_civitai=False, # Mark as local model since no hash yet
|
||||
hash_status="pending", # Mark hash as pending
|
||||
autov3=autov3 or "",
|
||||
)
|
||||
|
||||
# Save the created metadata
|
||||
@@ -120,7 +130,11 @@ class CheckpointScanner(ModelScanner):
|
||||
# that queries get_hash_by_filename first) will miss on every
|
||||
# lookup and keep calling back into this method, creating a
|
||||
# tight loop that never populates the index.
|
||||
self._hash_index.add_entry(metadata.sha256.lower(), file_path)
|
||||
self._hash_index.add_entry(
|
||||
metadata.sha256.lower(),
|
||||
file_path,
|
||||
getattr(metadata, "autov3", None) or None,
|
||||
)
|
||||
return metadata.sha256
|
||||
|
||||
async with self._hash_calculation_lock:
|
||||
@@ -132,7 +146,11 @@ class CheckpointScanner(ModelScanner):
|
||||
and metadata.hash_status == "completed"
|
||||
and metadata.sha256
|
||||
):
|
||||
self._hash_index.add_entry(metadata.sha256.lower(), file_path)
|
||||
self._hash_index.add_entry(
|
||||
metadata.sha256.lower(),
|
||||
file_path,
|
||||
getattr(metadata, "autov3", None) or None,
|
||||
)
|
||||
return metadata.sha256
|
||||
|
||||
task = self._hash_calculation_tasks.get(real_path)
|
||||
@@ -185,7 +203,11 @@ class CheckpointScanner(ModelScanner):
|
||||
if metadata.hash_status == "completed" and metadata.sha256:
|
||||
# Populate the in-memory hash index even for pre-computed
|
||||
# hashes, mirroring the fix in calculate_hash_for_model.
|
||||
self._hash_index.add_entry(metadata.sha256.lower(), file_path)
|
||||
self._hash_index.add_entry(
|
||||
metadata.sha256.lower(),
|
||||
file_path,
|
||||
getattr(metadata, "autov3", None) or None,
|
||||
)
|
||||
return metadata.sha256
|
||||
|
||||
# Update status to calculating
|
||||
@@ -202,7 +224,11 @@ class CheckpointScanner(ModelScanner):
|
||||
await MetadataManager.save_metadata(file_path, metadata)
|
||||
|
||||
# Update hash index
|
||||
self._hash_index.add_entry(sha256.lower(), file_path)
|
||||
self._hash_index.add_entry(
|
||||
sha256.lower(),
|
||||
file_path,
|
||||
getattr(metadata, "autov3", None) or None,
|
||||
)
|
||||
|
||||
# Update the in-memory cache entry so that subsequent
|
||||
# _persist_current_cache / _save_persistent_cache calls
|
||||
@@ -216,6 +242,7 @@ class CheckpointScanner(ModelScanner):
|
||||
if entry.get("file_path") == file_path:
|
||||
entry["sha256"] = sha256.lower()
|
||||
entry["hash_status"] = "completed"
|
||||
self.bump_cache_version()
|
||||
break
|
||||
|
||||
logger.info(f"Hash calculated for checkpoint: {file_path}")
|
||||
@@ -405,7 +432,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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -283,7 +283,7 @@ class CivitaiBaseModelService:
|
||||
return None
|
||||
|
||||
if isinstance(result, str):
|
||||
data = json.loads(result)
|
||||
data: Any = json.loads(result)
|
||||
else:
|
||||
data = result
|
||||
|
||||
|
||||
+128
-40
@@ -1,9 +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,
|
||||
@@ -19,6 +24,12 @@ from ..utils.civitai_utils import resolve_license_payload
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Best-effort cache for creator model counts, keyed by lowercase username.
|
||||
# Values are (monotonic timestamp, count or None); None results are cached
|
||||
# too so repeated failures don't hammer the API.
|
||||
_CREATOR_COUNT_CACHE_TTL_SECONDS = 600
|
||||
_creator_model_count_cache: Dict[str, Tuple[float, Optional[int]]] = {}
|
||||
|
||||
|
||||
class CivitaiClient:
|
||||
_instance = None
|
||||
@@ -51,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
|
||||
|
||||
@@ -65,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)."""
|
||||
|
||||
@@ -79,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:
|
||||
@@ -119,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
|
||||
@@ -166,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",
|
||||
@@ -213,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:
|
||||
@@ -268,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(
|
||||
@@ -276,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", []),
|
||||
@@ -310,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] = {}
|
||||
@@ -340,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] = {
|
||||
@@ -366,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:
|
||||
@@ -385,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
|
||||
@@ -404,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
|
||||
@@ -457,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
|
||||
|
||||
@@ -479,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
|
||||
@@ -487,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
|
||||
|
||||
@@ -496,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
|
||||
@@ -505,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}")
|
||||
@@ -525,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", {})
|
||||
@@ -535,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
|
||||
@@ -548,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 = {}
|
||||
@@ -564,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:
|
||||
@@ -589,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)
|
||||
@@ -619,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:
|
||||
@@ -652,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:
|
||||
@@ -692,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
|
||||
@@ -709,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]
|
||||
@@ -729,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)
|
||||
@@ -743,17 +755,34 @@ class CivitaiClient:
|
||||
|
||||
return all_versions if all_versions else None
|
||||
|
||||
async def get_user_models(self, username: str) -> Optional[List[Dict]]:
|
||||
"""Fetch all models for a specific Civitai user."""
|
||||
async def get_user_models(
|
||||
self, username: str, cursor: Optional[str] = None
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""Fetch one page (up to 100 models) for a specific Civitai user.
|
||||
|
||||
Returns ``{"items": [...], "nextCursor": <str|None>}`` on success,
|
||||
or None on failure. Pass ``cursor`` (from a previous response's
|
||||
``nextCursor``) to fetch subsequent pages.
|
||||
"""
|
||||
if not username:
|
||||
return None
|
||||
|
||||
params: Dict[str, Any] = {
|
||||
"username": username,
|
||||
"nsfw": "true",
|
||||
"limit": 100,
|
||||
"sort": "Newest",
|
||||
"period": "AllTime",
|
||||
}
|
||||
if cursor:
|
||||
params["cursor"] = cursor
|
||||
|
||||
try:
|
||||
success, result = await self._make_request(
|
||||
"GET",
|
||||
f"{self.base_url}/models",
|
||||
use_auth=True,
|
||||
params={"username": username, "nsfw": "true"},
|
||||
params=params,
|
||||
)
|
||||
|
||||
if not success:
|
||||
@@ -765,7 +794,7 @@ class CivitaiClient:
|
||||
|
||||
items = result.get("items") if isinstance(result, dict) else None
|
||||
if not isinstance(items, list):
|
||||
return []
|
||||
items = []
|
||||
|
||||
for model in items:
|
||||
versions = model.get("modelVersions")
|
||||
@@ -774,9 +803,68 @@ class CivitaiClient:
|
||||
for version in versions:
|
||||
self._remove_comfy_metadata(version)
|
||||
|
||||
return items
|
||||
next_cursor: Optional[str] = None
|
||||
metadata = result.get("metadata") if isinstance(result, dict) else None
|
||||
if isinstance(metadata, dict):
|
||||
raw_cursor = metadata.get("nextCursor")
|
||||
if raw_cursor is not None:
|
||||
next_cursor = str(raw_cursor)
|
||||
|
||||
return {"items": items, "nextCursor": next_cursor}
|
||||
except RateLimitError:
|
||||
raise
|
||||
except Exception as exc: # pragma: no cover - defensive logging
|
||||
logger.error("Error fetching models for %s: %s", username, exc)
|
||||
return None
|
||||
|
||||
async def get_creator_model_count(self, username: str) -> Optional[int]:
|
||||
"""Best-effort lookup of a creator's published model count.
|
||||
|
||||
Uses the ``/creators`` endpoint (a contains-match query), picking the
|
||||
entry whose username matches exactly (case-insensitive). Returns None
|
||||
on any failure; never raises. Results (including None) are cached
|
||||
for ``_CREATOR_COUNT_CACHE_TTL_SECONDS``.
|
||||
"""
|
||||
if not username:
|
||||
return None
|
||||
|
||||
cache_key = username.lower()
|
||||
cached = _creator_model_count_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
cached_at, cached_count = cached
|
||||
if time.monotonic() - cached_at < _CREATOR_COUNT_CACHE_TTL_SECONDS:
|
||||
return cached_count
|
||||
|
||||
count: Optional[int] = None
|
||||
try:
|
||||
success, result = await self._make_request(
|
||||
"GET",
|
||||
f"{self.base_url}/creators",
|
||||
use_auth=True,
|
||||
params={"query": username, "limit": 10},
|
||||
)
|
||||
|
||||
if success and isinstance(result, dict):
|
||||
creators = result.get("items")
|
||||
if isinstance(creators, list):
|
||||
for creator in creators:
|
||||
if not isinstance(creator, dict):
|
||||
continue
|
||||
creator_name = creator.get("username")
|
||||
if not isinstance(creator_name, str):
|
||||
continue
|
||||
if creator_name.lower() != cache_key:
|
||||
continue
|
||||
model_count = creator.get("modelCount")
|
||||
if isinstance(model_count, (int, float)) and not isinstance(
|
||||
model_count, bool
|
||||
):
|
||||
count = int(model_count)
|
||||
break
|
||||
except Exception as exc: # best-effort only, never propagate
|
||||
logger.debug(
|
||||
"Failed to fetch creator model count for %s: %s", username, exc
|
||||
)
|
||||
|
||||
_creator_model_count_cache[cache_key] = (time.monotonic(), count)
|
||||
return count
|
||||
|
||||
@@ -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
|
||||
|
||||
+135
-76
@@ -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 (
|
||||
@@ -18,7 +22,7 @@ from ..utils.constants import (
|
||||
VALID_LORA_TYPES,
|
||||
)
|
||||
from ..utils.civitai_utils import normalize_civitai_download_url, rewrite_preview_url
|
||||
from ..utils.file_utils import calculate_sha256
|
||||
from ..utils.file_utils import calculate_sha256, calculate_autov3
|
||||
from ..utils.preview_selection import resolve_mature_threshold, select_preview_media
|
||||
from ..utils.utils import sanitize_folder_name
|
||||
from ..utils.exif_utils import ExifUtils
|
||||
@@ -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,8 +1397,20 @@ 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)
|
||||
# Security: validate path containment after joining
|
||||
resolved_dir = os.path.abspath(os.path.normpath(save_dir))
|
||||
base_dir = os.path.abspath(os.path.normpath(base_save_dir))
|
||||
if not resolved_dir.startswith(base_dir + os.sep) and resolved_dir != base_dir:
|
||||
logger.warning(
|
||||
"Path traversal detected: %s escapes %s",
|
||||
resolved_dir, base_dir,
|
||||
)
|
||||
return {"success": False, "error": "Download path is outside allowed directory"}
|
||||
# Create directory if it doesn't exist
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
@@ -1551,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:
|
||||
@@ -1570,7 +1596,7 @@ class DownloadManager:
|
||||
},
|
||||
)
|
||||
|
||||
execute_kwargs = {
|
||||
execute_kwargs: Dict[str, Any] = {
|
||||
"download_urls": download_urls,
|
||||
"save_dir": save_dir,
|
||||
"metadata": metadata,
|
||||
@@ -1617,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
|
||||
@@ -1642,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:
|
||||
@@ -1673,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,
|
||||
)
|
||||
@@ -1691,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."""
|
||||
@@ -1715,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
|
||||
@@ -1726,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",
|
||||
@@ -1763,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
|
||||
|
||||
@@ -1785,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
|
||||
|
||||
@@ -1827,6 +1854,9 @@ class DownloadManager:
|
||||
model_tags, model_type
|
||||
)
|
||||
|
||||
if not first_tag:
|
||||
first_tag = "no tags" # Default if no tags available
|
||||
|
||||
# Format the template with available data
|
||||
formatted_path = path_template
|
||||
formatted_path = formatted_path.replace("{base_model}", mapped_base_model)
|
||||
@@ -1842,6 +1872,15 @@ class DownloadManager:
|
||||
if model_type == "embedding":
|
||||
formatted_path = formatted_path.replace(" ", "_")
|
||||
|
||||
# Sanitize the resolved path to prevent path traversal:
|
||||
# - Strip leading slashes (prevents os.path.join from treating path as absolute)
|
||||
# - Collapse double slashes from empty placeholder substitutions
|
||||
# - Strip trailing slashes for cleanliness
|
||||
formatted_path = formatted_path.lstrip("/")
|
||||
while "//" in formatted_path:
|
||||
formatted_path = formatted_path.replace("//", "/")
|
||||
formatted_path = formatted_path.rstrip("/")
|
||||
|
||||
return formatted_path
|
||||
|
||||
async def _execute_download(
|
||||
@@ -1849,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(
|
||||
@@ -1911,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:
|
||||
@@ -1937,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:
|
||||
@@ -1956,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(
|
||||
@@ -1982,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"
|
||||
)
|
||||
@@ -2034,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(
|
||||
@@ -2138,6 +2183,10 @@ class DownloadManager:
|
||||
"error": f"Zip archive does not contain any supported model files ({supported_text})",
|
||||
}
|
||||
actual_file_paths = extracted_paths
|
||||
# The archive entry's AutoV3 (if any) describes the zip itself,
|
||||
# not the extracted models; clear it so per-file header
|
||||
# resolution applies to every extracted model.
|
||||
metadata.autov3 = None
|
||||
try:
|
||||
os.remove(save_path)
|
||||
except OSError as exc:
|
||||
@@ -2213,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 = (
|
||||
@@ -2333,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
|
||||
@@ -2352,6 +2401,16 @@ class DownloadManager:
|
||||
sha256 = await calculate_sha256(file_path)
|
||||
if sha256:
|
||||
entry.sha256 = sha256.lower()
|
||||
# AutoV3: the Civitai-reported value for the downloaded file (set
|
||||
# by from_civitai_info) takes precedence. Only the un-checked
|
||||
# state (None) triggers a header read; '' (checked-unavailable)
|
||||
# is never re-read, honoring the three-state contract so rows
|
||||
# marked at download time stay untouched by later passes.
|
||||
if entry.autov3 is None:
|
||||
autov3 = await asyncio.get_running_loop().run_in_executor(
|
||||
None, calculate_autov3, file_path
|
||||
)
|
||||
entry.autov3 = (autov3 or "").lower()
|
||||
entries.append(entry)
|
||||
|
||||
return entries
|
||||
@@ -2370,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 []
|
||||
@@ -2429,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:
|
||||
@@ -2511,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
|
||||
@@ -2594,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()
|
||||
@@ -2641,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()
|
||||
@@ -2658,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():
|
||||
@@ -2785,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:
|
||||
|
||||
@@ -31,7 +31,7 @@ class DownloadQueueService:
|
||||
_instance: Optional[DownloadQueueService] = None
|
||||
_class_lock: asyncio.Lock = asyncio.Lock()
|
||||
|
||||
_SCHEMA = """
|
||||
_SCHEMA_TABLES = """
|
||||
CREATE TABLE IF NOT EXISTS download_queue (
|
||||
download_id TEXT PRIMARY KEY,
|
||||
model_id INTEGER,
|
||||
@@ -74,6 +74,9 @@ class DownloadQueueService:
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_dh_completed ON download_history(completed_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_dh_status ON download_history(status);
|
||||
"""
|
||||
|
||||
_CREATE_UNIQUE_INDEX = """
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS idx_dh_download_id
|
||||
ON download_history(download_id) WHERE download_id IS NOT NULL;
|
||||
"""
|
||||
@@ -115,10 +118,39 @@ class DownloadQueueService:
|
||||
if self._schema_initialized:
|
||||
return
|
||||
with self._connect() as conn:
|
||||
conn.executescript(self._SCHEMA)
|
||||
conn.executescript(self._SCHEMA_TABLES)
|
||||
|
||||
# Creating the unique index on download_history.download_id can
|
||||
# fail if pre-existing rows have duplicate values (e.g. from a
|
||||
# previous version that lacked the index). Deduplicate first so
|
||||
# that the migration does not crash on startup.
|
||||
if not self._index_exists(conn, "idx_dh_download_id"):
|
||||
self._remove_duplicate_download_ids(conn)
|
||||
conn.executescript(self._CREATE_UNIQUE_INDEX)
|
||||
|
||||
conn.commit()
|
||||
self._schema_initialized = True
|
||||
|
||||
@staticmethod
|
||||
def _index_exists(conn: sqlite3.Connection, name: str) -> bool:
|
||||
return conn.execute(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='index' AND name=?",
|
||||
(name,),
|
||||
).fetchone() is not None
|
||||
|
||||
@staticmethod
|
||||
def _remove_duplicate_download_ids(conn: sqlite3.Connection) -> None:
|
||||
conn.execute("""
|
||||
DELETE FROM download_history
|
||||
WHERE id NOT IN (
|
||||
SELECT MIN(id)
|
||||
FROM download_history
|
||||
WHERE download_id IS NOT NULL
|
||||
GROUP BY download_id
|
||||
)
|
||||
AND download_id IS NOT NULL
|
||||
""")
|
||||
|
||||
def get_database_path(self) -> str:
|
||||
"""Return the resolved database file path."""
|
||||
return self._db_path
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -201,6 +201,11 @@ PROVIDER_PRESETS: Dict[str, Dict[str, Any]] = {
|
||||
"api_base": "https://openrouter.ai/api/v1",
|
||||
"requires_key": True,
|
||||
},
|
||||
"google": {
|
||||
"name": "Gemini",
|
||||
"api_base": "https://generativelanguage.googleapis.com/v1beta/openai",
|
||||
"requires_key": True,
|
||||
},
|
||||
"opencode-go": {
|
||||
"name": "OpenCode Go",
|
||||
"api_base": "https://opencode.ai/zen/go/v1",
|
||||
@@ -566,18 +571,52 @@ class LLMService:
|
||||
if effective_max is None:
|
||||
effective_max = 4096
|
||||
|
||||
result = await self.chat_completion(
|
||||
messages=messages,
|
||||
model=model,
|
||||
temperature=temperature,
|
||||
response_format={"type": "json_object"},
|
||||
max_tokens=effective_max,
|
||||
)
|
||||
# Use json_schema (not json_object) for broader provider compatibility:
|
||||
# LM Studio and some other OpenAI-compatible servers reject
|
||||
# json_object but accept json_schema. {"type": "object"} is
|
||||
# functionally equivalent — it accepts any JSON object without
|
||||
# constraining specific fields.
|
||||
response_format = {
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "metadata",
|
||||
"schema": {"type": "object"},
|
||||
},
|
||||
}
|
||||
|
||||
try:
|
||||
result = await self.chat_completion(
|
||||
messages=messages,
|
||||
model=model,
|
||||
temperature=temperature,
|
||||
response_format=response_format,
|
||||
max_tokens=effective_max,
|
||||
)
|
||||
except LLMResponseError as e:
|
||||
# Only fall back when the provider rejects the response_format
|
||||
# type value (e.g. "'response_format.type' must be..."). Avoid
|
||||
# catching unrelated 400 errors whose body happens to mention
|
||||
# "response_format" (e.g. "model does not support
|
||||
# response_format restrictions on this endpoint").
|
||||
if "'response_format.type'" not in str(e).lower():
|
||||
raise
|
||||
logger.info(
|
||||
"Provider rejected response_format, retrying without it. "
|
||||
"Falling back to prompt-only JSON mode. Error: %s",
|
||||
e,
|
||||
)
|
||||
result = await self.chat_completion(
|
||||
messages=messages,
|
||||
model=model,
|
||||
temperature=temperature,
|
||||
response_format=None,
|
||||
max_tokens=effective_max,
|
||||
)
|
||||
|
||||
content = result.get("content", "") or ""
|
||||
if not content:
|
||||
raise LLMResponseError(
|
||||
"LLM returned empty content in json_object mode. "
|
||||
"LLM returned empty content. "
|
||||
f"Raw response: {json.dumps(result)[:500]}"
|
||||
)
|
||||
|
||||
|
||||
@@ -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")
|
||||
|
||||
|
||||
+41
-41
@@ -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.
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -6,25 +6,26 @@ 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
|
||||
from ..utils.model_utils import determine_base_model
|
||||
from ..utils.models import autov3_from_civitai_files
|
||||
from .connectivity_guard import OFFLINE_FRIENDLY_MESSAGE, is_expected_offline_error
|
||||
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]]:
|
||||
...
|
||||
|
||||
@@ -38,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
|
||||
@@ -152,6 +153,18 @@ class MetadataSyncService:
|
||||
civitai_metadata.get("baseModel")
|
||||
)
|
||||
|
||||
# Civitai-first AutoV3 propagation: the freshly fetched version
|
||||
# metadata may report an AutoV3 for the file whose SHA256 matches the
|
||||
# local model. Persist it now so recipe matching sees it immediately —
|
||||
# no full rescan or restart required (the header is never re-read to
|
||||
# upgrade the checked-unavailable '' state).
|
||||
sha256_value = (local_metadata.get("sha256") or "").lower()
|
||||
civitai_autov3 = autov3_from_civitai_files(
|
||||
local_metadata.get("civitai"), sha256_value
|
||||
)
|
||||
if civitai_autov3:
|
||||
local_metadata["autov3"] = civitai_autov3
|
||||
|
||||
await self._preview_service.ensure_preview_for_metadata(
|
||||
metadata_path, local_metadata, civitai_metadata.get("images", [])
|
||||
)
|
||||
@@ -479,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": {},
|
||||
|
||||
+35
-21
@@ -1,6 +1,7 @@
|
||||
import asyncio
|
||||
import time
|
||||
import logging
|
||||
import random
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
@@ -30,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) -> sorted list
|
||||
self._last_sort: Tuple[str, str] = (None, 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
|
||||
@@ -63,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):
|
||||
@@ -79,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":
|
||||
@@ -113,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
|
||||
@@ -142,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
|
||||
@@ -176,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]]:
|
||||
@@ -203,9 +209,9 @@ class ModelCache:
|
||||
async def resort(self):
|
||||
"""Resort cached data according to last sort mode if set"""
|
||||
async with self._lock:
|
||||
if self._last_sort != (None, None):
|
||||
sort_key, order = self._last_sort
|
||||
sorted_data = self._sort_data(self.raw_data, sort_key, order)
|
||||
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
|
||||
# else: do nothing
|
||||
@@ -218,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) -> 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')
|
||||
@@ -265,6 +271,13 @@ class ModelCache:
|
||||
),
|
||||
reverse=reverse
|
||||
)
|
||||
elif sort_key == 'random':
|
||||
# Random shuffle seeded for stable pagination: the same seed
|
||||
# always yields the same order, so successive page requests
|
||||
# stay consistent while browsing.
|
||||
rng = random.Random(seed or 'random')
|
||||
result = list(data)
|
||||
rng.shuffle(result)
|
||||
elif sort_key == 'versions_count':
|
||||
# Pre-dedup sort: fall back to name sort.
|
||||
# Actual re-sort by version_count happens in get_paginated_data after dedup.
|
||||
@@ -285,15 +298,16 @@ 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') -> 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:
|
||||
if (sort_key, order) == self._last_sort:
|
||||
cache_key = (sort_key, order, seed)
|
||||
if cache_key == self._last_sort:
|
||||
return self._last_sorted_data
|
||||
|
||||
start_time = time.perf_counter()
|
||||
sorted_data = self._sort_data(self.raw_data, sort_key, order)
|
||||
self._last_sort = (sort_key, order)
|
||||
sorted_data = self._sort_data(self.raw_data, sort_key, order, seed)
|
||||
self._last_sort = cache_key
|
||||
self._last_sorted_data = sorted_data
|
||||
|
||||
duration = time.perf_counter() - start_time
|
||||
@@ -312,9 +326,9 @@ class ModelCache:
|
||||
|
||||
self.name_display_mode = normalized
|
||||
|
||||
if self._last_sort[0] == 'name':
|
||||
sort_key, order = self._last_sort
|
||||
self._last_sorted_data = self._sort_data(self.raw_data, sort_key, order)
|
||||
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:
|
||||
"""Update preview_url for a specific model in all cached data
|
||||
|
||||
@@ -8,6 +8,7 @@ from abc import ABC, abstractmethod
|
||||
from ..utils.utils import calculate_relative_path_for_model, remove_empty_dirs
|
||||
from ..utils.constants import AUTO_ORGANIZE_BATCH_SIZE
|
||||
from ..services.settings_manager import get_settings_manager
|
||||
from ..services.model_lifecycle_service import _require_path_in_library_roots
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -40,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',
|
||||
@@ -417,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
|
||||
@@ -493,6 +496,9 @@ class ModelMoveService:
|
||||
Dictionary with move result
|
||||
"""
|
||||
try:
|
||||
_require_path_in_library_roots(file_path, self.scanner, label="Source path")
|
||||
_require_path_in_library_roots(target_path, self.scanner, label="Target path")
|
||||
|
||||
if use_default_paths:
|
||||
# Find the model in cache to get metadata
|
||||
cache = await self.scanner.get_cached_data()
|
||||
|
||||
@@ -8,11 +8,12 @@ class ModelHashIndex:
|
||||
self._hash_to_path: Dict[str, str] = {}
|
||||
self._filename_to_hash: Dict[str, str] = {}
|
||||
self._autov2_to_path: Dict[str, str] = {}
|
||||
self._autov3_to_path: Dict[str, str] = {}
|
||||
# New data structures for tracking duplicates
|
||||
self._duplicate_hashes: Dict[str, List[str]] = {} # sha256 -> list of paths
|
||||
self._duplicate_filenames: Dict[str, List[str]] = {} # filename -> list of paths
|
||||
|
||||
def add_entry(self, sha256: str, file_path: str) -> None:
|
||||
def add_entry(self, sha256: str, file_path: str, autov3: Optional[str] = None) -> None:
|
||||
"""Add or update hash index entry"""
|
||||
if not sha256 or not file_path:
|
||||
return
|
||||
@@ -33,9 +34,14 @@ class ModelHashIndex:
|
||||
self._duplicate_hashes.setdefault(sha256, []).append(file_path)
|
||||
|
||||
# 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)
|
||||
# Same path registered again (e.g. a file replaced in place with
|
||||
# new content) — used below to drop its stale autov3 mapping.
|
||||
is_re_registration = existing_path == file_path
|
||||
|
||||
# If this is a different file with the same filename
|
||||
if existing_path and existing_path != file_path:
|
||||
@@ -67,12 +73,36 @@ class ModelHashIndex:
|
||||
# AutoV2 = first 10 chars of SHA256
|
||||
if len(sha256) >= 10:
|
||||
self._autov2_to_path[sha256[:10]] = file_path
|
||||
# AutoV3 is an independent hash (not derived from SHA256), stored as-is.
|
||||
# Drop stale mappings for a path when it is re-registered with a NEW
|
||||
# sha256 (file replaced in place) or with an explicit new autov3 value
|
||||
# (correction). Re-registering the SAME file with the same sha256 and
|
||||
# no autov3 (e.g. lazy-hash completion) must never clear its existing
|
||||
# mapping. First-time registrations stay O(1).
|
||||
if autov3:
|
||||
autov3 = autov3.lower()
|
||||
if is_re_registration and (existing_hash != sha256 or autov3):
|
||||
stale_autov3_keys = [
|
||||
key for key, mapped_path in self._autov3_to_path.items()
|
||||
if mapped_path == file_path and key != autov3
|
||||
]
|
||||
for key in stale_autov3_keys:
|
||||
del self._autov3_to_path[key]
|
||||
if autov3:
|
||||
self._autov3_to_path[autov3] = file_path
|
||||
|
||||
def add_autov3(self, autov3: str, file_path: str) -> None:
|
||||
"""Add or update an AutoV3-only index entry (used when only AutoV3 is known)"""
|
||||
if not autov3:
|
||||
return
|
||||
autov3 = autov3.lower()
|
||||
self._autov3_to_path[autov3] = file_path
|
||||
|
||||
def _get_filename_from_path(self, file_path: str) -> str:
|
||||
"""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)
|
||||
|
||||
@@ -167,6 +197,11 @@ class ModelHashIndex:
|
||||
for k in autov2_keys_to_remove:
|
||||
del self._autov2_to_path[k]
|
||||
|
||||
# Remove from AutoV3 index
|
||||
autov3_keys_to_remove = [k for k, v in self._autov3_to_path.items() if v == file_path]
|
||||
for k in autov3_keys_to_remove:
|
||||
del self._autov3_to_path[k]
|
||||
|
||||
def remove_by_hash(self, sha256: str) -> None:
|
||||
"""Remove entry by hash"""
|
||||
sha256 = sha256.lower()
|
||||
@@ -189,6 +224,11 @@ class ModelHashIndex:
|
||||
autov2_key = sha256[:10]
|
||||
if autov2_key in self._autov2_to_path:
|
||||
del self._autov2_to_path[autov2_key]
|
||||
|
||||
# Remove AutoV3 entries pointing to any removed path
|
||||
autov3_keys_to_remove = [k for k, v in self._autov3_to_path.items() if v in paths_to_remove]
|
||||
for k in autov3_keys_to_remove:
|
||||
del self._autov3_to_path[k]
|
||||
|
||||
# Update filename-to-hash and duplicate filenames for all paths
|
||||
for path_to_remove in paths_to_remove:
|
||||
@@ -209,22 +249,26 @@ class ModelHashIndex:
|
||||
del self._duplicate_filenames[fname]
|
||||
|
||||
def has_hash(self, hash_value: str) -> bool:
|
||||
"""Check if hash exists in index (SHA256 or AutoV2)"""
|
||||
"""Check if hash exists in index (SHA256, AutoV2, or AutoV3)"""
|
||||
normalized = hash_value.lower()
|
||||
if normalized in self._hash_to_path:
|
||||
return True
|
||||
if len(normalized) == 10:
|
||||
return normalized in self._autov2_to_path
|
||||
if len(normalized) == 12:
|
||||
return normalized in self._autov3_to_path
|
||||
return False
|
||||
|
||||
def get_path(self, hash_value: str) -> Optional[str]:
|
||||
"""Get file path for a hash (SHA256 or AutoV2)"""
|
||||
"""Get file path for a hash (SHA256, AutoV2, or AutoV3)"""
|
||||
normalized = hash_value.lower()
|
||||
path = self._hash_to_path.get(normalized)
|
||||
if path is not None:
|
||||
return path
|
||||
if len(normalized) == 10:
|
||||
return self._autov2_to_path.get(normalized)
|
||||
if len(normalized) == 12:
|
||||
return self._autov3_to_path.get(normalized)
|
||||
return None
|
||||
|
||||
def get_hash(self, file_path: str) -> Optional[str]:
|
||||
@@ -243,6 +287,7 @@ class ModelHashIndex:
|
||||
self._hash_to_path.clear()
|
||||
self._filename_to_hash.clear()
|
||||
self._autov2_to_path.clear()
|
||||
self._autov3_to_path.clear()
|
||||
self._duplicate_hashes.clear()
|
||||
self._duplicate_filenames.clear()
|
||||
|
||||
@@ -253,6 +298,10 @@ class ModelHashIndex:
|
||||
def get_all_filenames(self) -> Set[str]:
|
||||
"""Get all filenames in the index"""
|
||||
return set(self._filename_to_hash.keys())
|
||||
|
||||
def get_all_autov3(self) -> Dict[str, str]:
|
||||
"""Get a snapshot of all AutoV3 hashes mapped to their file paths"""
|
||||
return dict(self._autov3_to_path)
|
||||
|
||||
def get_duplicate_hashes(self) -> Dict[str, List[str]]:
|
||||
"""Get dictionary of duplicate hashes and their paths"""
|
||||
|
||||
@@ -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
|
||||
@@ -48,6 +48,36 @@ async def delete_model_artifacts(
|
||||
return deleted
|
||||
|
||||
|
||||
def _require_path_in_library_roots(file_path: str, scanner, *, label: str = "path") -> None:
|
||||
"""Raise ``ValueError`` if *file_path* is not inside a configured model root.
|
||||
|
||||
Uses ``os.path.abspath()`` (NOT ``realpath``) to resolve ``..`` and ``.``
|
||||
while preserving symlinks — this keeps the check in business-path space.
|
||||
Skips when the scanner does not expose ``get_model_roots`` or the list
|
||||
is empty.
|
||||
"""
|
||||
|
||||
roots = None
|
||||
if hasattr(scanner, "get_model_roots"):
|
||||
try:
|
||||
roots = scanner.get_model_roots()
|
||||
except NotImplementedError:
|
||||
roots = None
|
||||
if not roots:
|
||||
return
|
||||
|
||||
resolved = os.path.abspath(os.path.normpath(file_path))
|
||||
|
||||
for root in roots:
|
||||
root_resolved = os.path.abspath(os.path.normpath(root))
|
||||
if resolved == root_resolved or resolved.startswith(root_resolved + os.sep):
|
||||
return
|
||||
|
||||
raise ValueError(
|
||||
f"{label} '{file_path}' is outside configured library directories"
|
||||
)
|
||||
|
||||
|
||||
class ModelLifecycleService:
|
||||
"""Co-ordinate destructive and mutating model operations."""
|
||||
|
||||
@@ -57,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
|
||||
@@ -74,6 +104,8 @@ class ModelLifecycleService:
|
||||
if not file_path:
|
||||
raise ValueError("Model path is required")
|
||||
|
||||
_require_path_in_library_roots(file_path, self._scanner, label="File path")
|
||||
|
||||
cache = await self._scanner.get_cached_data()
|
||||
|
||||
cached_entry = None
|
||||
@@ -106,6 +138,9 @@ class ModelLifecycleService:
|
||||
item for item in cache.raw_data if item.get("file_path") != file_path
|
||||
]
|
||||
await cache.resort()
|
||||
bump_cache_version = getattr(self._scanner, "bump_cache_version", None)
|
||||
if callable(bump_cache_version):
|
||||
bump_cache_version()
|
||||
|
||||
if hasattr(self._scanner, "_hash_index") and self._scanner._hash_index:
|
||||
self._scanner._hash_index.remove_by_path(file_path)
|
||||
@@ -114,7 +149,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}
|
||||
|
||||
@@ -182,6 +217,8 @@ class ModelLifecycleService:
|
||||
if not file_path:
|
||||
raise ValueError("Model path is required")
|
||||
|
||||
_require_path_in_library_roots(file_path, self._scanner, label="File path")
|
||||
|
||||
metadata_path = os.path.splitext(file_path)[0] + ".metadata.json"
|
||||
metadata = await self._metadata_loader(metadata_path)
|
||||
metadata["exclude"] = True
|
||||
@@ -210,6 +247,9 @@ class ModelLifecycleService:
|
||||
item for item in cache.raw_data if item["file_path"] != file_path
|
||||
]
|
||||
await cache.resort()
|
||||
bump_cache_version = getattr(self._scanner, "bump_cache_version", None)
|
||||
if callable(bump_cache_version):
|
||||
bump_cache_version()
|
||||
|
||||
excluded = getattr(self._scanner, "_excluded_models", None)
|
||||
if isinstance(excluded, list):
|
||||
@@ -218,7 +258,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}
|
||||
@@ -229,6 +269,8 @@ class ModelLifecycleService:
|
||||
if not file_path:
|
||||
raise ValueError("Model path is required")
|
||||
|
||||
_require_path_in_library_roots(file_path, self._scanner, label="File path")
|
||||
|
||||
if not os.path.exists(file_path):
|
||||
raise ValueError("Model file does not exist")
|
||||
|
||||
@@ -270,6 +312,9 @@ class ModelLifecycleService:
|
||||
if not file_paths:
|
||||
raise ValueError("No file paths provided for deletion")
|
||||
|
||||
for path in file_paths:
|
||||
_require_path_in_library_roots(path, self._scanner, label="File path")
|
||||
|
||||
return await self._scanner.bulk_delete_models(file_paths)
|
||||
|
||||
async def rename_model(
|
||||
@@ -280,6 +325,8 @@ class ModelLifecycleService:
|
||||
if not file_path or not new_file_name:
|
||||
raise ValueError("File path and new file name are required")
|
||||
|
||||
_require_path_in_library_roots(file_path, self._scanner, label="File path")
|
||||
|
||||
invalid_chars = {"/", "\\", ":", "*", "?", '"', "<", ">", "|"}
|
||||
if any(char in new_file_name for char in invalid_chars):
|
||||
raise ValueError("Invalid characters in file name")
|
||||
@@ -316,7 +363,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
|
||||
|
||||
@@ -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,50 +133,61 @@ 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) -> Optional[List[Dict]]:
|
||||
"""Fetch models owned by the specified user"""
|
||||
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,
|
||||
or None when unsupported/failed. ``cursor`` continues a previous page.
|
||||
"""
|
||||
pass
|
||||
|
||||
async def get_creator_model_count(self, username: str) -> Optional[int]:
|
||||
"""Published model count for the user; None when unsupported."""
|
||||
return None
|
||||
|
||||
class CivitaiModelMetadataProvider(ModelMetadataProvider):
|
||||
"""Provider that uses Civitai API for metadata"""
|
||||
|
||||
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) -> Optional[List[Dict]]:
|
||||
return await self.client.get_user_models(username)
|
||||
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]:
|
||||
return await self.client.get_creator_model_count(username)
|
||||
|
||||
class CivArchiveModelMetadataProvider(ModelMetadataProvider):
|
||||
"""Provider that uses CivArchive API for metadata"""
|
||||
@@ -184,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) -> Optional[List[Dict]]:
|
||||
async def get_user_models(self, username: str, cursor: Optional[str] = None) -> Optional[Dict[str, Any]]:
|
||||
"""Not supported by CivArchive provider"""
|
||||
return None
|
||||
|
||||
@@ -207,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
|
||||
@@ -232,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
|
||||
@@ -288,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
|
||||
@@ -328,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
|
||||
@@ -347,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) -> Optional[List[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 = ?"
|
||||
@@ -474,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(
|
||||
@@ -496,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:
|
||||
@@ -527,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(
|
||||
@@ -550,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(
|
||||
@@ -574,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(
|
||||
@@ -602,13 +613,14 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
continue
|
||||
return None
|
||||
|
||||
async def get_user_models(self, username: str) -> Optional[List[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(
|
||||
label,
|
||||
provider.get_user_models,
|
||||
username,
|
||||
cursor=cursor,
|
||||
)
|
||||
if result is not None:
|
||||
return result
|
||||
@@ -624,6 +636,19 @@ class FallbackMetadataProvider(ModelMetadataProvider):
|
||||
continue
|
||||
return None
|
||||
|
||||
async def get_creator_model_count(self, username: str) -> Optional[int]:
|
||||
for provider, label in self._iter_providers():
|
||||
try:
|
||||
result = await provider.get_creator_model_count(username)
|
||||
if result is not None:
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
"Provider %s failed for get_creator_model_count: %s", label, e
|
||||
)
|
||||
continue
|
||||
return None
|
||||
|
||||
def _iter_providers(self):
|
||||
return zip(self.providers, self._provider_labels)
|
||||
|
||||
@@ -656,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,
|
||||
@@ -673,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,
|
||||
@@ -682,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,
|
||||
@@ -697,20 +722,24 @@ 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) -> Optional[List[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,
|
||||
username,
|
||||
cursor=cursor,
|
||||
)
|
||||
|
||||
async def get_creator_model_count(self, username: str) -> Optional[int]:
|
||||
return await self._provider.get_creator_model_count(username)
|
||||
|
||||
class ModelMetadataProviderManager:
|
||||
"""Manager for selecting and using model metadata providers"""
|
||||
|
||||
@@ -733,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)
|
||||
@@ -746,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:
|
||||
@@ -755,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)
|
||||
@@ -768,20 +797,30 @@ 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)
|
||||
except NotImplementedError:
|
||||
return None
|
||||
|
||||
async def get_user_models(self, username: str, provider_name: str = None) -> Optional[List[Dict]]:
|
||||
"""Fetch models owned by the specified user"""
|
||||
async def get_user_models(
|
||||
self,
|
||||
username: str,
|
||||
provider_name: Optional[str] = None,
|
||||
cursor: Optional[str] = None,
|
||||
) -> 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)
|
||||
return await provider.get_user_models(username, cursor)
|
||||
|
||||
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:
|
||||
|
||||
@@ -12,6 +12,7 @@ from typing import (
|
||||
Tuple,
|
||||
Protocol,
|
||||
Callable,
|
||||
cast,
|
||||
)
|
||||
|
||||
from ..utils.constants import NSFW_LEVELS
|
||||
@@ -85,6 +86,7 @@ class SortParams:
|
||||
|
||||
key: str
|
||||
order: str
|
||||
seed: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -116,7 +118,7 @@ class ModelCacheRepository:
|
||||
async def fetch_sorted(self, params: SortParams) -> List[Dict[str, Any]]:
|
||||
"""Fetch cached data pre-sorted according to ``params``."""
|
||||
cache = await self.get_cache()
|
||||
return await cache.get_sorted_data(params.key, params.order)
|
||||
return await cache.get_sorted_data(params.key, params.order, params.seed)
|
||||
|
||||
@staticmethod
|
||||
def parse_sort(sort_by: str) -> SortParams:
|
||||
@@ -132,10 +134,17 @@ class ModelCacheRepository:
|
||||
sort_key = sort_by.strip().lower() or "name"
|
||||
order = "asc"
|
||||
|
||||
if order not in ("asc", "desc"):
|
||||
seed = None
|
||||
if sort_key == "random":
|
||||
# Random sort: the portion after ':' is the shuffle seed.
|
||||
# A stable seed keeps paginated requests consistent; order is
|
||||
# meaningless for a random shuffle.
|
||||
seed = order if order and order not in ("asc", "desc") else None
|
||||
order = "asc"
|
||||
elif order not in ("asc", "desc"):
|
||||
order = "asc"
|
||||
|
||||
return SortParams(key=sort_key, order=order)
|
||||
return SortParams(key=sort_key, order=order, seed=seed)
|
||||
|
||||
|
||||
class ModelFilterSet:
|
||||
@@ -301,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"
|
||||
|
||||
+289
-44
@@ -5,16 +5,16 @@ 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
|
||||
from ..utils.models import BaseModelMetadata, autov3_from_civitai_files
|
||||
from ..config import config
|
||||
from ..utils.file_utils import find_preview_file, get_preview_extension, calculate_sha256
|
||||
from ..utils.file_utils import find_preview_file, get_preview_extension, calculate_sha256, calculate_autov3
|
||||
from ..utils.metadata_manager import MetadataManager
|
||||
from ..utils.civitai_utils import resolve_license_info
|
||||
from .model_cache import ModelCache
|
||||
from .model_hash_index import ModelHashIndex
|
||||
from .model_lifecycle_service import delete_model_artifacts
|
||||
from .model_lifecycle_service import delete_model_artifacts, _require_path_in_library_roots
|
||||
from .service_registry import ServiceRegistry
|
||||
from .websocket_manager import ws_manager
|
||||
from .persistent_model_cache import get_persistent_cache
|
||||
@@ -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,8 @@ 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._cache_version: int = 0
|
||||
self._hash_index = hash_index or ModelHashIndex()
|
||||
self._tags_count = {} # Dictionary to store tag counts
|
||||
self._is_initializing = False # Flag to track initialization state
|
||||
@@ -86,6 +87,7 @@ class ModelScanner:
|
||||
self._persistent_cache = get_persistent_cache()
|
||||
self._name_display_mode = self._resolve_name_display_mode()
|
||||
self._cancel_requested = False # Flag for cancellation
|
||||
self._autov3_backfill_scheduled = False # One-time AutoV3 backfill trigger per process
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
@@ -97,6 +99,25 @@ class ModelScanner:
|
||||
# Register this service
|
||||
asyncio.create_task(self._register_service())
|
||||
|
||||
@property
|
||||
def cache_version(self) -> int:
|
||||
"""Monotonic version counter for the in-memory cache.
|
||||
|
||||
Every write path that mutates scanner cache state calls
|
||||
:meth:`bump_cache_version`, so consumers (e.g. RecipeScanner) can
|
||||
detect when a cached derivation of the raw data is stale. Reads never
|
||||
bump.
|
||||
"""
|
||||
return self._cache_version
|
||||
|
||||
def bump_cache_version(self) -> None:
|
||||
"""Invalidate derived caches by incrementing the cache version.
|
||||
|
||||
Public because external services (model lifecycle, route handlers)
|
||||
rewrite scanner raw_data directly and must be able to invalidate it.
|
||||
"""
|
||||
self._cache_version += 1
|
||||
|
||||
def on_library_changed(self) -> None:
|
||||
"""Reset caches when the active library changes."""
|
||||
self._persistent_cache = get_persistent_cache()
|
||||
@@ -106,6 +127,7 @@ class ModelScanner:
|
||||
self._excluded_models = []
|
||||
self._is_initializing = False
|
||||
self._name_display_mode = self._resolve_name_display_mode()
|
||||
self.bump_cache_version()
|
||||
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
@@ -182,7 +204,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()
|
||||
@@ -225,6 +247,19 @@ class ModelScanner:
|
||||
if not isinstance(notes, str):
|
||||
notes = str(notes)
|
||||
|
||||
# AutoV3 three-state contract: absent key / None = "not checked yet",
|
||||
# "" = "checked but unavailable" (never re-read the header), else the
|
||||
# 12-char lowercase hex value. A metadata object already follows the
|
||||
# contract and is passed through unchanged; a payload dict only carries
|
||||
# an explicit checked state when the key is present.
|
||||
if is_mapping:
|
||||
if 'autov3' in source:
|
||||
entry_autov3 = source['autov3'] or ''
|
||||
else:
|
||||
entry_autov3 = None
|
||||
else:
|
||||
entry_autov3 = get_value('autov3', None)
|
||||
|
||||
entry: Dict[str, Any] = {
|
||||
'file_path': normalized_path,
|
||||
# file_name is always stored WITHOUT extension (e.g. "OWSMianne_ANIMA_V1",
|
||||
@@ -238,6 +273,7 @@ class ModelScanner:
|
||||
'size': int(get_value('size', 0) or 0),
|
||||
'modified': float(get_value('modified', 0.0) or 0.0),
|
||||
'sha256': (get_value('sha256', '') or '').lower(),
|
||||
'autov3': entry_autov3,
|
||||
'base_model': get_value('base_model', '') or '',
|
||||
'preview_url': preview_url,
|
||||
'preview_nsfw_level': int(get_value('preview_nsfw_level', 0) or 0),
|
||||
@@ -473,6 +509,13 @@ class ModelScanner:
|
||||
if sha_value and path:
|
||||
hash_index.add_entry(sha_value.lower(), path)
|
||||
|
||||
# Rebuild the AutoV3 index from the persisted autov3_index rows. These
|
||||
# cover every known autov3 -> path mapping regardless of whether a
|
||||
# sha256 row also exists for the same file.
|
||||
for autov3_value, path in persisted.autov3_hash_rows:
|
||||
if autov3_value and path:
|
||||
hash_index.add_autov3(autov3_value.lower(), path)
|
||||
|
||||
tags_count: Dict[str, int] = {}
|
||||
adjusted_raw_data: List[Dict[str, Any]] = []
|
||||
for item in persisted.raw_data:
|
||||
@@ -541,8 +584,30 @@ class ModelScanner:
|
||||
'scanner_type': self.model_type,
|
||||
'pageType': page_type
|
||||
})
|
||||
|
||||
# Schedule the one-time AutoV3 backfill task (at most once per process)
|
||||
# so entries loaded from a persisted snapshot that predates autov3 get
|
||||
# their checked state computed in the background. The task never blocks
|
||||
# or crashes the load path.
|
||||
if not self._autov3_backfill_scheduled:
|
||||
self._autov3_backfill_scheduled = True
|
||||
try:
|
||||
loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
loop = None
|
||||
if loop is not None:
|
||||
loop.create_task(self._run_autov3_backfill())
|
||||
|
||||
return True
|
||||
|
||||
async def _run_autov3_backfill(self) -> None:
|
||||
"""Backfill autov3 for entries loaded from the persisted cache that lack it."""
|
||||
try:
|
||||
from ..services.autov3_backfill_service import Autov3BackfillService # lazy import (module created by another unit)
|
||||
await Autov3BackfillService.get_instance().backfill(self)
|
||||
except Exception as exc:
|
||||
logger.warning("AutoV3 backfill failed: %s", exc)
|
||||
|
||||
async def _save_persistent_cache(self, scan_result: CacheBuildResult) -> None:
|
||||
if not scan_result or not getattr(self, '_persistent_cache', None):
|
||||
return
|
||||
@@ -555,6 +620,7 @@ class ModelScanner:
|
||||
return
|
||||
|
||||
hash_snapshot = self._build_hash_index_snapshot(scan_result.hash_index)
|
||||
autov3_snapshot = self._build_autov3_index_snapshot(scan_result.hash_index)
|
||||
loop = asyncio.get_event_loop()
|
||||
try:
|
||||
await loop.run_in_executor(
|
||||
@@ -563,7 +629,8 @@ class ModelScanner:
|
||||
self.model_type,
|
||||
list(scan_result.raw_data),
|
||||
hash_snapshot,
|
||||
list(scan_result.excluded_models)
|
||||
list(scan_result.excluded_models),
|
||||
autov3_snapshot,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("%s Scanner: Failed to persist cache: %s", self.model_type.capitalize(), exc)
|
||||
@@ -589,6 +656,20 @@ class ModelScanner:
|
||||
bucket.append(path)
|
||||
return snapshot
|
||||
|
||||
def _build_autov3_index_snapshot(self, hash_index: Optional[ModelHashIndex]) -> Dict[str, List[str]]:
|
||||
"""Build the autov3 -> [paths] snapshot for the persisted cache."""
|
||||
snapshot: Dict[str, List[str]] = {}
|
||||
if not hash_index:
|
||||
return snapshot
|
||||
|
||||
for autov3_value, path in hash_index.get_all_autov3().items():
|
||||
if not autov3_value or not path:
|
||||
continue
|
||||
bucket = snapshot.setdefault(autov3_value.lower(), [])
|
||||
if path not in bucket:
|
||||
bucket.append(path)
|
||||
return snapshot
|
||||
|
||||
async def _persist_current_cache(self) -> None:
|
||||
if self._cache is None or not getattr(self, '_persistent_cache', None):
|
||||
return
|
||||
@@ -712,7 +793,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"""
|
||||
@@ -872,6 +953,8 @@ class ModelScanner:
|
||||
)
|
||||
continue
|
||||
model_data = validation_result.entry
|
||||
if model_data is None:
|
||||
continue
|
||||
|
||||
self._ensure_license_flags(model_data)
|
||||
# Add to cache
|
||||
@@ -880,7 +963,11 @@ class ModelScanner:
|
||||
|
||||
# Update hash index if available
|
||||
if 'sha256' in model_data and 'file_path' in model_data:
|
||||
self._hash_index.add_entry(model_data['sha256'].lower(), model_data['file_path'])
|
||||
self._hash_index.add_entry(
|
||||
model_data['sha256'].lower(),
|
||||
model_data['file_path'],
|
||||
model_data.get('autov3') or None
|
||||
)
|
||||
|
||||
# Update tags count
|
||||
if 'tags' in model_data and model_data['tags']:
|
||||
@@ -927,6 +1014,25 @@ class ModelScanner:
|
||||
# Update cache data
|
||||
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[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:
|
||||
seen_paths.add(path)
|
||||
deduped.append(item)
|
||||
else:
|
||||
for tag in item.get('tags', []):
|
||||
if tag in self._tags_count:
|
||||
self._tags_count[tag] = max(0, self._tags_count[tag] - 1)
|
||||
if self._tags_count[tag] == 0:
|
||||
del self._tags_count[tag]
|
||||
dedup_removed += 1
|
||||
if dedup_removed > 0:
|
||||
self._cache.raw_data = list(reversed(deduped))
|
||||
total_removed += dedup_removed
|
||||
|
||||
# Resort cache if changes were made
|
||||
if total_added > 0 or total_removed > 0:
|
||||
# Update folders list
|
||||
@@ -945,6 +1051,7 @@ class ModelScanner:
|
||||
logger.error(f"{self.model_type.capitalize()} Scanner: Error reconciling cache: {e}", exc_info=True)
|
||||
finally:
|
||||
self._is_initializing = False # Unset flag
|
||||
self.bump_cache_version()
|
||||
|
||||
def is_initializing(self) -> bool:
|
||||
"""Check if the scanner is currently initializing"""
|
||||
@@ -1025,7 +1132,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
|
||||
@@ -1049,7 +1156,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)")
|
||||
@@ -1086,6 +1193,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)
|
||||
|
||||
@@ -1111,6 +1220,36 @@ class ModelScanner:
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to compute SHA256 for {file_path}: {e}")
|
||||
|
||||
# AutoV3 resolution: prefer the Civitai AutoV3 reported for the file
|
||||
# whose SHA256 matches (authoritative for recipe matching), falling
|
||||
# back to the embedded safetensors header hash only for models never
|
||||
# checked before (autov3 is None). A checked-unavailable state ('')
|
||||
# is only upgraded by Civitai data — the header is never re-read.
|
||||
current_autov3 = model_data.get('autov3')
|
||||
if current_autov3 in (None, ''):
|
||||
try:
|
||||
civitai_data = None
|
||||
if isinstance(metadata, BaseModelMetadata):
|
||||
civitai_data = metadata.civitai
|
||||
elif isinstance(metadata, dict):
|
||||
civitai_data = metadata.get("civitai")
|
||||
autov3 = autov3_from_civitai_files(
|
||||
civitai_data, model_data.get("sha256") or ""
|
||||
) or ""
|
||||
if not autov3 and current_autov3 is None:
|
||||
autov3 = (calculate_autov3(os.path.realpath(file_path)) or '').lower()
|
||||
if autov3 != current_autov3:
|
||||
model_data['autov3'] = autov3
|
||||
if isinstance(metadata, BaseModelMetadata):
|
||||
metadata.autov3 = autov3
|
||||
await MetadataManager.save_metadata(file_path, metadata)
|
||||
elif isinstance(metadata, dict):
|
||||
# Dict payload: JSON null encodes the checked-unavailable state.
|
||||
metadata['autov3'] = autov3 or None
|
||||
await MetadataManager.save_metadata(file_path, metadata)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to resolve AutoV3 for {file_path}: {e}")
|
||||
|
||||
# Skip excluded models
|
||||
if model_data.get('exclude', False):
|
||||
excluded_models.append(model_data['file_path'])
|
||||
@@ -1150,6 +1289,8 @@ class ModelScanner:
|
||||
|
||||
self._log_duplicate_filename_summary()
|
||||
|
||||
self.bump_cache_version()
|
||||
|
||||
def _log_duplicate_filename_summary(self) -> None:
|
||||
"""Log a batched summary of duplicate filename conflicts once per scan."""
|
||||
# Duplicate filename detection is only relevant for LoRAs, which use
|
||||
@@ -1183,7 +1324,7 @@ class ModelScanner:
|
||||
|
||||
async def _sync_download_history(
|
||||
self,
|
||||
raw_data: List[Mapping[str, Any]],
|
||||
raw_data: Sequence[Mapping[str, Any]],
|
||||
*,
|
||||
source: str,
|
||||
) -> None:
|
||||
@@ -1232,7 +1373,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] = []
|
||||
@@ -1296,6 +1437,8 @@ class ModelScanner:
|
||||
)
|
||||
continue
|
||||
result = validation_result.entry
|
||||
if result is None:
|
||||
continue
|
||||
|
||||
self._ensure_license_flags(result)
|
||||
raw_data.append(result)
|
||||
@@ -1303,7 +1446,7 @@ class ModelScanner:
|
||||
sha_value = result.get('sha256')
|
||||
model_path = result.get('file_path')
|
||||
if sha_value and model_path:
|
||||
hash_index.add_entry(sha_value.lower(), model_path)
|
||||
hash_index.add_entry(sha_value.lower(), model_path, result.get('autov3') or None)
|
||||
|
||||
for tag in result.get('tags') or []:
|
||||
tags_count[tag] = tags_count.get(tag, 0) + 1
|
||||
@@ -1335,7 +1478,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:
|
||||
@@ -1348,31 +1491,44 @@ 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
|
||||
|
||||
# Add to cache
|
||||
self._cache.raw_data.append(metadata_dict)
|
||||
self._cache.add_to_version_index(metadata_dict)
|
||||
file_path = metadata_dict.get('file_path', '')
|
||||
if file_path:
|
||||
old_entries = [item for item in self._cache.raw_data if item.get('file_path') == file_path]
|
||||
for old_entry in old_entries:
|
||||
for tag in old_entry.get('tags', []):
|
||||
if tag in self._tags_count:
|
||||
self._tags_count[tag] = max(0, self._tags_count[tag] - 1)
|
||||
if self._tags_count[tag] == 0:
|
||||
del self._tags_count[tag]
|
||||
self._hash_index.remove_by_path(file_path)
|
||||
self._cache.raw_data = [item for item in self._cache.raw_data if item.get('file_path') != file_path]
|
||||
|
||||
for tag in metadata_dict.get('tags', []):
|
||||
self._tags_count[tag] = self._tags_count.get(tag, 0) + 1
|
||||
|
||||
self._cache.raw_data.append(metadata_dict)
|
||||
|
||||
# Resort cache data
|
||||
await self._cache.resort()
|
||||
|
||||
# Update folders list
|
||||
all_folders = set(self._cache.folders)
|
||||
all_folders.add(folder)
|
||||
self._cache.folders = sorted(list(all_folders), key=lambda x: x.lower())
|
||||
|
||||
# Update the hash index
|
||||
self._hash_index.add_entry(metadata_dict['sha256'], metadata_dict['file_path'])
|
||||
self._hash_index.add_entry(
|
||||
metadata_dict['sha256'],
|
||||
metadata_dict['file_path'],
|
||||
metadata_dict.get('autov3') or None,
|
||||
)
|
||||
await self._persist_current_cache()
|
||||
self.bump_cache_version()
|
||||
return True
|
||||
except Exception as e:
|
||||
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:
|
||||
@@ -1394,6 +1550,9 @@ class ModelScanner:
|
||||
|
||||
base_name = os.path.splitext(os.path.basename(source_path))[0]
|
||||
source_dir = os.path.dirname(source_path)
|
||||
|
||||
_require_path_in_library_roots(source_path, self, label="Source path")
|
||||
_require_path_in_library_roots(target_path, self, label="Target path")
|
||||
|
||||
os.makedirs(target_path, exist_ok=True)
|
||||
|
||||
@@ -1403,7 +1562,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, '/')
|
||||
@@ -1451,7 +1610,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)
|
||||
@@ -1469,7 +1628,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:
|
||||
@@ -1495,7 +1654,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()
|
||||
|
||||
@@ -1518,6 +1677,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, '/')
|
||||
@@ -1549,7 +1709,11 @@ class ModelScanner:
|
||||
|
||||
sha_value = cache_entry.get('sha256')
|
||||
if sha_value:
|
||||
self._hash_index.add_entry(sha_value.lower(), normalized_new_path)
|
||||
self._hash_index.add_entry(
|
||||
sha_value.lower(),
|
||||
normalized_new_path,
|
||||
cache_entry.get('autov3') or None,
|
||||
)
|
||||
|
||||
all_folders = set(item['folder'] for item in cache.raw_data)
|
||||
cache.folders = sorted(list(all_folders), key=lambda x: x.lower())
|
||||
@@ -1563,8 +1727,11 @@ class ModelScanner:
|
||||
|
||||
if cache_modified:
|
||||
await self._persist_current_cache()
|
||||
self.bump_cache_version()
|
||||
|
||||
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]
|
||||
@@ -1687,10 +1854,11 @@ class ModelScanner:
|
||||
# ---- In-place update of the cache entry ----
|
||||
existing_entry.clear()
|
||||
existing_entry.update(desired_entry)
|
||||
self.bump_cache_version()
|
||||
|
||||
# ---- 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:
|
||||
@@ -1707,7 +1875,11 @@ class ModelScanner:
|
||||
if old_sha:
|
||||
self._hash_index.remove_by_path(file_path)
|
||||
if new_sha:
|
||||
self._hash_index.add_entry(new_sha, file_path)
|
||||
self._hash_index.add_entry(
|
||||
new_sha,
|
||||
file_path,
|
||||
desired_entry.get('autov3') or None,
|
||||
)
|
||||
|
||||
# ---- Incremental version index update ----
|
||||
new_civitai = desired_entry.get("civitai")
|
||||
@@ -1723,7 +1895,7 @@ class ModelScanner:
|
||||
# ---- Conditional resort (only when sort-key fields changed) ----
|
||||
need_resort = False
|
||||
_last = cache._last_sort
|
||||
sort_key: Optional[str] = _last[0] if _last != (None, None) else None
|
||||
sort_key: Optional[str] = _last[0] if _last[0] is not None else None
|
||||
if sort_key == "name":
|
||||
if (
|
||||
old_model_name != desired_entry.get("model_name", "")
|
||||
@@ -1758,6 +1930,75 @@ class ModelScanner:
|
||||
|
||||
return True
|
||||
|
||||
async def update_autov3_for_model(self, model_type: str, file_path: str, autov3: str) -> bool:
|
||||
"""Persist an AutoV3 hash for a single model (single write path used by the backfill service).
|
||||
|
||||
Locates the in-memory cache entry by ``file_path`` and updates only its
|
||||
``autov3`` field: the in-memory hash index, the SQLite snapshot via
|
||||
:meth:`PersistentModelCache.update_single_model`, and the
|
||||
``.metadata.json`` sidecar. sha256, tags, and every other field are
|
||||
left untouched, so the persistent delta only ever differs in autov3.
|
||||
|
||||
Returns:
|
||||
``True`` when the entry was found and updated, ``False`` otherwise.
|
||||
Never raises — failures are logged and swallowed.
|
||||
"""
|
||||
try:
|
||||
if self._cache is None:
|
||||
return False
|
||||
|
||||
entry = next(
|
||||
(item for item in self._cache.raw_data if item.get('file_path') == file_path),
|
||||
None,
|
||||
)
|
||||
if entry is None:
|
||||
return False
|
||||
|
||||
# Normalize once so the memory entry, sidecar, and SQLite row agree.
|
||||
autov3 = (autov3 or "").lower()
|
||||
|
||||
# Capture the pre-mutation state so update_single_model only sees
|
||||
# an autov3 delta between old and new.
|
||||
old_item = dict(entry)
|
||||
|
||||
entry['autov3'] = autov3 or ''
|
||||
|
||||
# Prefer add_entry when a sha256 is known so the sha256 and autov3
|
||||
# maps stay in sync; fall back to an autov3-only registration.
|
||||
sha_value = entry.get('sha256')
|
||||
checked_autov3 = entry.get('autov3') or None
|
||||
if sha_value:
|
||||
self._hash_index.add_entry(sha_value.lower(), file_path, checked_autov3)
|
||||
elif checked_autov3:
|
||||
self._hash_index.add_autov3(checked_autov3, file_path)
|
||||
|
||||
persistent = getattr(self, '_persistent_cache', None)
|
||||
if persistent is not None:
|
||||
await asyncio.get_event_loop().run_in_executor(
|
||||
None,
|
||||
persistent.update_single_model,
|
||||
model_type,
|
||||
entry,
|
||||
old_item,
|
||||
)
|
||||
|
||||
# Sidecar write-back: JSON null encodes the checked-unavailable
|
||||
# state. Skip silently when the sidecar does not exist.
|
||||
metadata_path = f"{os.path.splitext(file_path)[0]}.metadata.json"
|
||||
if os.path.exists(metadata_path):
|
||||
with open(metadata_path, 'r', encoding='utf-8') as handle:
|
||||
payload = json.load(handle)
|
||||
if not isinstance(payload, dict):
|
||||
payload = {}
|
||||
payload['autov3'] = entry['autov3'] or None
|
||||
await MetadataManager.save_metadata(metadata_path, payload)
|
||||
|
||||
self.bump_cache_version()
|
||||
return True
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to update AutoV3 for %s: %s", file_path, exc)
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _cache_entries_differ(a: Dict[str, Any], b: Dict[str, Any]) -> bool:
|
||||
"""Return ``True`` when two cache-entry dicts differ in any field.
|
||||
@@ -1817,7 +2058,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()
|
||||
|
||||
@@ -1833,7 +2074,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``
|
||||
@@ -1856,7 +2097,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()
|
||||
|
||||
@@ -1937,7 +2178,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:
|
||||
@@ -1971,6 +2212,8 @@ class ModelScanner:
|
||||
break
|
||||
|
||||
try:
|
||||
_require_path_in_library_roots(file_path, self, label="File path")
|
||||
|
||||
target_dir = os.path.dirname(file_path)
|
||||
base_name = os.path.basename(file_path)
|
||||
file_name, main_extension = os.path.splitext(base_name)
|
||||
@@ -2083,6 +2326,8 @@ class ModelScanner:
|
||||
|
||||
await self._persist_current_cache()
|
||||
|
||||
self.bump_cache_version()
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
@@ -2133,7 +2378,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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -3,8 +3,8 @@ import logging
|
||||
import os
|
||||
import sqlite3
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Mapping, Optional, Sequence, Tuple
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
|
||||
|
||||
from ..utils.cache_paths import CacheType, resolve_cache_path_with_migration
|
||||
|
||||
@@ -15,9 +15,10 @@ 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)
|
||||
|
||||
|
||||
DEFAULT_LICENSE_FLAGS = 127 # 127 (0b1111111) encodes default CivitAI permissions with all commercial modes enabled.
|
||||
@@ -36,6 +37,7 @@ class PersistentModelCache:
|
||||
"size",
|
||||
"modified",
|
||||
"sha256",
|
||||
"autov3",
|
||||
"base_model",
|
||||
"preview_url",
|
||||
"preview_nsfw_level",
|
||||
@@ -68,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
|
||||
@@ -118,6 +120,10 @@ class PersistentModelCache:
|
||||
"SELECT sha256, file_path FROM hash_index WHERE model_type = ?",
|
||||
(model_type,),
|
||||
).fetchall()
|
||||
autov3_rows = conn.execute(
|
||||
"SELECT autov3, file_path FROM autov3_index WHERE model_type = ?",
|
||||
(model_type,),
|
||||
).fetchall()
|
||||
excluded = conn.execute(
|
||||
"SELECT file_path FROM excluded_models WHERE model_type = ?",
|
||||
(model_type,),
|
||||
@@ -128,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 = []
|
||||
@@ -139,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")
|
||||
@@ -191,6 +197,8 @@ class PersistentModelCache:
|
||||
"hash_status": row["hash_status"] or "completed",
|
||||
"hf_url": row["hf_url"] or "",
|
||||
}
|
||||
if row["autov3"] is not None:
|
||||
item["autov3"] = (row["autov3"] or "").lower()
|
||||
raw_data.append(item)
|
||||
|
||||
hash_pairs = [(entry["sha256"].lower(), entry["file_path"]) for entry in hash_rows if entry["sha256"]]
|
||||
@@ -201,10 +209,21 @@ class PersistentModelCache:
|
||||
if sha_value:
|
||||
hash_pairs.append((sha_value.lower(), item["file_path"]))
|
||||
|
||||
excluded_paths = [row["file_path"] for row in excluded]
|
||||
return PersistedCacheData(raw_data=raw_data, hash_rows=hash_pairs, excluded_models=excluded_paths)
|
||||
autov3_pairs = [
|
||||
(entry["autov3"].lower(), entry["file_path"])
|
||||
for entry in autov3_rows
|
||||
if entry["autov3"]
|
||||
]
|
||||
|
||||
def save_cache(self, model_type: str, raw_data: Sequence[Dict], hash_index: Dict[str, List[str]], excluded_models: Sequence[str]) -> None:
|
||||
excluded_paths = [row["file_path"] for row in excluded]
|
||||
return PersistedCacheData(
|
||||
raw_data=raw_data,
|
||||
hash_rows=hash_pairs,
|
||||
excluded_models=excluded_paths,
|
||||
autov3_hash_rows=autov3_pairs,
|
||||
)
|
||||
|
||||
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:
|
||||
@@ -219,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
|
||||
}
|
||||
|
||||
@@ -251,13 +270,17 @@ class PersistentModelCache:
|
||||
"DELETE FROM hash_index WHERE model_type = ? AND file_path = ?",
|
||||
to_remove_models,
|
||||
)
|
||||
conn.executemany(
|
||||
"DELETE FROM autov3_index WHERE model_type = ? AND file_path = ?",
|
||||
to_remove_models,
|
||||
)
|
||||
conn.executemany(
|
||||
"DELETE FROM excluded_models WHERE model_type = ? AND file_path = ?",
|
||||
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)
|
||||
@@ -289,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:
|
||||
@@ -332,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:
|
||||
@@ -373,6 +396,52 @@ class PersistentModelCache:
|
||||
hash_inserts,
|
||||
)
|
||||
|
||||
if autov3_hash_index is not None:
|
||||
existing_autov3_rows = conn.execute(
|
||||
"SELECT autov3, file_path FROM autov3_index WHERE model_type = ?",
|
||||
(model_type,),
|
||||
).fetchall()
|
||||
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[str]] = {}
|
||||
for autov3_value, paths in autov3_hash_index.items():
|
||||
normalized_autov3 = (autov3_value or "").lower()
|
||||
if not normalized_autov3:
|
||||
continue
|
||||
bucket = new_autov3_map.setdefault(normalized_autov3, set())
|
||||
for path in paths:
|
||||
if path:
|
||||
bucket.add(path)
|
||||
|
||||
autov3_inserts: List[Tuple[str, str, str]] = []
|
||||
autov3_deletes: List[Tuple[str, str, str]] = []
|
||||
|
||||
all_autov3 = set(existing_autov3_map.keys()) | set(new_autov3_map.keys())
|
||||
for autov3_value in all_autov3:
|
||||
existing_paths = existing_autov3_map.get(autov3_value, set())
|
||||
new_paths = new_autov3_map.get(autov3_value, set())
|
||||
|
||||
for path in existing_paths - new_paths:
|
||||
autov3_deletes.append((model_type, autov3_value, path))
|
||||
for path in new_paths - existing_paths:
|
||||
autov3_inserts.append((model_type, autov3_value, path))
|
||||
|
||||
if autov3_deletes:
|
||||
conn.executemany(
|
||||
"DELETE FROM autov3_index WHERE model_type = ? AND autov3 = ? AND file_path = ?",
|
||||
autov3_deletes,
|
||||
)
|
||||
if autov3_inserts:
|
||||
conn.executemany(
|
||||
"INSERT OR IGNORE INTO autov3_index (model_type, autov3, file_path) VALUES (?, ?, ?)",
|
||||
autov3_inserts,
|
||||
)
|
||||
|
||||
existing_excluded_rows = conn.execute(
|
||||
"SELECT file_path FROM excluded_models WHERE model_type = ?",
|
||||
(model_type,),
|
||||
@@ -435,6 +504,7 @@ class PersistentModelCache:
|
||||
size INTEGER,
|
||||
modified REAL,
|
||||
sha256 TEXT,
|
||||
autov3 TEXT,
|
||||
base_model TEXT,
|
||||
preview_url TEXT,
|
||||
preview_nsfw_level INTEGER,
|
||||
@@ -472,6 +542,13 @@ class PersistentModelCache:
|
||||
PRIMARY KEY (model_type, sha256, file_path)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS autov3_index (
|
||||
model_type TEXT NOT NULL,
|
||||
autov3 TEXT NOT NULL,
|
||||
file_path TEXT NOT NULL,
|
||||
PRIMARY KEY (model_type, autov3, file_path)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS excluded_models (
|
||||
model_type TEXT NOT NULL,
|
||||
file_path TEXT NOT NULL,
|
||||
@@ -504,6 +581,7 @@ class PersistentModelCache:
|
||||
"license_flags": f"INTEGER DEFAULT {DEFAULT_LICENSE_FLAGS}",
|
||||
"hash_status": "TEXT DEFAULT 'completed'",
|
||||
"hf_url": "TEXT DEFAULT ''",
|
||||
"autov3": "TEXT",
|
||||
}
|
||||
|
||||
for column, definition in required_columns.items():
|
||||
@@ -522,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):
|
||||
@@ -549,6 +627,12 @@ class PersistentModelCache:
|
||||
if license_flags is None:
|
||||
license_flags = DEFAULT_LICENSE_FLAGS
|
||||
|
||||
autov3_value = item.get("autov3")
|
||||
if autov3_value is None:
|
||||
autov3_column = None
|
||||
else:
|
||||
autov3_column = (autov3_value or "").lower()
|
||||
|
||||
return (
|
||||
model_type,
|
||||
item.get("file_path"),
|
||||
@@ -558,6 +642,7 @@ class PersistentModelCache:
|
||||
int(item.get("size") or 0),
|
||||
float(item.get("modified") or 0.0),
|
||||
(item.get("sha256") or "").lower() or None,
|
||||
autov3_column,
|
||||
item.get("base_model") or "",
|
||||
item.get("preview_url") or "",
|
||||
int(item.get("preview_nsfw_level") or 0),
|
||||
@@ -590,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.
|
||||
|
||||
@@ -630,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
|
||||
|
||||
@@ -663,6 +748,25 @@ class PersistentModelCache:
|
||||
(model_type, new_sha, file_path),
|
||||
)
|
||||
|
||||
# --- autov3_index ---
|
||||
new_autov3: Optional[str] = new_item.get("autov3")
|
||||
if new_autov3 is not None:
|
||||
new_autov3 = (new_autov3 or "").lower()
|
||||
old_autov3: Optional[str] = (old_item.get("autov3") if old_item else None)
|
||||
if old_autov3 is not None:
|
||||
old_autov3 = (old_autov3 or "").lower()
|
||||
if new_autov3 != old_autov3:
|
||||
if old_autov3:
|
||||
conn.execute(
|
||||
"DELETE FROM autov3_index WHERE model_type = ? AND autov3 = ? AND file_path = ?",
|
||||
(model_type, old_autov3, file_path),
|
||||
)
|
||||
if new_autov3:
|
||||
conn.execute(
|
||||
"INSERT OR IGNORE INTO autov3_index (model_type, autov3, file_path) VALUES (?, ?, ?)",
|
||||
(model_type, new_autov3, file_path),
|
||||
)
|
||||
|
||||
conn.execute("COMMIT")
|
||||
except Exception:
|
||||
conn.execute("ROLLBACK")
|
||||
@@ -676,6 +780,40 @@ class PersistentModelCache:
|
||||
exc,
|
||||
)
|
||||
|
||||
def get_models_missing_autov3(self, model_type: str) -> List[str]:
|
||||
"""Return file paths whose models lack an AutoV3 checked state.
|
||||
|
||||
Only rows with a completed sha256 and a NULL autov3 column qualify —
|
||||
rows with '' (checked-unavailable) or a value are never returned, so
|
||||
the backfill query self-terminates.
|
||||
"""
|
||||
if not self.is_enabled():
|
||||
return []
|
||||
if not self._schema_initialized:
|
||||
self._initialize_schema()
|
||||
if not self._schema_initialized:
|
||||
return []
|
||||
try:
|
||||
with self._db_lock:
|
||||
conn = self._connect(readonly=True)
|
||||
try:
|
||||
rows = conn.execute(
|
||||
"SELECT file_path FROM models "
|
||||
"WHERE model_type = ? AND autov3 IS NULL "
|
||||
"AND sha256 IS NOT NULL AND sha256 != ''",
|
||||
(model_type,),
|
||||
).fetchall()
|
||||
finally:
|
||||
conn.close()
|
||||
return [row["file_path"] for row in rows]
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to query models missing autov3 for %s: %s",
|
||||
model_type,
|
||||
exc,
|
||||
)
|
||||
return []
|
||||
|
||||
def _load_tags(self, conn: sqlite3.Connection, model_type: str) -> Dict[str, List[str]]:
|
||||
tag_rows = conn.execute(
|
||||
"SELECT file_path, tag FROM model_tags WHERE model_type = ?",
|
||||
|
||||
@@ -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"]:
|
||||
|
||||
@@ -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)
|
||||
|
||||
+13
-12
@@ -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:
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user