Files
AnyImage/tests/test_models.py
nan a6370d678d feat: Add MoGe-2 support and enhance depth plane functionality
- Introduced new MoGe-2 depth plane group with geometry and material displacement options.
- Updated existing depth plane group to improve background depth handling and subdivision settings.
- Enhanced tests to validate MoGe-2 models and ensure proper integration with existing systems.
- Modified installation environment to include new dependencies for MoGe-2 and utils3d.
- Updated vendor lock file to include new artifacts for MoGe and utils3d.
- Refactored download and model handling to accommodate new MoGe-2 models.
- Improved validation scripts for depth plane assets to ensure compliance with new specifications.
2026-08-13 15:55:58 +08:00

101 lines
3.6 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/models.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_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):
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, "PramaLLC/BEN2")
self.assertEqual(model.license_name, "MIT")
self.assertEqual(
model.download_patterns,
("config.json", "model.safetensors"),
)
self.assertNotIn("BEN2_BASE", self.models.MODELS)
def test_da3_model_directories_share_the_da3_prefix(self):
for model_key, model in self.models.MODELS.items():
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_ben2_prefix(self):
self.assertEqual(
self.models.model_directory_name("BEN2_BASE"),
"ben2_base",
)
def test_moge2_uses_single_checkpoint_and_mit_license(self):
model = self.models.get_downloadable_model("MOGE2_VITB_NORMAL")
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_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()],
["MoGe2ViTSNormal", "MoGe2ViTBNormal", "MoGe2ViTLNormal"],
)