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:
@@ -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"},
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user