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

64
tests/test_models.py
View File

@@ -30,12 +30,6 @@ class ModelTest(unittest.TestCase):
self.assertEqual(model.repository, "depth-anything/DA3-BASE")
self.assertEqual(model.license_name, "Apache 2.0")
def test_large_model_exposes_noncommercial_license(self):
model = self.models.get_model("LARGE_1_1")
self.assertEqual(model.license_name, "CC BY-NC 4.0")
self.assertFalse(model.enabled)
def test_giant_models_expose_gaussian_head(self):
for model_key in ("GIANT_1_1", "NESTED_GIANT_LARGE_1_1"):
with self.subTest(model_key=model_key):
@@ -54,19 +48,20 @@ class ModelTest(unittest.TestCase):
def test_ben2_download_only_includes_required_hugging_face_files(self):
model = self.models.get_downloadable_model("BEN2_BASE")
self.assertEqual(model.repository, "PramaLLC/BEN2")
self.assertEqual(model.license_name, "MIT")
self.assertEqual(
model.download_patterns,
("config.json", "model.safetensors"),
)
self.assertEqual(model.repository, "onnx-community/BEN2-ONNX")
self.assertEqual(model.license_name, "Apache 2.0")
self.assertEqual(model.download_patterns, ("onnx/model_fp16.onnx",))
self.assertEqual(model.ready_patterns, ("onnx/model_fp16.onnx",))
self.assertEqual(model.r2_files[0][0], "onnx/model_fp16.onnx")
self.assertNotIn("BEN2_BASE", self.models.MODELS)
def test_da3_model_directories_share_the_da3_prefix(self):
def test_da3_model_directories_use_repository_names(self):
for model_key, model in self.models.MODELS.items():
if not model.enabled:
continue
with self.subTest(model_key=model_key):
self.assertTrue(
model.directory_name.startswith("da3_"),
model.directory_name.startswith("DA3-"),
model.directory_name,
)
self.assertEqual(
@@ -74,20 +69,37 @@ class ModelTest(unittest.TestCase):
model.directory_name,
)
def test_ben2_model_directory_uses_ben2_prefix(self):
def test_ben2_model_directory_uses_repository_name(self):
self.assertEqual(
self.models.model_directory_name("BEN2_BASE"),
"ben2_base",
"BEN2-ONNX",
)
def test_moge2_uses_single_checkpoint_and_mit_license(self):
model = self.models.get_downloadable_model("MOGE2_VITB_NORMAL")
def test_enabled_da3_models_use_onnx_and_r2_checksums(self):
for model in self.models.MODELS.values():
if not model.enabled:
continue
with self.subTest(model=model.key):
self.assertTrue(model.download_patterns[0].endswith(".onnx"))
self.assertEqual(model.ready_patterns, ("model.onnx",))
self.assertEqual(model.r2_files[0][0], "model.onnx")
self.assertEqual(model.repository, "Ruicheng/moge-2-vitb-normal")
self.assertEqual(model.license_name, "MIT")
self.assertEqual(model.download_patterns, ("model.pt",))
self.assertEqual(model.ready_patterns, ("model.pt",))
self.assertFalse(model.supports_gs)
def test_moge2_models_use_official_onnx_models(self):
for key, variant in (("MOGE2_VITS_NORMAL", "vits"),
("MOGE2_VITB_NORMAL", "vitb"),
("MOGE2_VITL_NORMAL", "vitl")):
with self.subTest(key=key):
model = self.models.get_downloadable_model(key)
self.assertEqual(
model.repository,
f"Ruicheng/moge-2-{variant}-normal-onnx",
)
self.assertEqual(model.license_name, "MIT")
self.assertEqual(model.download_patterns, ("model.onnx",))
self.assertEqual(model.ready_patterns, ("model.onnx",))
self.assertEqual(model.r2_directory, model.repository.split("/")[-1])
self.assertEqual(model.r2_files[0][0], "model.onnx")
self.assertFalse(model.supports_gs)
def test_moge2_models_match_uploaded_r2_variants(self):
self.assertEqual(
@@ -96,5 +108,9 @@ class ModelTest(unittest.TestCase):
)
self.assertEqual(
[model.r2_directory for model in self.models.MOGE2_MODELS.values()],
["MoGe2ViTSNormal", "MoGe2ViTBNormal", "MoGe2ViTLNormal"],
[
"moge-2-vits-normal-onnx",
"moge-2-vitb-normal-onnx",
"moge-2-vitl-normal-onnx",
],
)