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()

23
tests/test_download_progress.py
View File

@@ -2,6 +2,7 @@ import os
import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
@@ -55,6 +56,28 @@ class DownloadProgressTest(unittest.TestCase):
)
return progress, reports
def test_https_connection_bounds_connect_and_idle_read_time(self):
connect_timeouts = []
read_timeouts = []
def connect(connection):
connect_timeouts.append(connection.timeout)
connection.sock = SimpleNamespace(
settimeout=read_timeouts.append,
)
connection = self.download._ShortConnectHTTPSConnection("example.com")
with patch.object(
self.download.http.client.HTTPSConnection,
"connect",
new=connect,
):
connection.connect()
self.assertEqual(connect_timeouts, [10.0])
self.assertEqual(read_timeouts, [60.0])
self.assertIsNone(connection.timeout)
def test_reports_aggregated_download_bytes(self):
progress, reports = self.create_progress()

121
tests/test_install_env.py
View File

@@ -1,121 +0,0 @@
import importlib.util
import json
import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
def load_environment():
path = PROJECT_ROOT / "src" / "snapform" / "blendjob" / "environment.py"
specification = importlib.util.spec_from_file_location(
"blendjob_environment_test",
path,
)
module = importlib.util.module_from_spec(specification)
specification.loader.exec_module(module)
return module
class EnvironmentTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.environment = load_environment()
def config(self):
return {
"python": "3.12",
"packages": ["numpy==2.4.2"],
"platform_packages": {
"windows": ["onnxruntime-directml==1.24.4"],
"default": ["onnxruntime==1.24.4"],
},
}
def test_normalizes_public_dictionary(self):
self.assertEqual(
self.environment.normalized_environment(self.config()),
self.config(),
)
def test_requires_python_and_list_packages(self):
with self.assertRaises(ValueError):
self.environment.normalized_environment({})
with self.assertRaises(TypeError):
self.environment.normalized_environment(
{"python": "3.12", "packages": "numpy"}
)
def test_environment_digest_tracks_configuration(self):
first = self.environment.environment_digest(self.config())
reordered = {
"platform_packages": self.config()["platform_packages"],
"packages": ["numpy==2.4.2"],
"python": "3.12",
}
self.assertEqual(first, self.environment.environment_digest(reordered))
reordered["packages"].append("pillow==12.1.1")
self.assertNotEqual(first, self.environment.environment_digest(reordered))
def test_selects_platform_packages_and_server_packages(self):
with patch.object(self.environment.sys, "platform", "win32"):
packages = self.environment.resolved_packages(self.config())
self.assertIn("fastapi==0.139.2", packages)
self.assertIn("uvicorn==0.51.0", packages)
self.assertIn("numpy==2.4.2", packages)
self.assertIn("onnxruntime-directml==1.24.4", packages)
self.assertNotIn("onnxruntime==1.24.4", packages)
def test_uses_default_platform_packages(self):
with patch.object(self.environment.sys, "platform", "linux"):
packages = self.environment.resolved_packages(self.config())
self.assertIn("onnxruntime==1.24.4", packages)
def test_install_writes_manifest_from_environment_hash(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
uv = root / "tools" / "uv"
uv.parent.mkdir()
uv.touch()
status = root / "environment-status.json"
with patch.object(
self.environment,
"install_uv",
return_value=uv,
), patch.object(
self.environment,
"run_command",
), patch.object(
self.environment,
"install_packages",
) as install_packages:
self.environment.install(root, self.config(), status, stages=3)
manifest = json.loads((root / "manifest.json").read_text("utf-8"))
final_status = json.loads(status.read_text("utf-8"))
self.assertEqual(
manifest["environment_hash"],
self.environment.environment_digest(self.config()),
)
self.assertIn("numpy==2.4.2", manifest["packages"])
self.assertEqual(final_status["stage"], 1)
self.assertEqual(final_status["stages"], 3)
self.assertEqual(final_status["progress"], 1.0)
self.assertTrue(install_packages.called)
def test_environment_python_is_inside_runtime_storage(self):
path = self.environment.environment_python("storage")
expected = (
Path("storage/.venv/Scripts/python.exe")
if sys.platform == "win32"
else Path("storage/.venv/bin/python")
)
self.assertEqual(path, expected)
if __name__ == "__main__":
unittest.main()

2
tests/test_models.py
View File

@@ -21,7 +21,7 @@ class ModelTest(unittest.TestCase):
def setUpClass(cls):
cls.models = load_module(
"snapform_models",
"src/snapform/server/model_catalog.py",
"src/snapform/model_catalog.py",
)
def test_base_model_is_default_friendly_apache_model(self):

10
tests/test_moge_onnx.py
View File

@@ -2,6 +2,7 @@ import importlib.util
import sys
import unittest
from pathlib import Path
from types import ModuleType
from unittest.mock import patch
import numpy as np
@@ -11,12 +12,17 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1]
def load_module():
snapform = ModuleType("snapform")
snapform.__path__ = [str(PROJECT_ROOT / "src" / "snapform")]
server = ModuleType("snapform.server")
server.__path__ = [str(PROJECT_ROOT / "src" / "snapform" / "server")]
specification = importlib.util.spec_from_file_location(
"snapform_onnx_moge2",
"snapform.server.onnx_moge2",
PROJECT_ROOT / "src" / "snapform" / "server" / "onnx_moge2.py",
)
module = importlib.util.module_from_spec(specification)
specification.loader.exec_module(module)
with patch.dict(sys.modules, {"snapform": snapform, "snapform.server": server}):
specification.loader.exec_module(module)
return module

