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

feat: Implement ONNX model inference for BEN2, Depth Anything 3, and MoGe

Browse Source 
- Added `ben2_onnx.py` for BEN2 ONNX inference with preprocessing.
- Introduced `da3_onnx.py` for Depth Anything 3 inference, including camera projection calculations.
- Created `moge_onnx.py` for MoGe-2 ONNX inference and post-processing.
- Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`.
- Added model synchronization script `sync_models.py` for downloading and verifying ONNX models.
- Developed unit tests for BEN2, Depth Anything 3, and MoGe ONNX models to ensure functionality.
- Updated `vendor.lock.toml` to remove unused artifacts.
This commit is contained in:
马南
2026-08-14 10:27:09 +08:00
parent 13fab9ce88
commit 2ad30bf4a3
41 changed files with 1569 additions and 2109 deletions
Download Patch File Download Diff File Expand all files Collapse all files

260
tests/test_install_env.py
View File

@@ -36,7 +36,7 @@ class InstallEnvTest(unittest.TestCase):
(vendor / "moge.tar.gz").write_bytes(b"moge")
return vendor
def test_installs_da3_from_archives_without_gsplat(self):
def test_installs_onnx_runtime_without_training_frameworks(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
@@ -71,77 +71,31 @@ class InstallEnvTest(unittest.TestCase):
commands += [
call.args[4] for call in mocked_run_install.call_args_list
]
da3_command = next(
command
for command in commands
if any("depth-anything-3 @" in str(value) for value in command)
)
self.assertIn(
f"addict=={bootstrap.ADDICT_VERSION}",
da3_command,
)
self.assertTrue(
any(
"depth-anything-3 @ file:" in str(value)
and "snapform_vendor_" in str(value)
for value in da3_command
)
)
self.assertFalse(
any("git+" in str(value) for value in da3_command)
)
self.assertFalse(
any(
"gsplat" in str(value)
for command in commands
for value in command
)
)
self.assertFalse(
any(
"--no-build-isolation" in str(value)
for command in commands
for value in command
)
)
ben2_command = next(
command
for command in commands
if any("ben2 @" in str(value) for value in command)
)
self.assertTrue(
any(
"ben2 @ file:" in str(value)
and "snapform_vendor_" in str(value)
for value in ben2_command
)
)
self.assertFalse(
any("git+" in str(value) for value in ben2_command)
)
service_command = next(
command
for command in commands
if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command
)
self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", service_command)
self.assertIn("setuptools", service_command)
self.assertNotIn("ninja", service_command)
self.assertIn(
f"onnxruntime-directml=={bootstrap.ONNXRUNTIME_VERSION}",
service_command,
)
for forbidden in ("torch", "torchvision", "xformers", "ben2 @", "moge @"):
self.assertFalse(
any(forbidden in str(value) for value in service_command),
forbidden,
)
check_command = bootstrap.verify_command(
bootstrap.venv_python(venv_dir)
)
self.assertTrue(
any("import torch;" in str(value) for value in check_command)
)
self.assertTrue(
any(
"from ben2 import BEN_Base" in str(value)
for value in check_command
)
)
self.assertFalse(any("torch" in str(value) for value in check_command))
self.assertTrue(
any("import fastapi, uvicorn" in str(value) for value in check_command)
)
self.assertTrue(
any("onnxruntime" in str(value) for value in check_command)
)
self.assertFalse(
any("import gsplat" in str(value) for value in check_command)
)
@@ -157,24 +111,24 @@ class InstallEnvTest(unittest.TestCase):
if any(
repository in command
for repository in (
"Ruicheng/moge-2-vitb-normal",
"PramaLLC/BEN2",
"Ruicheng/moge-2-vits-normal-onnx",
"onnx-community/BEN2-ONNX",
)
)
]
self.assertEqual(len(default_downloads), 2)
self.assertTrue(any("Ruicheng/moge-2-vitb-normal" in map(str, command) for command in default_downloads))
self.assertTrue(any("Ruicheng/moge-2-vits-normal-onnx" in map(str, command) for command in default_downloads))
ben2_download = next(
command
for command in default_downloads
if "PramaLLC/BEN2" in command
if "onnx-community/BEN2-ONNX" in command
)
patterns = [
ben2_download[index + 1]
for index, value in enumerate(ben2_download)
if value == "--allow-pattern"
]
self.assertEqual(patterns, ["config.json", "model.safetensors"])
self.assertEqual(patterns, ["onnx/model_fp16.onnx"])
self.assertNotIn("--stage", ben2_download)
self.assertNotIn("--stage-label", ben2_download)
self.assertEqual(
@@ -182,7 +136,7 @@ class InstallEnvTest(unittest.TestCase):
"BEN2 Base",
)
def test_darwin_install_skips_cuda_index_and_xformers(self):
def test_darwin_install_uses_default_onnx_runtime(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
@@ -224,15 +178,15 @@ class InstallEnvTest(unittest.TestCase):
commands = [
call.args[4] for call in mocked_run_install.call_args_list
]
torch_command = next(
runtime_command = next(
command
for command in commands
if any(
f"torch=={bootstrap.TORCH_VERSION}" in str(value)
f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}" in str(value)
for value in command
)
)
self.assertNotIn("--index-url", torch_command)
self.assertNotIn("--index-url", runtime_command)
self.assertFalse(
any(
"xformers" in str(value)
@@ -241,30 +195,27 @@ class InstallEnvTest(unittest.TestCase):
)
)
def test_intel_mac_is_rejected_with_clear_error(self):
def test_intel_mac_uses_portable_onnx_runtime(self):
bootstrap = load_install_env()
args = SimpleNamespace(
uv="uv",
venv_dir=".venv",
models_dir="models",
torch_variant="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status="status.json",
)
with patch.object(
bootstrap.sys,
"platform",
"darwin",
), patch.object(
bootstrap.platform,
"machine",
return_value="x86_64",
), patch.object(bootstrap.shutil, "rmtree") as mocked_rmtree:
with self.assertRaisesRegex(RuntimeError, "Apple Silicon"):
bootstrap.install(args)
mocked_rmtree.assert_not_called()
):
stages = bootstrap.install_stages(
SimpleNamespace(
models_dir="models",
status="status.json",
),
".venv",
"python",
"uv",
)
runtime = stages[1][1]
self.assertIn(
f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}",
runtime,
)
def test_run_stage_writes_stage_progress(self):
bootstrap = load_install_env()
@@ -414,10 +365,10 @@ class InstallEnvTest(unittest.TestCase):
)
self.assertEqual(ends[-1], 1.0)
self.assertEqual(starts[0], bootstrap.ENVIRONMENT_SLICE[0])
runtime_ranges = stage_ranges[:4]
runtime_ranges = stage_ranges[:2]
self.assertEqual(
[item[0] for item in runtime_ranges],
[1, 2, 3, 3],
[1, 2],
)
self.assertTrue(
all(item[0] is None for item in stage_ranges[4:])
@@ -606,19 +557,19 @@ class InstallEnvTest(unittest.TestCase):
self.assertEqual(bootstrap.format_elapsed(12), "12s")
self.assertEqual(bootstrap.format_elapsed(83), "1m 23s")
def test_ai_component_estimate_continues_across_internal_steps(self):
def test_onnx_runtime_estimate_decreases_with_progress(self):
bootstrap = load_install_env()
first = bootstrap.estimated_remaining(
"AI components",
bootstrap.XFORMERS_SLICE[0],
)
second = bootstrap.estimated_remaining(
"AI components",
"ONNX Runtime",
bootstrap.AI_COMPONENTS_SLICE[0],
)
second = bootstrap.estimated_remaining(
"ONNX Runtime",
sum(bootstrap.AI_COMPONENTS_SLICE) / 2,
)
self.assertEqual(first, bootstrap.INSTALL_ESTIMATES["AI components"])
self.assertEqual(first, bootstrap.INSTALL_ESTIMATES["ONNX Runtime"])
self.assertGreater(second, 0)
self.assertLess(second, first)
@@ -727,7 +678,7 @@ class InstallEnvTest(unittest.TestCase):
ready["packages"],
["numpy==1.26.4", "torch==2.9.1+cu128"],
)
self.assertEqual(ready["devices"], ["cuda", "cpu"])
self.assertEqual(ready["devices"], ["directml", "cpu"])
self.assertEqual(
ready["manifest_revision"],
bootstrap.MANIFEST_REVISION,
@@ -814,7 +765,7 @@ class InstallEnvTest(unittest.TestCase):
def fake_run(command, environment=None):
commands.append(command)
if any("moge-2-vitb-normal" in str(value) for value in command):
if any("moge-2-vits-normal" in str(value) for value in command):
raise subprocess.CalledProcessError(1, command)
with tempfile.TemporaryDirectory() as directory:
@@ -852,7 +803,7 @@ class InstallEnvTest(unittest.TestCase):
self.assertIn("Environment is installed", status["message"])
self.assertTrue(
any(
"PramaLLC/BEN2" in map(str, command)
"onnx-community/BEN2-ONNX" in map(str, command)
for command in commands
)
)
@@ -983,25 +934,13 @@ class InstallEnvTest(unittest.TestCase):
for value in command
)
)
ai_command = next(
runtime_command = next(
command
for command in install_commands
if any("depth-anything-3 @" in str(value) for value in command)
)
self.assertTrue(
any(
"depth-anything-3 @ file:" in str(value)
and "snapform_vendor_" in str(value)
for value in ai_command
)
)
self.assertTrue(
any(
"ben2 @ file:" in str(value)
and "snapform_vendor_" in str(value)
for value in ai_command
)
if any("onnxruntime" in str(value) for value in command)
)
self.assertFalse(any(" @ file:" in str(value) for value in runtime_command))
self.assertFalse(any("torch" in str(value) for value in runtime_command))
def test_install_reuses_existing_uv(self):
bootstrap = load_install_env()
@@ -1047,81 +986,12 @@ class InstallEnvTest(unittest.TestCase):
self.assertEqual(uv_path.read_bytes(), b"binary")
def test_copy_vendor_inputs_copies_missing_archives(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
bundled = root / "bundled"
bundled.mkdir()
(bundled / "depth-anything-3.tar.gz").write_bytes(b"da3")
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
(bundled / "utils3d.tar.gz").write_bytes(b"utils3d")
(bundled / "moge.tar.gz").write_bytes(b"moge")
vendor = root / "vendor"
bootstrap.copy_vendor_inputs(bundled, vendor)
self.assertEqual(
(vendor / "depth-anything-3.tar.gz").read_bytes(),
b"da3",
)
self.assertEqual(
(vendor / "ben2.tar.gz").read_bytes(),
b"ben2",
)
def test_copy_vendor_inputs_replaces_same_size_stale_archive(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
bundled = root / "bundled"
bundled.mkdir()
(bundled / "depth-anything-3.tar.gz").write_bytes(b"new")
(bundled / "ben2.tar.gz").write_bytes(b"ben2")
(bundled / "utils3d.tar.gz").write_bytes(b"utils3d")
(bundled / "moge.tar.gz").write_bytes(b"moge")
vendor = root / "vendor"
vendor.mkdir()
(vendor / "depth-anything-3.tar.gz").write_bytes(b"old")
bootstrap.copy_vendor_inputs(bundled, vendor)
self.assertEqual(
(vendor / "depth-anything-3.tar.gz").read_bytes(),
b"new",
)
def test_remove_legacy_vendor_removes_only_install_inputs(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
vendor = root / "vendor"
vendor.mkdir()
(vendor / "depth-anything-3.tar.gz").write_bytes(b"da3")
(root / "depth-anything-3.tar.gz").write_bytes(b"da3")
(root / "ben2.tar.gz").write_bytes(b"ben2")
models = root / "models"
models.mkdir()
bootstrap.remove_legacy_vendor(root)
self.assertFalse(vendor.exists())
self.assertFalse((root / "depth-anything-3.tar.gz").exists())
self.assertFalse((root / "ben2.tar.gz").exists())
self.assertTrue(models.is_dir())
def test_install_stages_uses_local_source_archives(self):
def test_install_stages_use_only_runtime_packages(self):
bootstrap = load_install_env()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
venv_dir = root / ".venv"
venv_dir.mkdir()
vendor = root / "vendor"
vendor.mkdir()
(vendor / "depth-anything-3.tar.gz").write_bytes(b"da3")
(vendor / "ben2.tar.gz").write_bytes(b"ben2")
(vendor / "utils3d.tar.gz").write_bytes(b"utils3d")
(vendor / "moge.tar.gz").write_bytes(b"moge")
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
@@ -1141,16 +1011,14 @@ class InstallEnvTest(unittest.TestCase):
uv_path,
)
ai_command = next(
runtime_command = next(
command
for _label, command, _environment, kind in stages
if kind == "install"
and any("depth-anything-3" in str(value) for value in command)
and any("onnxruntime" in str(value) for value in command)
)
da3_local = (vendor / "depth-anything-3.tar.gz").as_uri()
ben2_local = (vendor / "ben2.tar.gz").as_uri()
self.assertIn(f"depth-anything-3 @ {da3_local}", ai_command)
self.assertIn(f"ben2 @ {ben2_local}", ai_command)
self.assertFalse(any(" @ file:" in str(value) for value in runtime_command))
self.assertFalse(any("torch" in str(value) for value in runtime_command))
def test_install_stages_download_only_default_models(self):
bootstrap = load_install_env()
@@ -1189,13 +1057,13 @@ class InstallEnvTest(unittest.TestCase):
if kind == "download"
}
self.assertIn(
"moge2_vitb_normal",
" ".join(map(str, downloads["MoGe-2 ViT-B Normal"])),
"moge-2-vits-normal-onnx",
" ".join(map(str, downloads["MoGe-2 ViT-S Normal"])),
)
self.assertIn("ben2_base", " ".join(map(str, downloads["BEN2 Base"])))
self.assertIn("BEN2-ONNX", " ".join(map(str, downloads["BEN2 Base"])))
self.assertEqual(
set(downloads),
{"MoGe-2 ViT-B Normal", "BEN2 Base"},
{"MoGe-2 ViT-S Normal", "BEN2 Base"},
)
def test_install_does_not_verify_env(self):