nan
1c98b21235
- 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.
117 lines
4.5 KiB
Python
117 lines
4.5 KiB
Python
import importlib.util
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def load_module(name, relative_path):
|
|
specification = importlib.util.spec_from_file_location(
|
|
name,
|
|
PROJECT_ROOT / relative_path,
|
|
)
|
|
module = importlib.util.module_from_spec(specification)
|
|
specification.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class ModelTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.models = load_module(
|
|
"snapform_models",
|
|
"src/snapform/model_catalog.py",
|
|
)
|
|
|
|
def test_base_model_is_default_friendly_apache_model(self):
|
|
model = self.models.get_model("BASE")
|
|
|
|
self.assertEqual(model.repository, "depth-anything/DA3-BASE")
|
|
self.assertEqual(model.license_name, "Apache 2.0")
|
|
|
|
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):
|
|
self.assertTrue(self.models.get_model(model_key).supports_gs)
|
|
self.assertFalse(self.models.get_model(model_key).enabled)
|
|
|
|
def test_only_uploaded_da3_models_are_enabled(self):
|
|
enabled = [key for key, model in self.models.MODELS.items() if model.enabled]
|
|
|
|
self.assertEqual(enabled, ["SMALL", "BASE"])
|
|
|
|
def test_unknown_model_is_rejected(self):
|
|
with self.assertRaisesRegex(ValueError, "Unknown DA3 model"):
|
|
self.models.get_model("UNKNOWN")
|
|
|
|
def test_ben2_download_only_includes_required_hugging_face_files(self):
|
|
model = self.models.get_downloadable_model("BEN2_BASE")
|
|
|
|
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_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,
|
|
)
|
|
self.assertEqual(
|
|
self.models.model_directory_name(model_key),
|
|
model.directory_name,
|
|
)
|
|
|
|
def test_ben2_model_directory_uses_repository_name(self):
|
|
self.assertEqual(
|
|
self.models.model_directory_name("BEN2_BASE"),
|
|
"BEN2-ONNX",
|
|
)
|
|
|
|
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")
|
|
|
|
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(
|
|
list(self.models.MOGE2_MODELS),
|
|
["MOGE2_VITS_NORMAL", "MOGE2_VITB_NORMAL", "MOGE2_VITL_NORMAL"],
|
|
)
|
|
self.assertEqual(
|
|
[model.r2_directory for model in self.models.MOGE2_MODELS.values()],
|
|
[
|
|
"moge-2-vits-normal-onnx",
|
|
"moge-2-vitb-normal-onnx",
|
|
"moge-2-vitl-normal-onnx",
|
|
],
|
|
)
|