Add ONNX model inference implementations for MoGe-2, BEN2, and DA3

- Introduced `onnx_moge2.py` for MoGe-2 ONNX inference with NumPy post-processing.
- Added `onnx_ben2.py` for BEN2 ONNX inference with Pillow and NumPy preprocessing.
- Created `onnx_da3.py` for Depth Anything 3 inference with ONNX Runtime and NumPy.
- Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`.
- Developed unit tests for server functionalities in `test_server.py`, including model caching and job management.
- Enhanced job server to handle frame extraction without requiring models.
This commit is contained in:
2026-08-17 10:37:55 +08:00
parent b3a07cde32
commit c5b77c5c85
58 changed files with 5840 additions and 9931 deletions

View File

@@ -1,7 +1,142 @@
import platform
import sys
import bpy
from .models import gs_model_enum_items, model_enum_items, moge2_model_enum_items
from .client import device_enum_items
from .job_runtime import runtime
BEN2_MODEL_KEY = "BEN2_BASE"
DEFAULT_DA3_MODEL_KEY = "BASE"
DEFAULT_MOGE2_MODEL_KEY = "MOGE2_VITS_NORMAL"
DEVICE_LABELS = {
"cuda": ("CUDA", "NVIDIA CUDA"),
"directml": ("DirectML", "Windows GPU (DirectML)"),
"mps": ("MPS", "Apple Metal (MPS)"),
"cpu": ("CPU", "CPU"),
}
_models_cache = {}
def supported_devices():
devices = runtime.environment_manifest().get("devices")
if isinstance(devices, list) and devices:
return [str(device).lower() for device in devices]
if sys.platform == "win32":
return ["directml", "cpu"]
if sys.platform == "darwin" and platform.machine().lower() in {
"arm64",
"aarch64",
}:
return ["mps", "cpu"]
return ["cpu"]
def device_enum_items(_owner, _context):
return [
(name.upper(), *DEVICE_LABELS[name], index)
for index, name in enumerate(supported_devices())
]
def models_status(refresh=False):
global _models_cache
status = runtime.server_status().get("resources", {}).get("model_manager", {})
if status:
_models_cache = status
if status and not refresh:
return status
if _models_cache and not refresh:
return _models_cache
if not runtime.environment_ready():
return {}
try:
status = runtime.resource("model_manager")
_models_cache = status
return status
except RuntimeError:
return _models_cache
def model_ready(model_key):
return any(
model.get("key") == model_key and model.get("ready")
for model in models_status().get("catalog", ())
)
def model_catalog():
return models_status().get("catalog", ())
def model_record(model_key):
return next(
(model for model in model_catalog() if model.get("key") == model_key),
None,
)
def _model_items(models):
return [
(
model["key"],
model["label"] if model.get("ready") else f"{model['label']} (Missing)",
f"{model['description']}; license: {model['license']}",
index,
)
for index, model in enumerate(models)
]
def model_enum_items(_owner, _context):
return _model_items(
model
for model in model_catalog()
if model.get("family") == "da3" and model.get("enabled")
)
def gs_model_enum_items(_owner, _context):
return _model_items(
model
for model in model_catalog()
if model.get("family") == "da3"
and model.get("enabled")
and model.get("supports_gs")
)
def moge2_model_enum_items(_owner, _context):
return _model_items(
model for model in model_catalog() if model.get("family") == "moge2"
)
def depth_plane_model_enum_items(_owner, _context):
models = [
model
for model in model_catalog()
if model.get("family") in {"moge2", "da3"} and model.get("enabled")
]
models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
items = _model_items(models)
items.append(
("ORIGINAL", "Original", "Use the original image as depth", len(items))
)
return items
def surface_model_enum_items(_owner, _context):
models = [
model
for model in model_catalog()
if model.get("family") in {"moge2", "da3"}
and model.get("enabled")
]
models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
return _model_items(models)
class SnapFormSettings(bpy.types.PropertyGroup):
@@ -185,16 +320,3 @@ class SnapFormSettings(bpy.types.PropertyGroup):
default=1,
min=1,
)
status: bpy.props.StringProperty(
name="Status",
default="Ready",
options={"SKIP_SAVE"},
)
progress: bpy.props.FloatProperty(
name="Progress",
default=0.0,
min=0.0,
max=1.0,
subtype="FACTOR",
options={"SKIP_SAVE"},
)