AnyImage/src/anyimage/server/model_catalog.py

from dataclasses import dataclass


@dataclass(frozen=True)
class ModelSpec:
    key: str
    label: str
    directory_name: str
    repository: str
    license_name: str
    description: str
    huggingface_repository: str = ""
    download_patterns: tuple[str, ...] = ()
    ready_patterns: tuple[str, ...] = ("config.json", "*.safetensors")
    r2_directory: str = ""
    r2_files: tuple[tuple[str, int, str], ...] = ()


BEN2_MODEL = ModelSpec(
    key="BEN2_BASE",
    label="BEN2 Base",
    directory_name="BEN2-ONNX",
    repository="onnx-community/BEN2-ONNX",
    license_name="Apache 2.0",
    huggingface_repository="onnx-community/BEN2-ONNX",
    description="Fast high-quality FP16 ONNX background removal model",
    download_patterns=("onnx/model_fp16.onnx",),
    ready_patterns=("onnx/model_fp16.onnx",),
    r2_directory="BEN2-ONNX",
    r2_files=((
        "onnx/model_fp16.onnx",
        219121675,
        "dfdc25f421f32a0d1268e0f2ff2153d340e8f1d52d3dd16f5dc33c1ce85cedf1",
    ),),
)


REALESRGAN_GENERAL_WDN_X4V3 = ModelSpec(
    key="REALESRGAN_GENERAL_WDN_X4V3",
    label="Real-ESRGAN General WDN x4v3",
    directory_name="realesr-general-wdn-x4v3",
    repository="xinntao/Real-ESRGAN",
    license_name="BSD 3-Clause",
    description="Fast and conservative 2x output from a 4x ONNX model",
    download_patterns=("realesr-general-wdn-x4v3.onnx",),
    ready_patterns=("realesr-general-wdn-x4v3.onnx",),
    r2_directory="realesr-general-wdn-x4v3",
    r2_files=((
        "realesr-general-wdn-x4v3.onnx",
        4867419,
        "b021a837f7a0d3f287f9a39e367f581f1e6c3a759a08dc0c1c79c22442729774",
    ),),
)

SAFMN_MODEL = ModelSpec(
    key="SAFMN_REAL_X4",
    label="SAFMN Real x4",
    directory_name="safmn-real-x4",
    repository="sunny2109/SAFMN",
    license_name="Apache 2.0",
    description="Balanced 2x output from a real-world 4x ONNX model",
    download_patterns=("safmn-real-x4.onnx",),
    ready_patterns=("safmn-real-x4.onnx",),
    r2_directory="safmn-real-x4",
    r2_files=((
        "safmn-real-x4.onnx",
        22657477,
        "8f13867b66c0300941b902bf6a5c6e0b545794faabe8b0478bf55668bc517621",
    ),),
)

REALESRGAN_X4PLUS = ModelSpec(
    key="REALESRGAN_X4PLUS",
    label="Real-ESRGAN x4plus",
    directory_name="realesrgan-x4plus",
    repository="jonathanst29/tinier-upscale-models",
    license_name="BSD 3-Clause",
    huggingface_repository="jonathanst29/tinier-upscale-models",
    description="Detailed 2x output from a general-purpose 4x ONNX model",
    download_patterns=("realesrgan-x4plus.onnx",),
    ready_patterns=("realesrgan-x4plus.onnx",),
    r2_directory="realesrgan-x4plus",
    r2_files=((
        "realesrgan-x4plus.onnx",
        67051618,
        "4ed6a45a8185bde6c8d7c7790a6e80e3c5b9f608946f634021068dd0fd0473c8",
    ),),
)


PISA_SR = ModelSpec(
    key="PISA_SR",
    label="PiSA-SR",
    directory_name="pisa-sr",
    repository="csslc/PiSA-SR",
    license_name="Apache 2.0 / CreativeML Open RAIL++-M",
    description="Generative 2x upscale with fixed detail strength 0.7",
    download_patterns=("encoder.onnx", "decoder.onnx", "unet-pixel.onnx", "unet-full.onnx", "noise.npy", "LICENSE-APACHE.txt", "LICENSE-MODEL.txt", "NOTICE.txt"),
    ready_patterns=("encoder.onnx", "decoder.onnx", "unet-pixel.onnx", "unet-full.onnx", "noise.npy", "LICENSE-APACHE.txt", "LICENSE-MODEL.txt", "NOTICE.txt"),
    r2_directory="pisa-sr",
    r2_files=(
        ("encoder.onnx", 68426712, "4b93e4a3d81c3349a78ab96619501e39ee8f815d14a20b73fc4f969e30a1ceeb"),
        ("decoder.onnx", 99084589, "c9021384c402e267654907527fa5f0ee86cab1814bfd96233216a15ae8e83c4c"),
        ("unet-pixel.onnx", 1732713007, "4f381574ad584050d0534c8cca794fb5828cb8e99e148f8efc1aaac20833c0b1"),
        ("unet-full.onnx", 1732713007, "3b0fd44ff1093608a65865c3f26d0060775337e1fc117d8ac29aacbe748d3f9c"),
        ("noise.npy", 32896, "a94ed01e2e963022cfa8ec01e04d2e96e599408ef0c385878347eaacd25a663b"),
        ("LICENSE-APACHE.txt", 11358, "cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30"),
        ("LICENSE-MODEL.txt", 14206, "e514164be7f2ca2c0fb39e7283b01ab4a06d802837eaf171aa4645cfd7ee1157"),
        ("NOTICE.txt", 1081, "ccb55fcb98cce42401368b88e3ccd555f4ff49ac1347c6ce063061dd48726050"),
    ),
)


