icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add model catalog with specifications for DA3 and MoGe-2 models

Browse Source 
- Introduced ModelSpec dataclass to define model attributes.
- Added multiple models including DA3 Small, Base, Giant 1.1, Nested Giant Large 1.1, and MoGe-2 variants.
- Implemented functions to retrieve model details and metadata.
- Included BEN2 model specifications for background removal.
This commit is contained in:
马南
2026-08-17 14:53:25 +08:00
parent c5b77c5c85
commit 1c98b21235
44 changed files with 637 additions and 2709 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
src/snapform/server/__init__.py
View File

@@ -1,19 +1,6 @@
"""SnapForm Job Server without Blender dependencies."""
import sys
from pathlib import Path
SERVER_DIRECTORY = Path(__file__).resolve().parent
ADDON_DIRECTORY = SERVER_DIRECTORY.parent
for import_directory in (SERVER_DIRECTORY, ADDON_DIRECTORY):
if str(import_directory) not in sys.path:
sys.path.insert(0, str(import_directory))
try:
from ..blendjob import JobServer
except ImportError:
from blendjob import JobServer
from blendjob import JobServer
from .model_manager import ModelManager, ModelSessionCache
from .background import remove_background as _remove_background
@@ -36,6 +23,12 @@ def download_model(context, parameters):
return manager.download(context, parameters["model"])
@server.job("download-required-models")
def download_required_models(context, parameters):
manager = context.resource("model_manager")
return manager.download_missing(context, parameters["models"])
@server.job("debug-ben2")
@server.job("remove-background")
def remove_background(context, parameters):

1
src/snapform/server/da3.py
View File

@@ -121,7 +121,6 @@ def generate_da3(context, parameters):
inference_options = {
"process_res": parameters["process_res"],
"process_res_method": parameters["process_res_method"],
"use_ray_pose": bool(parameters["use_ray_pose"]),
"ref_view_strategy": parameters["ref_view_strategy"],
}
try:

189
src/snapform/server/model_catalog.py
View File

@@ -1,189 +0,0 @@
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),
}

7
src/snapform/server/model_download.py
View File

@@ -15,12 +15,13 @@ from urllib.request import (
HF_ENDPOINT = "https://huggingface.co"
R2_ENDPOINT = "https://models.omoolab.xyz"
CONNECT_TIMEOUT = 3.0
CONNECT_TIMEOUT = 10.0
READ_TIMEOUT = 60.0
DOWNLOAD_CHUNK_SIZE = 256 * 1024
class _ShortConnectHTTPSConnection(http.client.HTTPSConnection):
"""Use a short timeout for connection setup, without limiting downloads."""
"""Bound connection and idle reads without limiting total download time."""
def connect(self):
try:
@@ -28,7 +29,7 @@ class _ShortConnectHTTPSConnection(http.client.HTTPSConnection):
super().connect()
finally:
self.timeout = None
self.sock.settimeout(None)
self.sock.settimeout(READ_TIMEOUT)
class _ShortConnectHTTPSHandler(HTTPSHandler):

21
src/snapform/server/model_manager.py
View File

@@ -3,7 +3,7 @@ import time
from pathlib import Path
from types import SimpleNamespace
from . import model_catalog
import model_catalog
from .model_download import download_model
@@ -157,11 +157,28 @@ class ModelManager:
self.sessions.close()
def download(self, context, model_key):
return self._download(context, model_key, context.progress)
def download_missing(self, context, model_keys):
missing = [key for key in model_keys if not self.ready(key)]
if not missing:
context.progress(1.0, "Required models are ready")
return {"downloaded": []}
total = len(missing)
for index, model_key in enumerate(missing):
def report(value, message, *, offset=index):
context.progress((offset + value) / total, message)
self._download(context, model_key, report)
return {"downloaded": missing}
def _download(self, context, model_key, progress):
model = model_catalog.get_downloadable_model(model_key)
args = SimpleNamespace(
repository=model.repository,
model_dir=str(self.directory(model_key)),
progress=context.progress,
progress=progress,
cancel_check=context.check_cancelled,
label=model.label,
allow_pattern=list(model.download_patterns),

3
src/snapform/server/onnx_da3.py
View File

@@ -195,9 +195,7 @@ def infer(
process_res=504,
process_res_method="upper_bound_resize",
ref_view_strategy="saddle_balanced",
use_ray_pose=False,
):
del use_ray_pose
inputs = session.get_inputs()
model_input, processed_images = preprocess(
image_paths,
@@ -240,4 +238,3 @@ def infer(
processed_images=processed_images,
aux={},
)

15
src/snapform/server/onnx_moge2.py
View File

@@ -1,19 +1,9 @@
"""MoGe-2 ONNX inference and NumPy post-processing."""
from pathlib import Path
import sys
SCRIPT_DIRECTORY = Path(__file__).resolve().parent
if str(SCRIPT_DIRECTORY) not in sys.path:
sys.path.insert(0, str(SCRIPT_DIRECTORY))
try:
from .onnx_runtime import create_session as create_runtime_session
from .onnx_runtime import provider_label, select_providers
except ImportError:
from onnx_runtime import create_session as create_runtime_session
from onnx_runtime import provider_label, select_providers
from .onnx_runtime import create_session as create_runtime_session
from .onnx_runtime import provider_label, select_providers
ONNX_FILENAME = "model.onnx"
@@ -163,4 +153,3 @@ def infer(session, image, resolution_level):
raw_output = dict(zip(names, values))
raw_output["metric_scale"] = raw_output.pop(scale_name)
return postprocess(raw_output)