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

"""Shared download routing logic.

Decides whether a download initiated from the checkpoint library should be
routed to the unet/diffusion-model roots instead of the checkpoint roots.
Used by both the download manager (at download time) and the download
routing HTTP endpoint (when the user picks a location in the UI), so the
two can never disagree.
"""

from __future__ import annotations

import logging
from typing import Iterable, Mapping, Optional, Sequence

from ..utils.constants import (
    CHECKPOINT_BASE_MODELS,
    CIVITAI_FILE_TYPE_TO_OTHER_SUB_TYPE,
    CIVITAI_TYPE_TO_OTHER_SUB_TYPE,
    DIFFUSION_MODEL_BASE_MODELS,
)

logger = logging.getLogger(__name__)

# File types reported by the CivitAI API that indicate a raw diffusion
# model (loaded via UNETLoader in ComfyUI) rather than a full checkpoint.
DIFFUSION_FILE_TYPES = frozenset({"UNet", "Diffusion Model"})

# Allowed values for the "unknown_base_model_routing" setting / the
# unknown_base_model_default parameter below.
ROUTING_DIFFUSION_MODEL = "diffusion_model"
ROUTING_CHECKPOINT = "checkpoint"

# Download-time model type -> download_root_rules folder key. The rule set is
# keyed by folder_paths keys, so a checkpoint download that resolves to the
# unet roots consults the "unet" rules, not the "checkpoints" ones.
MODEL_TYPE_TO_ROOT_RULE_KEY = {
    "lora": "loras",
    "checkpoint": "checkpoints",
    "unet": "unet",
    "embedding": "embeddings",
}

# Download-time model type -> the setting holding its default download root.
MODEL_TYPE_TO_DEFAULT_ROOT_SETTING = {
    "lora": "default_lora_root",
    "checkpoint": "default_checkpoint_root",
    "unet": "default_unet_root",
    "embedding": "default_embedding_root",
}


def resolve_download_root(
    rules: Optional[Sequence[Mapping[str, object]]],
    base_model: str,
    default_root: str,
) -> str:
    """Return the root from the first matching rule, else ``default_root``.

    Rules are ordered; the first rule whose ``match`` list contains
    ``base_model`` wins. Only ``match_type == "base_model"`` rules are
    supported (tag matching is a P1 concern).
    """
    if not rules or not base_model:
        return default_root

    for rule in rules:
        if not isinstance(rule, Mapping):
            continue
        if rule.get("match_type") != "base_model":
            continue
        match = rule.get("match")
        if not isinstance(match, Sequence) or isinstance(match, (str, bytes)):
            continue
        if base_model not in match:
            continue
        root = rule.get("root")
        if isinstance(root, str) and root.strip():
            logger.info(
                "Download root rule matched base model '%s', routing to %s",
                base_model,
                root,
            )
            return root

    return default_root


def is_diffusion_model_download(
    model_type: str,
    file_types: Iterable[str] = (),
    base_model: str = "",
    unknown_base_model_default: str = ROUTING_DIFFUSION_MODEL,
) -> bool:
    """Return True when a download should be routed to the unet roots.

    Only applies to downloads initiated from the checkpoint library.
    Priority: (1) any file has type "UNet" or "Diffusion Model" (the more
    direct signal from CivitAI), (2) baseModel is a known diffusion model,
    (3) baseModel is a known full checkpoint -> not diffusion, (4) unknown
    or empty baseModel -> the ``unknown_base_model_default`` setting, which
    defaults to diffusion because the set of true checkpoint families is
    closed while new DiT base models appear all the time.
    """
    if model_type != "checkpoint":
        return False

    for file_type in file_types:
        if file_type in DIFFUSION_FILE_TYPES:
            logger.info(
                "File type '%s' detected, routing checkpoint to unet folder",
                file_type,
            )
            return True

    if base_model in DIFFUSION_MODEL_BASE_MODELS:
        logger.info(
            "baseModel '%s' is a known diffusion model, routing to unet folder",
            base_model,
        )
        return True

    if base_model in CHECKPOINT_BASE_MODELS:
        return False

    is_diffusion = unknown_base_model_default != ROUTING_CHECKPOINT
    logger.info(
        "baseModel '%s' is unknown, routing to %s folder (unknown_base_model_routing)",
        base_model,
        "unet" if is_diffusion else "checkpoint",
    )
    return is_diffusion


def resolve_other_download_sub_type(
    civitai_model_type: str,
    file_types: Iterable[str] = (),
    selected_file_type: Optional[str] = None,
) -> Optional[str]:
    """Resolve the "other"-page sub_type for a download.

    Fixed priority (locked design, docs/plans/other-models-page.md §9.2):

    1. Explicit user file pick — when the picked file's type maps, it wins
       even when model.type maps to something else.
    2. model.type via CIVITAI_TYPE_TO_OTHER_SUB_TYPE.
    3. file.type fallback — only when model.type maps to nothing. Must NOT
       override a mapped model.type: checkpoint models routinely bundle
       VAE/Text Encoder component files.
    4. Still undecidable -> None (caller must ask the user for a folder).
    """
    if selected_file_type:
        mapped = CIVITAI_FILE_TYPE_TO_OTHER_SUB_TYPE.get(selected_file_type)
        if mapped:
            logger.info(
                "Explicit file pick type '%s' routes other download to '%s'",
                selected_file_type,
                mapped,
            )
            return mapped

    normalized_model_type = (civitai_model_type or "").strip().lower()
    mapped = CIVITAI_TYPE_TO_OTHER_SUB_TYPE.get(normalized_model_type)
    if mapped:
        return mapped

    for file_type in file_types:
        mapped = CIVITAI_FILE_TYPE_TO_OTHER_SUB_TYPE.get(file_type)
        if mapped:
            logger.info(
                "model.type '%s' unmapped; file type '%s' routes other download to '%s'",
                civitai_model_type,
                file_type,
                mapped,
            )
            return mapped

    return None