11
tests/test_pack.py
View File

@@ -65,15 +65,20 @@ class PackTest(unittest.TestCase):
self.assertIn("server/da3.py", names)
self.assertNotIn("server/processing.py", names)
self.assertNotIn("server/progress.py", names)
self.assertIn("server/model_catalog.py", names)
self.assertIn("model_catalog.py", names)
self.assertIn("server/model_download.py", names)
self.assertIn("server/onnx_ben2.py", names)
self.assertIn("server/onnx_da3.py", names)
self.assertIn("server/onnx_moge2.py", names)
self.assertNotIn("server/jobs.py", names)
self.assertIn("server/__init__.py", names)
self.assertIn("blendjob/environment.py", names)
self.assertIn("blendjob/installer/environment.py", names)
self.assertIn(
"wheels/blendjob-0.1.5-py3-none-any.whl",
names,
)
self.assertFalse(
any(name.startswith("blendjob/") for name in names)
)
self.assertNotIn("install_env.py", names)
self.assertNotIn("snapform/__init__.py", names)
self.assertNotIn("importer.py", names)

30
tests/test_server.py
View File

@@ -44,6 +44,35 @@ class ServerTest(unittest.TestCase):
self.addCleanup(self.jobs_temp.cleanup)
self.storage_root = Path(self.jobs_temp.name)
def test_required_model_download_skips_ready_model(self):
manager = self.server_module.ModelManager(Path("."))
downloads = []
context = SimpleNamespace(progress=lambda *_args: None)
with patch.object(
manager,
"ready",
side_effect=lambda key: key == "BEN2_BASE",
), patch.object(
manager,
"_download",
side_effect=lambda _context, key, _progress: downloads.append(key),
):
result = manager.download_missing(
context,
("MOGE2_VITS_NORMAL", "BEN2_BASE"),
)
self.assertEqual(downloads, ["MOGE2_VITS_NORMAL"])
self.assertEqual(result, {"downloaded": ["MOGE2_VITS_NORMAL"]})
def test_models_are_owned_by_the_server_resource(self):
self.assertFalse((PROJECT_ROOT / "src" / "snapform" / "models.py").exists())
with tempfile.TemporaryDirectory() as directory:
manager = self.server_module.ModelManager(Path(directory))
snapshot = manager.snapshot()
self.assertIn("catalog", snapshot)
self.assertTrue(snapshot["catalog"])
def create_cache(self):
cache = self.server_module.ModelSessionCache.__new__(
self.server_module.ModelSessionCache
@@ -178,6 +207,7 @@ class ServerTest(unittest.TestCase):
set(server.handlers),
{
"download-model",
"download-required-models",
"remove-background",
"extract-frames",
"generate-da3-depth",

3
tests/test_worker.py
View File

@@ -799,7 +799,6 @@ class WorkerTest(unittest.TestCase):
"device": "cpu",
"process_res": 518,
"process_res_method": "lower_bound_resize",
"use_ray_pose": 1,
"ref_view_strategy": "middle",
"conf_thresh_percentile": 35.0,
"num_max_points": 12345,
@@ -823,7 +822,6 @@ class WorkerTest(unittest.TestCase):
options["process_res_method"],
"lower_bound_resize",
)
self.assertTrue(options["use_ray_pose"])
self.assertEqual(options["ref_view_strategy"], "middle")
self.assertEqual(mesh_calls[0][1], [input_path])
self.assertEqual(mesh_calls[0][2:], (35.0, True, 0.2))
@@ -874,7 +872,6 @@ class WorkerTest(unittest.TestCase):
"device": "cpu",
"process_res": 518,
"process_res_method": "upper_bound_resize",
"use_ray_pose": 1,
"ref_view_strategy": "first",
"conf_thresh_percentile": 40.0,
"num_max_points": 1000,