Add tools and tests for AnyImage model management and documentation
- Implement sync_models.py for preparing and uploading ONNX artifacts. - Create build_node.py for building and validating Blender scripts. - Add docs.py for building and previewing documentation with Mike or MkDocs. - Introduce model_exports.py for exporting PyTorch models to ONNX format. - Develop pack.py for packaging the project with dependencies for different platforms. - Add test_sync_models.py to validate sync_models functionality. - Add test_tool_commands.py to ensure repository tools have correct command scripts. - Initialize tools package with __init__.py.
This commit is contained in:
138
tests/test_model_catalog.py
Normal file
138
tests/test_model_catalog.py
Normal file
@@ -0,0 +1,138 @@
|
||||
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 ModelCatalogTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.models = load_module(
|
||||
"anyimage_server_models",
|
||||
"src/anyimage/server/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_upscale_models_use_verified_onnx_artifacts(self):
|
||||
expected = {
|
||||
"REALESRGAN_GENERAL_WDN_X4V3": ("realesr-general-wdn-x4v3.onnx", 4867419, "b021a837"),
|
||||
"SAFMN_REAL_X4": ("safmn-real-x4.onnx", 22657477, "8f13867b"),
|
||||
"REALESRGAN_X4PLUS": ("realesrgan-x4plus.onnx", 67051618, "4ed6a45a"),
|
||||
}
|
||||
self.assertEqual(set(self.models.UPSCALE_MODELS), set(expected))
|
||||
for key, (filename, size, checksum_prefix) in expected.items():
|
||||
with self.subTest(model=key):
|
||||
model = self.models.get_downloadable_model(key)
|
||||
self.assertEqual(model.download_patterns, (filename,))
|
||||
self.assertEqual(model.ready_patterns, (filename,))
|
||||
self.assertEqual(model.r2_files[0][:2], (filename, size))
|
||||
self.assertTrue(model.r2_files[0][2].startswith(checksum_prefix))
|
||||
self.assertEqual(self.models.model_record(model)["family"], "upscale")
|
||||
wdn = self.models.UPSCALE_MODELS["REALESRGAN_GENERAL_WDN_X4V3"]
|
||||
self.assertEqual(wdn.repository, "xinntao/Real-ESRGAN")
|
||||
self.assertEqual(wdn.huggingface_repository, "")
|
||||
self.assertEqual(
|
||||
self.models.UPSCALE_MODELS["SAFMN_REAL_X4"].huggingface_repository, ""
|
||||
)
|
||||
|
||||
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",
|
||||
],
|
||||
)
|
||||
Reference in New Issue
Block a user