UPSCALE_MODELS = {
    model.key: model
    for model in (
        REALESRGAN_GENERAL_WDN_X4V3,
        SAFMN_MODEL,
        REALESRGAN_X4PLUS,
        PISA_SR,
    )
}


DEFAULT_UPSCALE_MODEL_KEY = REALESRGAN_GENERAL_WDN_X4V3.key


MOGE2_MODELS = {
    "MOGE2_VITS_NORMAL": ModelSpec(
        key="MOGE2_VITS_NORMAL",
        label="MoGe-2 ViT-S Normal",
        directory_name="moge-2-vits-normal-onnx",
        repository="Ruicheng/moge-2-vits-normal-onnx",
        license_name="MIT",
        huggingface_repository="Ruicheng/moge-2-vits-normal-onnx",
        description="Fastest MoGe-2 model using ONNX Runtime",
        download_patterns=("model.onnx",),
        ready_patterns=("model.onnx",),
        r2_directory="moge-2-vits-normal-onnx",
        r2_files=((
            "model.onnx",
            140852051,
            "24eacb5dc7a2c54c7bc98f7de085ffbed79ad006ea5b664c2c2cdc02ff3a52f0",
        ),),
    ),
    "MOGE2_VITB_NORMAL": ModelSpec(
        key="MOGE2_VITB_NORMAL",
        label="MoGe-2 ViT-B Normal",
        directory_name="moge-2-vitb-normal-onnx",
        repository="Ruicheng/moge-2-vitb-normal-onnx",
        license_name="MIT",
        huggingface_repository="Ruicheng/moge-2-vitb-normal-onnx",
        description="Balanced MoGe-2 model using ONNX Runtime",
        download_patterns=("model.onnx",),
        ready_patterns=("model.onnx",),
        r2_directory="moge-2-vitb-normal-onnx",
        r2_files=((
            "model.onnx",
            419411850,
            "bbf14e07a30f11e69d36ab861590123f5598ababcbc8946a063eb4a966f35a21",
        ),),
    ),
    "MOGE2_VITL_NORMAL": ModelSpec(
        key="MOGE2_VITL_NORMAL",
        label="MoGe-2 ViT-L Normal",
        directory_name="moge-2-vitl-normal-onnx",
        repository="Ruicheng/moge-2-vitl-normal-onnx",
        license_name="MIT",
        huggingface_repository="Ruicheng/moge-2-vitl-normal-onnx",
        description="Highest-quality MoGe-2 model using ONNX Runtime",
        download_patterns=("model.onnx",),
        ready_patterns=("model.onnx",),
        r2_directory="moge-2-vitl-normal-onnx",
        r2_files=((
            "model.onnx",
            1324265014,
            "afbc4ccc3450298f3afb35b90f015f4c4f552dea21dc6470d5f7b78b77e2d751",
        ),),
    ),
}
MOGE3_VITL = ModelSpec(
    key="MOGE3_VITL",
    label="MoGe-3 ViT-L",
    directory_name="moge-3-vitl-onnx",
    repository="Ruicheng/moge-3-vitl",
    license_name="MIT",
    description="MoGe-3 geometry with three refinement steps using ONNX Runtime",
    download_patterns=("backbone.onnx", "refiner.onnx"),
    ready_patterns=("backbone.onnx", "refiner.onnx"),
    r2_directory="moge-3-vitl-onnx",
    r2_files=(
        ("backbone.onnx", 1324190318, "f21246c6921373c8c79ecdb8c78818a79bd2506c2fb88f4d4cdec21b92f725e3"),
        ("refiner.onnx", 157729628, "67ad7297aba18f8ef7939e8f700a5b25c41a19ddb68656bc2cd08b54c248fe2a"),
    ),
)

MOGE_MODELS = {**MOGE2_MODELS, MOGE3_VITL.key: MOGE3_VITL}

DOWNLOADABLE_MODELS = {
    BEN2_MODEL.key: BEN2_MODEL,
    **UPSCALE_MODELS,
    **MOGE_MODELS,
}


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_record(model, ready=False):
    """Return public model metadata exposed by the Models Resource."""
    return {
        "key": model.key,
        "label": model.label,
        "description": model.description,
        "license": model.license_name,
        "family": (
            "moge"
            if model.key in MOGE_MODELS
            else "ben2"
            if model.key == BEN2_MODEL.key
            else "upscale"
        ),
        "ready": bool(ready),
    }