alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-09-21 03:01:27 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat(backend): add Other model type (VAE/upscaler/text encoder) scanner, service and routes

Browse Source
This commit is contained in:
Will Miao
2026-09-12 11:25:51 +08:00
parent 3070838a42
commit 27da7b3ca3
22 changed files with 1771 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/config.py
+161 -1
View File
@@ -17,6 +17,10 @@ import types as _types
import time
from .utils.cache_paths import CacheType, get_cache_file_path, get_legacy_cache_paths
from .utils.constants import (
DEFAULT_OTHER_MODEL_FOLDERS,
OTHER_MODEL_FOLDER_SUBTYPES,
)
from .utils.settings_paths import (
ensure_settings_file,
get_settings_dir,
@@ -172,6 +176,13 @@ class Config:
self.embeddings_roots = None
self.base_models_roots = self._init_checkpoint_paths()
self.embeddings_roots = self._init_embedding_paths()
# Other-model roots (VAE, upscalers, text encoders, ...): flat deduped
# list plus a normalized root -> sub_type map and per-folder_paths-key
# roots for settings persistence.
self.other_roots: Optional[List[str]] = None
self.other_root_subtypes: Dict[str, str] = {}
self.other_folder_roots: Dict[str, List[str]] = {}
self.other_roots = self._init_other_paths()
# Extra paths (only for LoRA Manager, not shared with ComfyUI)
self.extra_loras_roots: List[str] = []
self.extra_checkpoints_roots: List[str] = []
@@ -336,6 +347,10 @@ class Config:
"unet": list(self.unet_roots or []),
"embeddings": list(self.embeddings_roots or []),
}
# Persist the other-model roots under their original folder_paths
# keys so library switching round-trips them.
for key, roots in (self.other_folder_roots or {}).items():
target_folder_paths[key] = list(roots)
normalized_target_paths = _normalize_folder_paths_for_comparison(
target_folder_paths
@@ -522,6 +537,7 @@ class Config:
roots.extend(self.loras_roots or [])
roots.extend(self.base_models_roots or [])
roots.extend(self.embeddings_roots or [])
roots.extend(self.other_roots or [])
# Include extra paths for scanning symlinks
roots.extend(self.extra_loras_roots or [])
roots.extend(self.extra_checkpoints_roots or [])
@@ -862,6 +878,8 @@ class Config:
preview_roots.update(self._expand_preview_root(root))
for root in self.embeddings_roots or []:
preview_roots.update(self._expand_preview_root(root))
for root in self.other_roots or []:
preview_roots.update(self._expand_preview_root(root))
# Include extra paths for preview access
for root in self.extra_loras_roots or []:
preview_roots.update(self._expand_preview_root(root))
@@ -882,7 +900,7 @@ class Config:
path for path in preview_roots if path.is_absolute()
}
logger.debug(
"Preview roots rebuilt: %d paths from %d lora roots (%d extra), %d checkpoint roots (%d extra), %d embedding roots (%d extra), %d symlink mappings",
"Preview roots rebuilt: %d paths from %d lora roots (%d extra), %d checkpoint roots (%d extra), %d embedding roots (%d extra), %d other roots, %d symlink mappings",
len(self._preview_root_paths),
len(self.loras_roots or []),
len(self.extra_loras_roots or []),
@@ -890,6 +908,7 @@ class Config:
len(self.extra_checkpoints_roots or []),
len(self.embeddings_roots or []),
len(self.extra_embeddings_roots or []),
len(self.other_roots or []),
len(self._path_mappings),
)
@@ -1128,6 +1147,102 @@ class Config:
return unique_paths
def _get_enabled_other_folder_keys(self) -> List[str]:
"""Return the OTHER_MODEL_FOLDER_SUBTYPES keys that are enabled.
Default-enabled categories come from DEFAULT_OTHER_MODEL_FOLDERS;
opt-in categories (e.g. controlnet) are added via the
``enabled_other_folders`` setting (a list of folder_paths keys).
"""
keys = list(DEFAULT_OTHER_MODEL_FOLDERS)
try:
from .services.settings_manager import get_settings_manager
extra = get_settings_manager().get("enabled_other_folders", [])
except Exception:
extra = []
if isinstance(extra, str):
extra = [extra]
if isinstance(extra, Iterable):
for key in extra:
if (
isinstance(key, str)
and key in OTHER_MODEL_FOLDER_SUBTYPES
and key not in keys
):
keys.append(key)
return keys
def _prepare_other_paths(
self, folder_path_map: Mapping[str, Iterable[str]]
) -> Tuple[List[str], Dict[str, str], Dict[str, List[str]]]:
"""Prepare other-model paths from a folder_paths-key -> raw paths map.
Returns:
Tuple of (all_unique_roots, business_root -> sub_type map,
folder_paths key -> business roots). This method does NOT modify
instance variables - callers must set them.
"""
unique_paths: List[str] = []
sub_type_map: Dict[str, str] = {}
per_key_roots: Dict[str, List[str]] = {}
seen_real_paths: Dict[str, str] = {} # real path -> business path
# Cross-scanner overlap detection: warn when an "other" root is
# already covered by the checkpoints/unet or embeddings scanners.
# Kept (not dropped) on purpose - duplicate cards across pages are
# cosmetic, while dropping would silently unmanage the files.
covered_real_paths = {
os.path.normpath(os.path.realpath(path)).replace(os.sep, "/"): path
for path in [
*(self.base_models_roots or []),
*(self.embeddings_roots or []),
]
if isinstance(path, str) and path.strip() and os.path.exists(path)
}
for key, sub_type in OTHER_MODEL_FOLDER_SUBTYPES.items():
raw_paths = folder_path_map.get(key)
if not raw_paths:
continue
path_map = self._dedupe_existing_paths(raw_paths)
key_roots: List[str] = []
for real_path, business_path in sorted(
path_map.items(), key=lambda item: item[1].lower()
):
if real_path in seen_real_paths:
logger.warning(
"Detected the same folder '%s' under multiple other-model "
"categories ('%s' is already mapped). Keeping the first "
"category; please fix your path configuration.",
business_path,
seen_real_paths[real_path],
)
continue
seen_real_paths[real_path] = business_path
unique_paths.append(business_path)
key_roots.append(business_path)
sub_type_map[business_path] = sub_type
if real_path != business_path:
self.add_path_mapping(business_path, real_path)
covered_by = covered_real_paths.get(real_path)
if covered_by:
logger.warning(
"Detected an other-model root ('%s', category '%s') that "
"overlaps an existing checkpoints/embeddings root ('%s'). "
"The same files will appear on both pages; please review "
"your path configuration.",
business_path,
key,
covered_by,
)
if key_roots:
per_key_roots[key] = key_roots
return unique_paths, sub_type_map, per_key_roots
def _apply_library_paths(
self,
folder_paths: Mapping[str, Any],
@@ -1151,6 +1266,16 @@ class Config:
) = self._prepare_checkpoint_paths(checkpoint_paths, unet_paths)
self.embeddings_roots = self._prepare_embedding_paths(embedding_paths)
other_path_map = {
key: folder_paths.get(key, []) or []
for key in self._get_enabled_other_folder_keys()
}
(
self.other_roots,
self.other_root_subtypes,
self.other_folder_roots,
) = self._prepare_other_paths(other_path_map)
# Process extra paths (only for LoRA Manager, not shared with ComfyUI)
extra_paths = extra_folder_paths or {}
extra_lora_paths = extra_paths.get("loras", []) or []
@@ -1267,6 +1392,41 @@ class Config:
logger.warning(f"Error initializing embedding paths: {e}")
return []
def _init_other_paths(self) -> List[str]:
"""Initialize and validate other-model paths from ComfyUI settings.
Iterates the enabled OTHER_MODEL_FOLDER_SUBTYPES keys and pulls each
from ``folder_paths.get_folder_paths(key)`` (in standalone mode the
mock serves arbitrary keys from ``settings.json.folder_paths``).
"""
try:
folder_path_map: Dict[str, List[str]] = {}
for key in self._get_enabled_other_folder_keys():
try:
folder_path_map[key] = folder_paths.get_folder_paths(key)
except Exception as exc:
logger.debug("Error reading folder paths for '%s': %s", key, exc)
(
unique_paths,
self.other_root_subtypes,
self.other_folder_roots,
) = self._prepare_other_paths(folder_path_map)
logger.info(
"Found other model roots:"
+ ("\n - " + "\n - ".join(unique_paths) if unique_paths else "[]")
)
if not unique_paths:
logger.info("No valid other-model folders found in configuration")
return []
return unique_paths
except Exception as e:
logger.warning(f"Error initializing other model paths: {e}")
return []
def get_preview_static_url(self, preview_path: str) -> str:
if not preview_path:
return ""
py/lora_manager.py
+7 -1
View File
@@ -219,6 +219,7 @@ class LoraManager:
lora_scanner = await ServiceRegistry.get_lora_scanner()
checkpoint_scanner = await ServiceRegistry.get_checkpoint_scanner()
embedding_scanner = await ServiceRegistry.get_embedding_scanner()
other_scanner = await ServiceRegistry.get_other_scanner()
# Initialize recipe scanner if needed
recipe_scanner = await ServiceRegistry.get_recipe_scanner()
@@ -236,6 +237,10 @@ class LoraManager:
embedding_scanner.initialize_in_background(),
name="embedding_cache_init",
),
asyncio.create_task(
other_scanner.initialize_in_background(),
name="other_cache_init",
),
asyncio.create_task(
recipe_scanner.initialize_in_background(), name="recipe_cache_init"
),
@@ -328,6 +333,7 @@ class LoraManager:
all_roots.update(config.loras_roots)
all_roots.update(config.base_models_roots or [])
all_roots.update(config.embeddings_roots or [])
all_roots.update(config.other_roots or [])
total_deleted = 0
total_size_freed = 0
@@ -460,7 +466,7 @@ class LoraManager:
# Cancel any in-flight scanner initialization tasks so thread-pool
# workers (e.g. _initialize_cache_sync) can break out of their loops
# when the server shuts down (e.g. Ctrl+C on WSL).
for name in ("lora_scanner", "checkpoint_scanner", "embedding_scanner"):
for name in ("lora_scanner", "checkpoint_scanner", "embedding_scanner", "other_scanner"):
scanner = ServiceRegistry.get_service_sync(name)
if scanner is not None and hasattr(scanner, "cancel_task"):
scanner.cancel_task()
py/metadata_ops/__init__.py
+3 -2
View File
@@ -36,6 +36,7 @@ SCANNER_TYPE_MAP: dict[str, str] = {
"get_lora_scanner": "lora",
"get_checkpoint_scanner": "checkpoint",
"get_embedding_scanner": "embedding",
"get_other_scanner": "other",
}
SCANNER_GETTER_NAMES = tuple(SCANNER_TYPE_MAP.keys())
@@ -80,8 +81,8 @@ async def _find_scanner_for_model(
async def identify_model_type(model_path: str) -> str:
"""Determine the model type (``\"lora\"``, ``\"checkpoint\"``, or
``\"embedding\"``) for *model_path*.
"""Determine the model type (``\"lora\"``, ``\"checkpoint\"``,
``\"embedding\"``, or ``\"other\"``) for *model_path*.
Falls back to ``\"lora\"`` when unknown.
"""
py/routes/handlers/misc_handlers.py
+2
View File
@@ -658,6 +658,7 @@ class HealthCheckHandler:
"lora": ServiceRegistry.get_lora_scanner,
"checkpoint": ServiceRegistry.get_checkpoint_scanner,
"embedding": ServiceRegistry.get_embedding_scanner,
"other": ServiceRegistry.get_other_scanner,
"recipe": ServiceRegistry.get_recipe_scanner,
}
@@ -757,6 +758,7 @@ class DoctorHandler:
("lora", "LoRAs", ServiceRegistry.get_lora_scanner),
("checkpoint", "Checkpoints", ServiceRegistry.get_checkpoint_scanner),
("embedding", "Embeddings", ServiceRegistry.get_embedding_scanner),
("other", "Other Models", ServiceRegistry.get_other_scanner),
)
)
self._app_version_getter = app_version_getter
py/routes/handlers/pending_delete_handler.py
+1
View File
@@ -35,6 +35,7 @@ _MODEL_TYPE_GETTER_NAMES: Dict[str, str] = {
"loras": "get_lora_scanner",
"checkpoints": "get_checkpoint_scanner",
"embeddings": "get_embedding_scanner",
"other": "get_other_scanner",
}
# Staged batch ids are ``uuid.uuid4().hex`` (32 lowercase hex chars). The id is
py/routes/other_routes.py
+74
View File
@@ -0,0 +1,74 @@
import logging
from typing import Any, Dict
from aiohttp import web
from .base_model_routes import BaseModelRoutes
from .model_route_registrar import ModelRouteRegistrar
from ..services.other_model_service import OtherModelService
from ..services.service_registry import ServiceRegistry
from ..utils.constants import VALID_OTHER_CIVITAI_TYPES
logger = logging.getLogger(__name__)
class OtherRoutes(BaseModelRoutes):
"""Other-model-specific route controller (VAE, upscaler, text encoder, ...)"""
def __init__(self):
"""Initialize Other-model routes with OtherModel service"""
super().__init__()
self.template_name = "other.html"
async def initialize_services(self):
"""Initialize services from ServiceRegistry"""
other_scanner = await ServiceRegistry.get_other_scanner()
update_service = await ServiceRegistry.get_model_update_service()
self.service = OtherModelService(other_scanner, update_service=update_service)
self.set_model_update_service(update_service)
# Attach service dependencies
self.attach_service(self.service)
def setup_routes(self, app: web.Application, prefix: str = "other"):
"""Setup Other-model routes"""
# Schedule service initialization on app startup
app.on_startup.append(lambda _: self.initialize_services())
# Setup common routes with 'other' prefix (includes page route)
super().setup_routes(app, prefix)
def setup_specific_routes(self, registrar: ModelRouteRegistrar, prefix: str):
"""Setup Other-model-specific routes"""
# Other-model info by name
registrar.add_prefixed_route('GET', '/api/lm/{prefix}/info/{name}', prefix, self.get_other_model_info)
def _validate_civitai_model_type(self, model_type: str) -> bool:
"""Validate CivitAI model type for other models.
Accepts retired CivitAI types (CLIP, CLIPVision) as well — grandfathered
models on CivitAI still carry them.
"""
return model_type.lower() in VALID_OTHER_CIVITAI_TYPES
def _get_expected_model_types(self) -> str:
"""Get expected model types string for error messages"""
return "VAE, Upscaler, TextEncoder, CLIPVision, Controlnet, or Other"
def _parse_specific_params(self, request: web.Request) -> Dict[str, Any]:
"""Parse other-model-specific parameters (none in Phase 1)."""
return {}
async def get_other_model_info(self, request: web.Request) -> web.Response:
"""Get detailed information for a specific other model by name"""
try:
name = request.match_info.get('name', '')
model_info = await self.service.get_model_info_by_name(name) # pyright: ignore[reportAttributeAccessIssue]
if model_info:
return web.json_response(model_info)
else:
return web.json_response({"error": "Model not found"}, status=404)
except Exception as e:
logger.error(f"Error in get_other_model_info: {e}", exc_info=True)
return web.json_response({"error": str(e)}, status=500)
py/services/base_model_service.py
+6 -1
View File
@@ -7,7 +7,7 @@ import logging
import os
import time
from ..utils.constants import VALID_LORA_SUB_TYPES, VALID_CHECKPOINT_SUB_TYPES
from ..utils.constants import VALID_LORA_SUB_TYPES, VALID_CHECKPOINT_SUB_TYPES, VALID_OTHER_SUB_TYPES
from ..utils.models import BaseModelMetadata
from ..utils.metadata_manager import MetadataManager
from ..utils.usage_stats import UsageStats
@@ -904,6 +904,11 @@ class BaseModelService(ABC):
and normalized_type not in VALID_CHECKPOINT_SUB_TYPES
):
continue
if (
self.model_type == "other"
and normalized_type not in VALID_OTHER_SUB_TYPES
):
continue
type_counts[normalized_type] = type_counts.get(normalized_type, 0) + 1
py/services/model_scanner.py
+1
View File
@@ -68,6 +68,7 @@ PAGE_TYPE_MAP = {
'lora': 'loras',
'checkpoint': 'checkpoints',
'embedding': 'embeddings',
'other': 'other',
}
py/services/model_service_factory.py
+10 -5
View File
@@ -118,19 +118,24 @@ class ModelServiceFactory:
def register_default_model_types():
"""Register the default model types (LoRA, Checkpoint, and Embedding)"""
"""Register the default model types (LoRA, Checkpoint, Embedding, and Other)"""
from ..services.lora_service import LoraService
from ..services.checkpoint_service import CheckpointService
from ..services.embedding_service import EmbeddingService
from ..services.other_model_service import OtherModelService
from ..routes.lora_routes import LoraRoutes
from ..routes.checkpoint_routes import CheckpointRoutes
from ..routes.embedding_routes import EmbeddingRoutes
from ..routes.other_routes import OtherRoutes
# Register LoRA model type
ModelServiceFactory.register_model_type('lora', LoraService, LoraRoutes)
# Register Checkpoint model type
ModelServiceFactory.register_model_type('checkpoint', CheckpointService, CheckpointRoutes)
# Register Embedding model type
ModelServiceFactory.register_model_type('embedding', EmbeddingService, EmbeddingRoutes)
ModelServiceFactory.register_model_type('embedding', EmbeddingService, EmbeddingRoutes)
# Register Other model type (VAE, upscaler, text encoder, ...)
ModelServiceFactory.register_model_type('other', OtherModelService, OtherRoutes)
py/services/other_model_service.py
+79
View File
@@ -0,0 +1,79 @@
import os
import logging
from typing import Any, Dict, Optional
from .base_model_service import BaseModelService
from .auto_tag_service import extract_auto_tags
from ..utils.models import OtherModelMetadata
from ..config import config
logger = logging.getLogger(__name__)
class OtherModelService(BaseModelService):
"""Other-model-specific service implementation (VAE, upscaler, text encoder, ...)"""
def __init__(self, scanner, update_service=None):
"""Initialize Other-model service
Args:
scanner: Other-model scanner instance
update_service: Optional service for remote update tracking.
"""
super().__init__("other", scanner, OtherModelMetadata, update_service=update_service)
async def format_response(self, model_data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""Format other-model 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 = model_data.get("file_path")
if not file_path or not isinstance(file_path, str):
logger.warning(
"Skipping corrupted other-model entry (missing file_path): %s",
model_data.get("file_name", "<unknown>"),
)
return None
# Get sub_type from cache entry (new canonical field)
sub_type = model_data.get("sub_type", "vae")
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(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": model_data.get("sha256", ""),
"autov3": model_data.get("autov3"),
"file_path": file_path.replace(os.sep, "/"),
"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),
"notes": model_data.get("notes", ""),
"sub_type": sub_type,
"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[str, Any]:
"""Find other models with duplicate SHA256 hashes"""
return self.scanner._hash_index.get_duplicate_hashes()
def find_duplicate_filenames(self) -> Dict[str, Any]:
"""Find other models with conflicting filenames"""
return self.scanner._hash_index.get_duplicate_filenames()
py/services/other_scanner.py
+468
View File
@@ -0,0 +1,468 @@
# 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
import os
from datetime import datetime
from typing import Any, Dict, List, Optional
from ..utils.models import OtherModelMetadata
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, _is_excluded_dir
from .model_hash_index import ModelHashIndex
logger = logging.getLogger(__name__)
class OtherScanner(ModelScanner):
"""Service for scanning and managing "other" model files.
Aggregates every enabled folder_paths category from
OTHER_MODEL_FOLDER_SUBTYPES (VAE, upscalers, text encoders, CLIP vision,
opt-in ControlNet) into one scanner; sub_type is derived from the root
containing the file (mirrors CheckpointScanner's checkpoints/unet split).
Hashing is lazy (checkpoint-style): text encoders can be ~10 GB, so the
initial scan records hash_status="pending" and the SHA256 is computed
on-demand via calculate_hash_for_model (e.g. when fetching CivitAI
metadata).
"""
def __init__(self):
# Same extension set as CheckpointScanner (ComfyUI's
# supported_pt_extensions plus ".gguf").
file_extensions = {
".ckpt",
".pt",
".pt2",
".bin",
".pth",
".safetensors",
".pkl",
".sft",
".gguf",
}
super().__init__(
model_type="other",
model_class=OtherModelMetadata,
file_extensions=file_extensions,
hash_index=ModelHashIndex(),
)
if not hasattr(self, "_hash_calculation_lock"):
self._hash_calculation_lock = asyncio.Lock()
self._hash_calculation_tasks: dict[str, asyncio.Task[Optional[str]]] = {}
async def _create_default_metadata(
self, file_path: str
) -> Optional[OtherModelMetadata]:
"""Create default metadata without calculating hash (lazy hash).
Other models include multi-GB text encoders, so hash calculation is
deferred until on-demand (e.g. CivitAI metadata fetch).
"""
try:
real_path = os.path.realpath(file_path)
if not os.path.exists(real_path):
logger.error(f"File not found: {file_path}")
return None
base_name = os.path.splitext(os.path.basename(file_path))[0]
dir_path = os.path.dirname(file_path)
# 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 files; record the checked state at creation time ("" =
# checked but unavailable).
autov3 = calculate_autov3(real_path)
# Create metadata WITHOUT calculating hash
metadata = OtherModelMetadata(
file_name=base_name,
model_name=base_name,
file_path=normalize_path(file_path),
size=os.path.getsize(real_path),
modified=datetime.now().timestamp(),
sha256="", # Empty hash - will be calculated on-demand
base_model="Unknown",
preview_url=normalize_path(preview_url),
tags=[],
modelDescription="",
sub_type=self.resolve_sub_type_for_path(file_path) or "vae",
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
logger.info(f"Creating other-model metadata (hash pending) for {file_path}")
await MetadataManager.save_metadata(file_path, metadata)
return metadata
except Exception as e:
logger.error(
f"Error creating default other-model metadata for {file_path}: {e}"
)
return None
async def calculate_hash_for_model(self, file_path: str) -> Optional[str]:
"""Calculate hash for a model on-demand with per-file singleflight.
Args:
file_path: Path to the model file
Returns:
SHA256 hash string, or None if calculation failed
"""
try:
real_path = os.path.realpath(file_path)
if not os.path.exists(real_path):
logger.error(f"File not found for hash calculation: {file_path}")
return None
metadata, _ = await MetadataManager.load_metadata(
file_path, self.model_class
)
if (
metadata is not None
and metadata.hash_status == "completed"
and metadata.sha256
):
# Ensure the in-memory hash index is populated even when
# the hash was already computed and persisted to the metadata
# file. Without this, usage tracking (and any other caller
# 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,
getattr(metadata, "autov3", None) or None,
)
return metadata.sha256
async with self._hash_calculation_lock:
metadata, _ = await MetadataManager.load_metadata(
file_path, self.model_class
)
if (
metadata is not None
and metadata.hash_status == "completed"
and metadata.sha256
):
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)
if task is None:
task = asyncio.create_task(
self._run_hash_calculation_task(file_path, real_path)
)
self._hash_calculation_tasks[real_path] = task
return await asyncio.shield(task)
except Exception as e:
logger.error(f"Error calculating hash for {file_path}: {e}")
return None
async def _run_hash_calculation_task(
self, file_path: str, real_path: str
) -> Optional[str]:
"""Run a hash calculation task and remove it from the in-flight map."""
try:
return await self._calculate_hash_for_model_uncached(file_path, real_path)
finally:
task = asyncio.current_task()
async with self._hash_calculation_lock:
if self._hash_calculation_tasks.get(real_path) is task:
del self._hash_calculation_tasks[real_path]
async def _calculate_hash_for_model_uncached(
self, file_path: str, real_path: str
) -> Optional[str]:
"""Calculate hash for a model without checking in-flight tasks."""
from ..utils.file_utils import calculate_sha256
try:
# Load current metadata
metadata, should_skip = await MetadataManager.load_metadata(
file_path, self.model_class
)
if metadata is None:
if should_skip:
logger.error(f"Invalid metadata found for {file_path}")
return None
created_metadata = await self._create_default_metadata(file_path)
if created_metadata is None:
logger.error(f"No metadata found for {file_path}")
return None
metadata = created_metadata
# Check if hash is already calculated
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,
getattr(metadata, "autov3", None) or None,
)
return metadata.sha256
# Update status to calculating
metadata.hash_status = "calculating"
await MetadataManager.save_metadata(file_path, metadata)
# Calculate hash
logger.info(f"Calculating hash for other model: {file_path}")
sha256 = await calculate_sha256(real_path)
# Update metadata with hash
metadata.sha256 = sha256
metadata.hash_status = "completed"
await MetadataManager.save_metadata(file_path, metadata)
# Update hash index
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
# write the hash back to the SQLite models table. Without
# this the hash only lives in the metadata file and the
# in-memory hash index, both of which are lost across
# restarts, causing the same re-computation loop on the
# next session.
if self._cache is not None and self._cache.raw_data:
for entry in self._cache.raw_data:
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 other model: {file_path}")
return sha256
except Exception as e:
logger.error(f"Error calculating hash for {file_path}: {e}")
# Update status to failed
try:
metadata, _ = await MetadataManager.load_metadata(
file_path, self.model_class
)
if metadata:
metadata.hash_status = "failed"
await MetadataManager.save_metadata(file_path, metadata)
except Exception:
pass
return None
async def calculate_all_pending_hashes(
self, progress_callback=None
) -> Dict[str, int]:
"""Calculate hashes for all other models with pending hash status.
If cache is not initialized, scans filesystem directly for metadata files
with hash_status != 'completed'.
Args:
progress_callback: Optional callback(progress, total, current_file)
Returns:
Dict with 'completed', 'failed', 'total' counts
"""
# Try to get from cache first
cache = await self.get_cached_data()
if cache and cache.raw_data:
# Use cache if available
pending_models = [
item
for item in cache.raw_data
if item.get("hash_status") != "completed" or not item.get("sha256")
]
else:
# Cache not initialized, scan filesystem directly
pending_models = await self._find_pending_models_from_filesystem()
if not pending_models:
return {"completed": 0, "failed": 0, "total": 0}
total = len(pending_models)
completed = 0
failed = 0
for i, model_data in enumerate(pending_models):
file_path = model_data.get("file_path")
if not file_path:
continue
try:
sha256 = await self.calculate_hash_for_model(file_path)
if sha256:
completed += 1
else:
failed += 1
except Exception as e:
logger.error(f"Error calculating hash for {file_path}: {e}")
failed += 1
if progress_callback:
try:
await progress_callback(i + 1, total, file_path)
except Exception:
pass
return {"completed": completed, "failed": failed, "total": total}
async def _find_pending_models_from_filesystem(self) -> List[Dict[str, Any]]:
"""Scan filesystem for other-model metadata files with pending hash status."""
pending_models = []
for root_path in self.get_model_roots():
if not os.path.exists(root_path):
continue
for dirpath, dirnames, filenames in os.walk(root_path):
dirnames[:] = [d for d in dirnames if not _is_excluded_dir(d)]
for filename in filenames:
if not filename.endswith(".metadata.json"):
continue
metadata_path = os.path.join(dirpath, filename)
try:
with open(metadata_path, "r", encoding="utf-8") as f:
data = json.load(f)
# Check if hash is pending
hash_status = data.get("hash_status", "completed")
sha256 = data.get("sha256", "")
if hash_status != "completed" or not sha256:
# Find corresponding model file
model_name = filename.replace(".metadata.json", "")
model_path = None
# Look for model file with matching name
for ext in self.file_extensions:
potential_path = os.path.join(dirpath, model_name + ext)
if os.path.exists(potential_path):
model_path = potential_path
break
if model_path:
pending_models.append(
{
"file_path": model_path.replace(os.sep, "/"),
"hash_status": hash_status,
"sha256": sha256,
**{
k: v
for k, v in data.items()
if k
not in [
"file_path",
"hash_status",
"sha256",
]
},
}
)
except (json.JSONDecodeError, Exception) as e:
logger.debug(
f"Error reading metadata file {metadata_path}: {e}"
)
continue
return pending_models
def _root_sub_type_map(self) -> Dict[str, str]:
"""Return the configured business root -> sub_type map."""
root_map = getattr(config, "other_root_subtypes", None)
return root_map if isinstance(root_map, dict) else {}
def _resolve_sub_type(self, root_path: Optional[str]) -> Optional[str]:
"""Resolve the sub_type for a configured root path."""
if not root_path:
return None
normalized_root = self._normalize_path_value(root_path)
for root, sub_type in self._root_sub_type_map().items():
if self._normalize_path_value(root) == normalized_root:
return sub_type
return None
def resolve_sub_type_for_path(self, file_path: Optional[str]) -> Optional[str]:
"""Resolve sub_type from the configured root that contains the file.
Uses the longest-prefix match so nested roots (e.g. a controlnet root
inside a vae root) resolve to the most specific category.
"""
normalized_path = self._normalize_path_value(file_path)
if not normalized_path:
return None
best_length = 0
best_sub_type: Optional[str] = None
for root, sub_type in self._root_sub_type_map().items():
normalized_root = self._normalize_path_value(root)
if not normalized_root:
continue
if (
normalized_path == normalized_root
or normalized_path.startswith(f"{normalized_root}/")
) and len(normalized_root) > best_length:
best_length = len(normalized_root)
best_sub_type = sub_type
return best_sub_type
def adjust_metadata(self, metadata, file_path, root_path):
"""Adjust metadata during scanning to set sub_type."""
sub_type = self._resolve_sub_type(root_path) or self.resolve_sub_type_for_path(
file_path
)
if sub_type:
metadata.sub_type = sub_type
return metadata
def adjust_cached_entry(self, entry: Dict[str, Any]) -> Dict[str, Any]:
"""Adjust entries loaded from the persisted cache to ensure sub_type is set.
sub_type is location-derived: it is re-derived on cache load, never
trusted from the persisted snapshot.
"""
sub_type = self.resolve_sub_type_for_path(entry.get("file_path"))
if sub_type:
entry["sub_type"] = sub_type
return entry
def get_model_roots(self) -> List[str]:
"""Get other-model root directories"""
roots: List[str] = []
roots.extend(config.other_roots or [])
# Remove duplicates while preserving order
seen: set[str] = set()
unique_roots: List[str] = []
for root in roots:
if root and root not in seen:
seen.add(root)
unique_roots.append(root)
return unique_roots
py/services/pending_delete_service.py
+2
View File
@@ -59,6 +59,7 @@ _MODEL_TYPE_PAGE_MAP = {
"lora": "loras",
"checkpoint": "checkpoints",
"embedding": "embeddings",
"other": "other",
}
# Module-level alias so tests can spy on timer task creation without patching
@@ -983,6 +984,7 @@ class PendingDeleteService:
"get_lora_scanner",
"get_checkpoint_scanner",
"get_embedding_scanner",
"get_other_scanner",
):
getter = getattr(ServiceRegistry, getter_name, None)
if not callable(getter):
py/services/service_registry.py
+26 -5
View File
@@ -297,23 +297,44 @@ class ServiceRegistry:
async def get_embedding_scanner(cls):
"""Get or create Embedding scanner instance"""
service_name = "embedding_scanner"
if service_name in cls._services:
return cls._services[service_name]
async with cls._get_lock(service_name):
# Double-check after acquiring lock
if service_name in cls._services:
return cls._services[service_name]
# Import here to avoid circular imports
from .embedding_scanner import EmbeddingScanner
scanner = await EmbeddingScanner.get_instance()
cls._services[service_name] = scanner
logger.debug(f"Created and registered {service_name}")
return scanner
@classmethod
async def get_other_scanner(cls):
"""Get or create Other-model scanner instance"""
service_name = "other_scanner"
if service_name in cls._services:
return cls._services[service_name]
async with cls._get_lock(service_name):
# Double-check after acquiring lock
if service_name in cls._services:
return cls._services[service_name]
# Import here to avoid circular imports
from .other_scanner import OtherScanner
scanner = await OtherScanner.get_instance()
cls._services[service_name] = scanner
logger.debug(f"Created and registered {service_name}")
return scanner
@classmethod
def clear_services(cls):
"""Clear all registered services - mainly for testing"""
py/utils/constants.py
+43
View File
@@ -83,6 +83,49 @@ VALID_LORA_SUB_TYPES = ["lora", "locon", "dora"]
VALID_CHECKPOINT_SUB_TYPES = ["checkpoint", "diffusion_model"]
VALID_EMBEDDING_SUB_TYPES = ["embedding"]
# folder_paths key -> sub_type; single source of truth for extensibility.
# Adding support for a new ComfyUI folder category is a one-line change here.
OTHER_MODEL_FOLDER_SUBTYPES = {
"vae": "vae",
"upscale_models": "upscaler",
"text_encoders": "text_encoder",
"clip": "text_encoder", # legacy ComfyUI key
"clip_vision": "clip_vision",
"controlnet": "controlnet",
}
# folder_paths keys scanned by default; anything else in
# OTHER_MODEL_FOLDER_SUBTYPES (e.g. controlnet) is opt-in via the
# "enabled_other_folders" setting.
DEFAULT_OTHER_MODEL_FOLDERS = (
"vae",
"upscale_models",
"text_encoders",
"clip",
"clip_vision",
)
VALID_OTHER_SUB_TYPES = ["vae", "upscaler", "text_encoder", "clip_vision", "controlnet"]
# CivitAI model.type values accepted by the "other" page's fetch-metadata
# validation (lowercased). CLIP/CLIPVision are retired upstream but still
# appear on grandfathered models.
VALID_OTHER_CIVITAI_TYPES = {
"vae",
"upscaler",
"textencoder",
"clip",
"clipvision",
"controlnet",
"other",
}
# CivitAI model.type -> internal sub_type for the "other" model page.
CIVITAI_TYPE_TO_OTHER_SUB_TYPE = {
"vae": "vae",
"upscaler": "upscaler",
"textencoder": "text_encoder",
"clip": "text_encoder",
"clipvision": "clip_vision",
"controlnet": "controlnet",
}
# Backward compatibility alias
VALID_LORA_TYPES = VALID_LORA_SUB_TYPES
py/utils/models.py
+55 -1
View File
@@ -2,7 +2,7 @@ from dataclasses import dataclass, asdict, field
from typing import Callable, Dict, Optional, List, Any
from datetime import datetime
import os
from .constants import INVALID_AUTOV3_EMPTY_HASH
from .constants import CIVITAI_TYPE_TO_OTHER_SUB_TYPE, INVALID_AUTOV3_EMPTY_HASH
from .model_utils import determine_base_model
@@ -318,6 +318,60 @@ class CheckpointMetadata(BaseModelMetadata):
)
@dataclass
class OtherModelMetadata(BaseModelMetadata):
"""Represents the metadata structure for an "other" model (VAE, upscaler,
text encoder, CLIP vision, ControlNet, ...).
The sub_type is location-derived: the OtherScanner sets it from the
folder_paths category whose root contains the file. The dataclass default
is only a placeholder.
"""
sub_type: str = "vae" # Placeholder; overridden by the scanner hooks
@classmethod
def from_civitai_info(
cls, version_info: Dict[str, Any], file_info: Dict[str, Any], save_path: str
) -> "OtherModelMetadata":
"""Create OtherModelMetadata instance from Civitai version info"""
file_name = file_info.get("name", "")
base_model = determine_base_model(version_info.get("baseModel", ""))
sha256_value = (file_info.get("hashes") or {}).get("SHA256", "").lower()
# Map the CivitAI model type onto our sub_types; unknown types keep the
# placeholder until the scanner re-derives sub_type from the location.
civitai_type = str(version_info.get("type", "") or "").lower()
sub_type = CIVITAI_TYPE_TO_OTHER_SUB_TYPE.get(civitai_type, "vae")
# Extract tags and description if available
tags = []
description = ""
model_data = version_info.get("model") or {}
if "tags" in model_data:
tags = model_data["tags"]
if "description" in model_data:
description = model_data["description"]
return cls(
file_name=os.path.splitext(file_name)[0],
model_name=model_data.get("name", os.path.splitext(file_name)[0]),
file_path=save_path.replace(os.sep, "/"),
size=file_info.get("sizeKB", 0) * 1024,
modified=datetime.now().timestamp(),
sha256=sha256_value,
base_model=base_model,
preview_url="", # Will be updated after preview download
preview_nsfw_level=0,
from_civitai=True,
civitai=version_info,
sub_type=sub_type,
tags=tags,
modelDescription=description,
# Direct read: the downloaded file IS file_info, no SHA256 matching.
autov3=normalize_autov3((file_info.get("hashes") or {}).get("AutoV3")),
)
@dataclass
class EmbeddingMetadata(BaseModelMetadata):
"""Represents the metadata structure for an Embedding model"""