ComfyUI-Lora-Manager/py/services/other_scanner.py

# 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