Remove obsolete implementation and simplify project structure
This commit is contained in:
188
src/anyimage/server/model_catalog.py
Normal file
188
src/anyimage/server/model_catalog.py
Normal file
@@ -0,0 +1,188 @@
|
||||
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, ...] = ()
|
||||
ready_patterns: tuple[str, ...] = ("config.json", "*.safetensors")
|
||||
enabled: bool = True
|
||||
r2_directory: str = ""
|
||||
r2_files: tuple[tuple[str, int, 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 using ONNX Runtime",
|
||||
supports_gs=False,
|
||||
download_patterns=("model.onnx",),
|
||||
ready_patterns=("model.onnx",),
|
||||
r2_directory="DA3-SMALL",
|
||||
r2_files=(
|
||||
("model.onnx", 105484029, "7a9d81475b3e2334ace0d6d2c2a816b3c93081437e1866cc434ce1789125f631"),
|
||||
),
|
||||
),
|
||||
"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 using ONNX Runtime",
|
||||
supports_gs=False,
|
||||
download_patterns=("model.onnx",),
|
||||
ready_patterns=("model.onnx",),
|
||||
r2_directory="DA3-BASE",
|
||||
r2_files=(
|
||||
("model.onnx", 412856697, "31f38fc4990539d7902ba267924b9731db1acb00d840c248a60b1865810ff9e6"),
|
||||
),
|
||||
),
|
||||
"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,
|
||||
enabled=False,
|
||||
),
|
||||
"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,
|
||||
enabled=False,
|
||||
),
|
||||
}
|
||||
|
||||
BEN2_MODEL = ModelSpec(
|
||||
key="BEN2_BASE",
|
||||
label="BEN2 Base",
|
||||
directory_name="BEN2-ONNX",
|
||||
repository="onnx-community/BEN2-ONNX",
|
||||
license_name="Apache 2.0",
|
||||
description="Fast high-quality FP16 ONNX background removal model",
|
||||
supports_gs=False,
|
||||
download_patterns=("onnx/model_fp16.onnx",),
|
||||
ready_patterns=("onnx/model_fp16.onnx",),
|
||||
r2_directory="BEN2-ONNX",
|
||||
r2_files=((
|
||||
"onnx/model_fp16.onnx",
|
||||
219121675,
|
||||
"dfdc25f421f32a0d1268e0f2ff2153d340e8f1d52d3dd16f5dc33c1ce85cedf1",
|
||||
),),
|
||||
)
|
||||
|
||||
|
||||
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",
|
||||
description="Fastest MoGe-2 model using ONNX Runtime",
|
||||
supports_gs=False,
|
||||
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",
|
||||
description="Balanced MoGe-2 model using ONNX Runtime",
|
||||
supports_gs=False,
|
||||
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",
|
||||
description="Highest-quality MoGe-2 model using ONNX Runtime",
|
||||
supports_gs=False,
|
||||
download_patterns=("model.onnx",),
|
||||
ready_patterns=("model.onnx",),
|
||||
r2_directory="moge-2-vitl-normal-onnx",
|
||||
r2_files=((
|
||||
"model.onnx",
|
||||
1324265014,
|
||||
"afbc4ccc3450298f3afb35b90f015f4c4f552dea21dc6470d5f7b78b77e2d751",
|
||||
),),
|
||||
),
|
||||
}
|
||||
MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITS_NORMAL"]
|
||||
DOWNLOADABLE_MODELS = {
|
||||
**MODELS,
|
||||
BEN2_MODEL.key: BEN2_MODEL,
|
||||
**MOGE2_MODELS,
|
||||
}
|
||||
MODEL_NUMBERS = {
|
||||
model_key: number for number, model_key in enumerate(MODELS)
|
||||
}
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
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,
|
||||
"supports_gs": model.supports_gs,
|
||||
"enabled": model.enabled,
|
||||
"family": "moge2" if model.key.startswith("MOGE2_") else (
|
||||
"ben2" if model.key == BEN2_MODEL.key else "da3"
|
||||
),
|
||||
"ready": bool(ready),
|
||||
}
|
||||
Reference in New Issue
Block a user