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

Add model catalog with specifications for DA3 and MoGe-2 models

Browse Source 
- 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.
This commit is contained in:
马南
2026-08-17 14:53:25 +08:00
parent c5b77c5c85
commit 1c98b21235
44 changed files with 637 additions and 2709 deletions
Download Patch File Download Diff File Expand all files Collapse all files

157
tests/test_addon.py
View File

@@ -69,7 +69,9 @@ class AddonTest(unittest.TestCase):
)
sys.modules["bpy"] = fake_bpy
sys.path.insert(0, str(SOURCE_ROOT))
cls.paths_before_import = tuple(sys.path)
cls.addon = importlib.import_module("snapform")
cls.paths_after_import = tuple(sys.path)
cls.fake_bpy = fake_bpy
@classmethod
@@ -116,8 +118,17 @@ class AddonTest(unittest.TestCase):
self.assertFalse((SOURCE_ROOT / "snapform" / name).exists(), name)
self.assertFalse(hasattr(self.addon, "client"))
def test_import_does_not_add_addon_directories_to_sys_path(self):
self.assertEqual(self.paths_after_import, self.paths_before_import)
self.assertFalse(
any(
name == "snapform.server" or name.startswith("snapform.server.")
for name in sys.modules
)
)
def test_job_request_is_not_part_of_blendjob(self):
blendjob = importlib.import_module("snapform.blendjob")
blendjob = importlib.import_module("blendjob")
self.assertFalse(hasattr(blendjob, "JobRequest"))
def test_job_operators_are_flat_and_use_plain_properties(self):
@@ -217,8 +228,8 @@ class AddonTest(unittest.TestCase):
with patch.object(runtime, "storage_root_factory", root):
self.assertEqual(runtime.storage_root(), root.resolve())
self.assertEqual(runtime.environment_directory(), root / ".venv")
self.assertEqual(runtime.environment_status_path(), root / "environment-status.json")
self.assertEqual(runtime.environment_log_path(), root / "environment.log")
self.assertEqual(runtime.install_status_path(), root / "install-status.json")
self.assertEqual(runtime.install_log_path(), root / "install.log")
self.assertEqual(runtime.server_log_path(), root / "server.log")
def test_runtime_checks_its_environment_dictionary_hash(self):
@@ -274,8 +285,25 @@ class AddonTest(unittest.TestCase):
],
)
def test_snapform_post_install_names_each_failed_model(self):
calls = []
def request(_job_type, parameters):
calls.append(parameters["model"])
raise RuntimeError("network unavailable")
with self.assertRaisesRegex(
RuntimeError,
"MoGe-2 ViT-S Normal: network unavailable; BEN2 Base: network unavailable",
):
self.addon.job_runtime.post_install(
SimpleNamespace(request=request)
)
self.assertEqual(calls, ["MOGE2_VITS_NORMAL", "BEN2_BASE"])
def test_environment_controller_runs_arbitrary_post_install_callable(self):
runtime_module = importlib.import_module("snapform.blendjob.runtime")
runtime_module = importlib.import_module("blendjob.runtime")
runtime = self.addon.runtime
calls = []
started = threading.Event()
@@ -294,7 +322,7 @@ class AddonTest(unittest.TestCase):
post_install,
), patch.object(
runtime,
"environment_status_path",
"install_status_path",
return_value=Path(directory) / "missing.json",
):
self.assertEqual(controller.status("environment")["state"], "running")
@@ -306,8 +334,33 @@ class AddonTest(unittest.TestCase):
self.assertEqual(calls, [runtime])
self.assertEqual(status["state"], "succeeded")
def test_environment_controller_reports_post_install_failure(self):
runtime_module = importlib.import_module("blendjob.runtime")
runtime = self.addon.runtime
def post_install(_runtime):
raise RuntimeError("BEN2 download failed")
controller = runtime_module.EnvironmentController(runtime)
controller.process = SimpleNamespace(poll=lambda: 0, returncode=0)
with tempfile.TemporaryDirectory() as directory, patch.object(
runtime,
"post_install",
post_install,
), patch.object(
runtime,
"install_status_path",
return_value=Path(directory) / "missing.json",
):
controller.status("environment")
controller.post_thread.join(1.0)
status = controller.status("environment")
self.assertEqual(status["state"], "failed")
self.assertEqual(status["error"], "BEN2 download failed")
def test_post_install_progress_resets_without_stage_fraction(self):
runtime_module = importlib.import_module("snapform.blendjob.runtime")
runtime_module = importlib.import_module("blendjob.runtime")
controller = runtime_module.EnvironmentController(self.addon.runtime)
controller._update_post_progress(
@@ -336,6 +389,78 @@ class AddonTest(unittest.TestCase):
self.assertEqual(depth_items[0][0], "MOGE2_VITS_NORMAL")
self.assertEqual(surface_items[0][0], "MOGE2_VITS_NORMAL")
def test_required_model_operator_requests_both_models(self):
operator = self.addon.operators.DownloadRequiredModels()
with patch.object(self.addon.operators, "require_environment"):
parameters = operator.request(None)
self.assertEqual(
parameters,
{"models": ["MOGE2_VITS_NORMAL", "BEN2_BASE"]},
)
def test_image_menu_offers_required_model_repair_when_missing(self):
calls = []
class Layout:
operator_context = ""
enabled = True
def __init__(self, enabled=True):
self.enabled = enabled
def row(self):
return Layout(self.enabled)
def operator(self, identifier, **_options):
calls.append((identifier, self.enabled))
return SimpleNamespace()
def separator(self):
calls.append(("separator", self.enabled))
menu = self.addon.SnapFormImageMenu()
menu.layout = Layout()
with patch.object(
self.addon.runtime,
"environment_ready",
return_value=True,
), patch(
"snapform.properties.required_models_ready",
return_value=False,
):
menu.draw(None)
self.assertEqual(
calls[0],
("snapform.download_required_models", True),
)
self.assertIn(("snapform.convert_to_depth_plane", False), calls)
self.assertIn(("snapform.convert_to_surface", True), calls)
def test_ui_reads_required_model_files_without_starting_server(self):
properties = importlib.import_module("snapform.properties")
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
moge = root / "models" / "moge-2-vits-normal-onnx"
ben2 = root / "models" / "BEN2-ONNX" / "onnx"
moge.mkdir(parents=True)
ben2.mkdir(parents=True)
(moge / "model.onnx").touch()
(ben2 / "model_fp16.onnx").touch()
with patch.object(
self.addon.runtime,
"storage_root",
return_value=root,
), patch.object(
self.addon.runtime,
"resource",
side_effect=AssertionError("UI must not start the Server"),
):
self.assertTrue(properties.required_models_ready())
def test_runtime_request_uses_request_response_shape(self):
result = object()
responses = []
@@ -359,7 +484,7 @@ class AddonTest(unittest.TestCase):
)
def test_server_request_returns_job_result_and_calls_response(self):
blendjob = importlib.import_module("snapform.blendjob")
blendjob = importlib.import_module("blendjob")
controller = blendjob.ServerController.__new__(blendjob.ServerController)
statuses = iter(
(
@@ -389,10 +514,11 @@ class AddonTest(unittest.TestCase):
self.assertEqual(result.value, {"value": 7})
self.assertEqual(responses, [result])
def test_runtime_locates_the_package_exporting_server(self):
entrypoint = self.addon.runtime._server_entrypoint()
self.assertEqual(entrypoint.name, "__init__.py")
self.assertEqual(entrypoint.parent.name, "server")
def test_runtime_keeps_server_as_an_entrypoint_string(self):
path, attribute = self.addon.runtime.server_entrypoint.rsplit(":", 1)
self.assertEqual(Path(path).name, "__init__.py")
self.assertEqual(Path(path).parent.name, "server")
self.assertEqual(attribute, "server")
def test_runtime_registers_system_operators_and_status_bar(self):
self.addon.register()
@@ -510,14 +636,5 @@ class AddonTest(unittest.TestCase):
self.assertEqual(calls[0]["source_object_name"], "Reference")
self.assertFalse(any(key.startswith("snapform_") for key in calls[0]))
def test_models_are_owned_by_the_server_resource(self):
self.assertFalse((SOURCE_ROOT / "snapform" / "models.py").exists())
with tempfile.TemporaryDirectory() as directory:
manager = self.addon.server.ModelManager(Path(directory))
snapshot = manager.snapshot()
self.assertIn("catalog", snapshot)
self.assertTrue(snapshot["catalog"])
if __name__ == "__main__":
unittest.main()