# 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 _should_keep_cached_entry(self, entry: Dict[str, Any]) -> bool: """Drop persisted entries whose folder is no longer a managed root. sub_type is location-derived and config only maps enabled roots, so a file under a disabled sub_type - or under any other root while the feature is off - resolves to None here and is filtered out while the persisted cache is hydrated. """ return self.resolve_sub_type_for_path(entry.get("file_path")) is not None 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