ComfyUI-Lora-Manager/py/services/checkpoint_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 CheckpointMetadata
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
from .model_hash_index import ModelHashIndex

logger = logging.getLogger(__name__)


class CheckpointScanner(ModelScanner):
    """Service for scanning and managing checkpoint files"""

    def __init__(self):
        # Define supported file extensions
        file_extensions = {
            ".ckpt",
            ".pt",
            ".pt2",
            ".bin",
            ".pth",
            ".safetensors",
            ".pkl",
            ".sft",
            ".gguf",
        }
        super().__init__(
            model_type="checkpoint",
            model_class=CheckpointMetadata,
            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[CheckpointMetadata]:
        """Create default metadata for checkpoint without calculating hash (lazy hash).

        Checkpoints are typically large (10GB+), so we skip hash calculation during initial
        scanning to improve startup performance. Hash will be calculated on-demand when
        fetching metadata from Civitai.
        """
        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 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,
                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="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
            logger.info(f"Creating checkpoint 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 checkpoint metadata for {file_path}: {e}"
            )
            return None

    async def calculate_hash_for_model(self, file_path: str) -> Optional[str]:
        """Calculate hash for a checkpoint 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 checkpoint 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 checkpoint: {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"
                        break

            logger.info(f"Hash calculated for checkpoint: {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 checkpoints 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 checkpoint 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):
                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 _resolve_sub_type(self, root_path: Optional[str]) -> Optional[str]:
        """Resolve the sub-type based on the root path.

        Checks both standard ComfyUI paths and LoRA Manager's extra folder paths.
        """
        if not root_path:
            return None

        # Check standard ComfyUI checkpoint paths
        if config.checkpoints_roots and root_path in config.checkpoints_roots:
            return "checkpoint"

        # Check extra checkpoint paths
        if (
            config.extra_checkpoints_roots
            and root_path in config.extra_checkpoints_roots
        ):
            return "checkpoint"

        # Check standard ComfyUI unet paths
        if config.unet_roots and root_path in config.unet_roots:
            return "diffusion_model"

        # Check extra unet paths
        if config.extra_unet_roots and root_path in config.extra_unet_roots:
            return "diffusion_model"

        return None

    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)
        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 = self._resolve_sub_type(
            self._find_root_for_file(entry.get("file_path"))
        )
        if sub_type:
            entry["sub_type"] = sub_type
        return entry

    def get_model_roots(self) -> List[str]:
        """Get checkpoint root directories (including extra paths)"""
        roots: List[str] = []
        roots.extend(config.base_models_roots or [])
        roots.extend(config.extra_checkpoints_roots or [])
        roots.extend(config.extra_unet_roots or [])
        # Remove duplicates while preserving order
        seen: set[str] = set()
        unique_roots: List[str] = []
        for root in roots:
            if root not in seen:
                seen.add(root)
                unique_roots.append(root)
        return unique_roots