ComfyUI-Lora-Manager/py/nodes/random_checkpoint_loader.py

import logging
import os
import random
from typing import Any, List, Optional, 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__)


class RandomCheckpointLoaderLM:
    """Checkpoint Loader that can randomly pick a checkpoint from the pool

    Loads checkpoints from both standard ComfyUI folders and LoRA Manager's
    extra folder paths. When select_at_random is enabled, ignores ckpt_name
    and picks a random checkpoint (optionally filtered by base_model) on
    every run.
    """

    NAME = "Random Checkpoint Loader (LoraManager)"
    CATEGORY = "Lora Manager/loaders"

    @classmethod
    def INPUT_TYPES(cls):
        # Get list of checkpoint names from scanner (includes extra folder paths)
        checkpoint_names = cls._get_checkpoint_names()
        base_models = cls._get_available_base_models()
        return {
            "required": {
                "ckpt_name": (
                    checkpoint_names,
                    {"tooltip": "The name of the checkpoint (model) to load."},
                ),
                "select_at_random": (
                    "BOOLEAN",
                    {
                        "default": False,
                        "tooltip": (
                            "Ignore ckpt_name and pick a random checkpoint from the "
                            "pool (optionally filtered by base_model) on every run."
                        ),
                    },
                ),
                "base_model": (
                    base_models,
                    {
                        "default": "Any",
                        "tooltip": "Restrict random selection to this base model. 'Any' uses the full pool.",
                    },
                ),
            }
        }

    RETURN_TYPES = ("MODEL", "CLIP", "VAE", "STRING")
    RETURN_NAMES = ("MODEL", "CLIP", "VAE", "model_name")
    OUTPUT_TOOLTIPS = (
        "The model used for denoising latents.",
        "The CLIP model used for encoding text prompts.",
        "The VAE model used for encoding and decoding images to and from latent space.",
        "The name of the checkpoint that was loaded (useful when select_at_random is enabled).",
    )
    FUNCTION = "load_checkpoint"

    @classmethod
    def IS_CHANGED(cls, ckpt_name, select_at_random=False, base_model="Any"):
        # Force re-execution on every run while randomizing, since the widget
        # values themselves don't change between queue runs.
        if select_at_random:
            return float("nan")
        return ckpt_name

    @staticmethod
    def _run_async(coro_fn):
        """Run an async fetcher, handling the case where an event loop is already running."""
        import asyncio

        try:
            asyncio.get_running_loop()
            import concurrent.futures

            def run_in_thread():
                new_loop = asyncio.new_event_loop()
                asyncio.set_event_loop(new_loop)
                try:
                    return new_loop.run_until_complete(coro_fn())
                finally:
                    new_loop.close()

            with concurrent.futures.ThreadPoolExecutor() as executor:
                future = executor.submit(run_in_thread)
                return future.result()
        except RuntimeError:
            return asyncio.run(coro_fn())

    @classmethod
    def _get_checkpoint_names(cls, base_model: Optional[str] = None) -> List[str]:
        """Get list of checkpoint names from scanner cache in ComfyUI format (relative path with extension)

        Args:
            base_model: If given (and not "Any"), only include checkpoints matching this base model.
        """
        try:
            from ..services.service_registry import ServiceRegistry

            async def _get_names():
                scanner = await ServiceRegistry.get_checkpoint_scanner()
                cache = await scanner.get_cached_data()

                # Get all model roots for calculating relative paths
                model_roots = scanner.get_model_roots()

                # Filter only checkpoint type (not diffusion_model) and format names
                names = []
                for item in cache.raw_data:
                    if item.get("sub_type") != "checkpoint":
                        continue
                    if (
                        base_model
                        and base_model != "Any"
                        and item.get("base_model") != base_model
                    ):
                        continue
                    file_path = item.get("file_path", "")
                    # Only offer models that still exist on disk so ComfyUI
                    # flags missing checkpoints at queue time via
                    # "value not in list" (the scanner cache can be stale).
                    if file_path and os.path.exists(file_path):
                        # Format using relative path with OS-native separator
                        formatted_name = _format_model_name_for_comfyui(
                            file_path, model_roots
                        )
                        if formatted_name:
                            names.append(formatted_name)

                return sorted(names)

            return cls._run_async(_get_names)
        except Exception as e:
            logger.error(f"Error getting checkpoint names: {e}")
            return []

    @classmethod
    def _get_available_base_models(cls) -> List[str]:
        """Get distinct base_model values present among indexed checkpoints, for the random-selection filter."""
        try:
            from ..services.service_registry import ServiceRegistry

            async def _get_base_models():
                scanner = await ServiceRegistry.get_checkpoint_scanner()
                cache = await scanner.get_cached_data()

                base_models = set()
                for item in cache.raw_data:
                    if item.get("sub_type") != "checkpoint":
                        continue
                    base_model = item.get("base_model")
                    file_path = item.get("file_path", "")
                    if base_model and file_path and os.path.exists(file_path):
                        base_models.add(base_model)

                return sorted(base_models)

            return ["Any"] + cls._run_async(_get_base_models)
        except Exception as e:
            logger.error(f"Error getting available base models: {e}")
            return ["Any"]

    def load_checkpoint(
        self,
        ckpt_name: str,
        select_at_random: bool = False,
        base_model: str = "Any",
    ) -> Tuple[Any, Any, Any]:
        """Load a checkpoint by name, supporting extra folder paths

        Args:
            ckpt_name: The name of the checkpoint to load (relative path with extension)
            select_at_random: If True, ignore ckpt_name and pick randomly from the pool
            base_model: Restricts random selection to this base model ("Any" = no filter)

        Returns:
            Tuple of (MODEL, CLIP, VAE, model_name)
        """
        if select_at_random:
            pool = self._get_checkpoint_names(base_model)
            if not pool:
                raise FileNotFoundError(
                    f"No checkpoints found for base model '{base_model}'. "
                    "Pick a different base model or disable 'select_at_random'."
                )
            ckpt_name = random.choice(pool)
            logger.info(
                f"[RandomCheckpointLoaderLM] Randomly selected checkpoint: {ckpt_name}"
            )

        # Get absolute path from cache using ComfyUI-style name
        ckpt_path, metadata = get_checkpoint_info_absolute(ckpt_name)

        if metadata is None:
            raise FileNotFoundError(
                f"Checkpoint '{ckpt_name}' not found in LoRA Manager cache. "
                "Make sure the checkpoint is indexed and try again."
            )

        # Load regular checkpoint using ComfyUI's API
        logger.info(f"Loading checkpoint from: {ckpt_path}")
        out = comfy.sd.load_checkpoint_guess_config(
            ckpt_path,
            output_vae=True,
            output_clip=True,
            embedding_directory=folder_paths.get_folder_paths("embeddings"),
        )
        return out[:3] + (ckpt_name,)