AnyImage/src/snapform/models.py

from dataclasses import dataclass


@dataclass(frozen=True)
class ModelSpec:
    key: str
    label: str
    directory_name: str
    repository: str
    license_name: str
    description: str
    supports_gs: bool
    download_patterns: tuple[str, ...] = ()


MODELS = {
    "SMALL": ModelSpec(
        key="SMALL",
        label="DA3 Small",
        directory_name="da3_small",
        repository="depth-anything/DA3-SMALL",
        license_name="Apache 2.0",
        description="Fastest any-view model with Apache 2.0 weights",
        supports_gs=False,
    ),
    "BASE": ModelSpec(
        key="BASE",
        label="DA3 Base",
        directory_name="da3_base",
        repository="depth-anything/DA3-BASE",
        license_name="Apache 2.0",
        description="Balanced any-view model with Apache 2.0 weights",
        supports_gs=False,
    ),
    "LARGE_1_1": ModelSpec(
        key="LARGE_1_1",
        label="DA3 Large 1.1",
        directory_name="da3_large_1_1",
        repository="depth-anything/DA3-LARGE-1.1",
        license_name="CC BY-NC 4.0",
        description="Higher quality non-commercial model with larger VRAM use",
        supports_gs=False,
    ),
    "GIANT_1_1": ModelSpec(
        key="GIANT_1_1",
        label="DA3 Giant 1.1",
        directory_name="da3_giant_1_1",
        repository="depth-anything/DA3-GIANT-1.1",
        license_name="CC BY-NC 4.0",
        description="Non-commercial Giant model with Gaussian head",
        supports_gs=True,
    ),
    "NESTED_GIANT_LARGE_1_1": ModelSpec(
        key="NESTED_GIANT_LARGE_1_1",
        label="DA3 Nested Giant Large 1.1",
        directory_name="da3_nested_giant_large_1_1",
        repository="depth-anything/DA3NESTED-GIANT-LARGE-1.1",
        license_name="CC BY-NC 4.0",
        description="Non-commercial nested model with metric depth and Gaussian head",
        supports_gs=True,
    ),
}

BEN2_MODEL = ModelSpec(
    key="BEN2_BASE",
    label="BEN2 Base",
    directory_name="ben2_base",
    repository="PramaLLC/BEN2",
    license_name="MIT",
    description="Fast high-quality background removal model",
    supports_gs=False,
    download_patterns=("config.json", "model.safetensors"),
)
DOWNLOADABLE_MODELS = {**MODELS, BEN2_MODEL.key: BEN2_MODEL}
MODEL_NUMBERS = {
    model_key: number for number, model_key in enumerate(MODELS)
}


_ENUM_ITEMS_CACHE = {}


def enum_items(supports_gs=None, show_download_status=False):
    model_ready = None
    if show_download_status:
        from .runtime import model_ready

    models = [
        model
        for model in MODELS.values()
        if supports_gs is None or model.supports_gs == supports_gs
    ]
    download_status = (
        tuple(model_ready(model.key) for model in models)
        if model_ready is not None
        else ()
    )
    cache_key = (
        supports_gs,
        show_download_status,
        download_status,
    )
    cached_items = _ENUM_ITEMS_CACHE.get(cache_key)
    if cached_items is not None:
        return cached_items

    items = []
    for index, model in enumerate(models):
        label = model.label
        if download_status and not download_status[index]:
            label = f"{label} (Missing)"
        items.append(
            (
                model.key,
                label,
                f"{model.description}; license: {model.license_name}",
                MODEL_NUMBERS[model.key],
            )
        )

    _ENUM_ITEMS_CACHE[cache_key] = items
    return items


def model_enum_items(_owner, _context):
    return enum_items(show_download_status=True)


def gs_model_enum_items(_owner, _context):
    return enum_items(supports_gs=True, show_download_status=True)


def selected_model_key(settings, export_format=None):
    selected_format = export_format or settings.export_format
    if selected_format == "gs_ply":
        return settings.gs_model
    return settings.model


def get_model(model_key):
    try:
        return MODELS[model_key]
    except KeyError as error:
        raise ValueError(f"Unknown DA3 model: {model_key}") from error


def get_downloadable_model(model_key):
    try:
        return DOWNLOADABLE_MODELS[model_key]
    except KeyError as error:
        raise ValueError(f"Unknown model: {model_key}") from error


def model_directory_name(model_key):
    """Return the model's folder name inside the shared models directory."""
    return get_downloadable_model(model_key).directory_name