nan
c5b77c5c85
- Introduced `onnx_moge2.py` for MoGe-2 ONNX inference with NumPy post-processing. - Added `onnx_ben2.py` for BEN2 ONNX inference with Pillow and NumPy preprocessing. - Created `onnx_da3.py` for Depth Anything 3 inference with ONNX Runtime and NumPy. - Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`. - Developed unit tests for server functionalities in `test_server.py`, including model caching and job management. - Enhanced job server to handle frame extraction without requiring models.
524 lines
19 KiB
Python
524 lines
19 KiB
Python
import importlib
|
|
import json
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
SOURCE_ROOT = PROJECT_ROOT / "src"
|
|
|
|
|
|
class AddonTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.remove_modules()
|
|
|
|
def append_status_bar(draw):
|
|
draw._owner = "snapform"
|
|
cls.status_bar_draws.append(draw)
|
|
|
|
fake_bpy = ModuleType("bpy")
|
|
fake_bpy.props = SimpleNamespace(
|
|
BoolProperty=lambda **options: options,
|
|
EnumProperty=lambda **options: options,
|
|
FloatProperty=lambda **options: options,
|
|
IntProperty=lambda **options: options,
|
|
PointerProperty=lambda **options: options,
|
|
StringProperty=lambda **options: options,
|
|
)
|
|
fake_bpy.types = SimpleNamespace(
|
|
AddonPreferences=object,
|
|
Menu=object,
|
|
Operator=object,
|
|
Panel=object,
|
|
PropertyGroup=object,
|
|
Scene=type("Scene", (), {}),
|
|
STATUSBAR_HT_header=SimpleNamespace(
|
|
append=append_status_bar,
|
|
remove=lambda draw: cls.status_bar_draws.remove(draw),
|
|
),
|
|
VIEW3D_MT_object_context_menu=SimpleNamespace(
|
|
append=lambda draw: cls.context_menu_draws.append(draw),
|
|
remove=lambda draw: cls.context_menu_draws.remove(draw),
|
|
),
|
|
OUTLINER_MT_object=SimpleNamespace(
|
|
append=lambda draw: cls.outliner_menu_draws.append(draw),
|
|
remove=lambda draw: cls.outliner_menu_draws.remove(draw),
|
|
),
|
|
)
|
|
cls.registered = []
|
|
cls.status_bar_draws = []
|
|
cls.context_menu_draws = []
|
|
cls.outliner_menu_draws = []
|
|
fake_bpy.utils = SimpleNamespace(
|
|
register_class=cls.registered.append,
|
|
unregister_class=cls.registered.remove,
|
|
)
|
|
fake_bpy.context = SimpleNamespace(
|
|
preferences=SimpleNamespace(addons={}),
|
|
)
|
|
fake_bpy.app = SimpleNamespace(online_access=True)
|
|
fake_bpy.path = SimpleNamespace(abspath=lambda path: path)
|
|
fake_bpy.data = SimpleNamespace(
|
|
objects=SimpleNamespace(get=lambda _name: None),
|
|
)
|
|
sys.modules["bpy"] = fake_bpy
|
|
sys.path.insert(0, str(SOURCE_ROOT))
|
|
cls.addon = importlib.import_module("snapform")
|
|
cls.fake_bpy = fake_bpy
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
sys.path.remove(str(SOURCE_ROOT))
|
|
sys.modules.pop("bpy", None)
|
|
cls.remove_modules()
|
|
|
|
@staticmethod
|
|
def remove_modules():
|
|
for module_name in tuple(sys.modules):
|
|
if module_name == "snapform" or module_name.startswith("snapform."):
|
|
sys.modules.pop(module_name)
|
|
|
|
def setUp(self):
|
|
preferences = SimpleNamespace(
|
|
storage_root="",
|
|
device="CUDA",
|
|
)
|
|
self.fake_bpy.context.preferences.addons = {
|
|
"snapform": SimpleNamespace(preferences=preferences)
|
|
}
|
|
|
|
def tearDown(self):
|
|
self.addon.runtime.close_active()
|
|
self.registered.clear()
|
|
self.status_bar_draws.clear()
|
|
self.context_menu_draws.clear()
|
|
self.outliner_menu_draws.clear()
|
|
self.fake_bpy.data = SimpleNamespace(
|
|
objects=SimpleNamespace(get=lambda _name: None),
|
|
)
|
|
if hasattr(self.fake_bpy.types.Scene, "snapform_settings"):
|
|
del self.fake_bpy.types.Scene.snapform_settings
|
|
|
|
def test_snapform_has_no_runtime_boundary_modules(self):
|
|
for name in (
|
|
"client.py",
|
|
"environment.py",
|
|
"storage.py",
|
|
"models.py",
|
|
"status_bar.py",
|
|
):
|
|
self.assertFalse((SOURCE_ROOT / "snapform" / name).exists(), name)
|
|
self.assertFalse(hasattr(self.addon, "client"))
|
|
|
|
def test_job_request_is_not_part_of_blendjob(self):
|
|
blendjob = importlib.import_module("snapform.blendjob")
|
|
self.assertFalse(hasattr(blendjob, "JobRequest"))
|
|
|
|
def test_job_operators_are_flat_and_use_plain_properties(self):
|
|
operators = self.addon.operators
|
|
for name in (
|
|
"SnapFormJobMixin",
|
|
"ImageInputJobOperator",
|
|
"GeometryImageJobOperator",
|
|
"Da3ImageJobOperator",
|
|
"Moge2ImageJobOperator",
|
|
"BackgroundRequestMixin",
|
|
):
|
|
self.assertFalse(hasattr(operators, name), name)
|
|
for operator in operators.JOB_OPERATOR_CLASSES:
|
|
for name in getattr(operator, "__annotations__", {}):
|
|
self.assertFalse(name.startswith("snapform_"), (operator, name))
|
|
|
|
def test_concrete_job_operator_is_the_blender_operator(self):
|
|
operator = self.addon.operators.GenerateMoge2Depth
|
|
self.assertEqual(operator.bl_idname, "snapform.generate_moge2_depth")
|
|
self.assertIn(self.addon.job_runtime.JobOperatorBase, operator.__mro__)
|
|
self.assertIn(self.fake_bpy.types.Operator, operator.__mro__)
|
|
|
|
def test_default_request_sends_every_operator_property(self):
|
|
class Example(self.addon.job_runtime.JobOperatorBase):
|
|
value: object()
|
|
enabled: object()
|
|
|
|
operator = Example()
|
|
operator.value = 42
|
|
operator.enabled = True
|
|
self.assertEqual(
|
|
operator.request(object()),
|
|
{"value": 42, "enabled": True},
|
|
)
|
|
|
|
def test_overridden_request_purely_replaces_default_properties(self):
|
|
class Example(self.addon.job_runtime.JobOperatorBase):
|
|
value: object()
|
|
|
|
def request(self, context):
|
|
return {"replacement": context.value}
|
|
|
|
operator = Example()
|
|
operator.value = 42
|
|
self.assertEqual(
|
|
operator.request(SimpleNamespace(value=7)),
|
|
{"replacement": 7},
|
|
)
|
|
|
|
def test_runtime_submits_job_type_and_plain_dict(self):
|
|
calls = []
|
|
submitted = threading.Event()
|
|
|
|
class Controller:
|
|
def submit(self, job_type, parameters):
|
|
calls.append((job_type, parameters))
|
|
submitted.set()
|
|
return {"job_id": "job", "directory": "."}
|
|
|
|
def cancel(self, _job_id):
|
|
pass
|
|
|
|
def mark_job_complete(self, _job_id):
|
|
pass
|
|
|
|
class Example(self.addon.job_runtime.JobOperatorBase):
|
|
job_type = "example"
|
|
value: object()
|
|
|
|
def controller(self, _runtime):
|
|
return Controller()
|
|
|
|
timers = []
|
|
window_manager = SimpleNamespace(
|
|
progress_begin=lambda *_args: None,
|
|
progress_update=lambda _value: None,
|
|
event_timer_add=lambda *_args, **_kwargs: timers.append(object()) or timers[-1],
|
|
event_timer_remove=lambda _timer: None,
|
|
modal_handler_add=lambda _operator: None,
|
|
)
|
|
context = SimpleNamespace(
|
|
window_manager=window_manager,
|
|
window=object(),
|
|
)
|
|
operator = Example()
|
|
operator.value = 9
|
|
|
|
self.assertEqual(operator.execute(context), {"RUNNING_MODAL"})
|
|
self.assertTrue(submitted.wait(1.0))
|
|
self.assertEqual(calls, [("example", {"value": 9})])
|
|
|
|
def test_runtime_owns_environment_and_storage_paths(self):
|
|
runtime = self.addon.runtime
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
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.server_log_path(), root / "server.log")
|
|
|
|
def test_runtime_checks_its_environment_dictionary_hash(self):
|
|
runtime = self.addon.runtime
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
python = root / ".venv" / (
|
|
"Scripts/python.exe" if sys.platform == "win32" else "bin/python"
|
|
)
|
|
python.parent.mkdir(parents=True)
|
|
python.touch()
|
|
(root / "manifest.json").write_text(
|
|
json.dumps({"environment_hash": runtime.environment_hash}),
|
|
encoding="utf-8",
|
|
)
|
|
with patch.object(runtime, "storage_root_factory", root):
|
|
self.assertTrue(runtime.environment_ready())
|
|
|
|
def test_runtime_builds_the_standard_installer_command(self):
|
|
runtime = self.addon.runtime
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
with patch.object(runtime, "storage_root_factory", root):
|
|
command = runtime.install_command()
|
|
self.assertEqual(command[0], sys.executable)
|
|
self.assertEqual(Path(command[1]).name, "environment.py")
|
|
self.assertEqual(Path(command[1]).parent.name, "installer")
|
|
self.assertIn("--storage-root", command)
|
|
self.assertIn("--config", command)
|
|
self.assertNotIn("--venv-dir", command)
|
|
self.assertNotIn("--models-dir", command)
|
|
self.assertEqual(command[command.index("--stages") + 1], "1")
|
|
|
|
def test_runtime_exposes_declared_environment_and_post_install(self):
|
|
runtime = self.addon.runtime
|
|
self.assertEqual(runtime.environment["python"], "3.12")
|
|
self.assertIn("numpy==2.4.2", runtime.environment["packages"])
|
|
self.assertTrue(callable(runtime.post_install))
|
|
|
|
def test_snapform_post_install_is_ordinary_code_using_runtime_request(self):
|
|
calls = []
|
|
fake_runtime = SimpleNamespace(
|
|
request=lambda job_type, parameters: calls.append(
|
|
(job_type, parameters)
|
|
)
|
|
)
|
|
self.addon.job_runtime.post_install(fake_runtime)
|
|
self.assertEqual(
|
|
calls,
|
|
[
|
|
("download-model", {"model": "MOGE2_VITS_NORMAL"}),
|
|
("download-model", {"model": "BEN2_BASE"}),
|
|
],
|
|
)
|
|
|
|
def test_environment_controller_runs_arbitrary_post_install_callable(self):
|
|
runtime_module = importlib.import_module("snapform.blendjob.runtime")
|
|
runtime = self.addon.runtime
|
|
calls = []
|
|
started = threading.Event()
|
|
release = threading.Event()
|
|
|
|
def post_install(received):
|
|
calls.append(received)
|
|
started.set()
|
|
release.wait(1.0)
|
|
|
|
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,
|
|
"environment_status_path",
|
|
return_value=Path(directory) / "missing.json",
|
|
):
|
|
self.assertEqual(controller.status("environment")["state"], "running")
|
|
self.assertTrue(started.wait(1.0))
|
|
release.set()
|
|
controller.post_thread.join(1.0)
|
|
status = controller.status("environment")
|
|
|
|
self.assertEqual(calls, [runtime])
|
|
self.assertEqual(status["state"], "succeeded")
|
|
|
|
def test_post_install_progress_resets_without_stage_fraction(self):
|
|
runtime_module = importlib.import_module("snapform.blendjob.runtime")
|
|
controller = runtime_module.EnvironmentController(self.addon.runtime)
|
|
|
|
controller._update_post_progress(
|
|
{"job_id": "first", "job_type": "download-model", "progress": 1.0}
|
|
)
|
|
first_stage = controller.post_status["stage"]
|
|
controller._update_post_progress(
|
|
{"job_id": "second", "job_type": "download-model", "progress": 0.1}
|
|
)
|
|
|
|
self.assertGreater(controller.post_status["stage"], first_stage)
|
|
self.assertNotIn("stages", controller.post_status)
|
|
|
|
def test_conversion_model_lists_default_to_moge2_vits(self):
|
|
properties = importlib.import_module("snapform.properties")
|
|
catalog = (
|
|
{"key": "SMALL", "family": "da3", "enabled": True,
|
|
"label": "DA3 Small", "description": "DA3", "license": "Apache"},
|
|
{"key": "MOGE2_VITS_NORMAL", "family": "moge2", "enabled": True,
|
|
"label": "MoGe-2 ViT-S", "description": "MoGe", "license": "MIT"},
|
|
)
|
|
with patch.object(properties, "model_catalog", return_value=catalog):
|
|
depth_items = properties.depth_plane_model_enum_items(None, None)
|
|
surface_items = properties.surface_model_enum_items(None, None)
|
|
|
|
self.assertEqual(depth_items[0][0], "MOGE2_VITS_NORMAL")
|
|
self.assertEqual(surface_items[0][0], "MOGE2_VITS_NORMAL")
|
|
|
|
def test_runtime_request_uses_request_response_shape(self):
|
|
result = object()
|
|
responses = []
|
|
with patch.object(
|
|
self.addon.runtime.server,
|
|
"request",
|
|
return_value=result,
|
|
) as request:
|
|
returned = self.addon.runtime.request(
|
|
"example",
|
|
{"value": 3},
|
|
response=responses.append,
|
|
)
|
|
self.assertIs(returned, result)
|
|
request.assert_called_once_with(
|
|
"example",
|
|
{"value": 3},
|
|
response=responses.append,
|
|
progress=None,
|
|
cancel_check=None,
|
|
)
|
|
|
|
def test_server_request_returns_job_result_and_calls_response(self):
|
|
blendjob = importlib.import_module("snapform.blendjob")
|
|
controller = blendjob.ServerController.__new__(blendjob.ServerController)
|
|
statuses = iter(
|
|
(
|
|
{"state": "running", "progress": 0.5},
|
|
{
|
|
"state": "succeeded",
|
|
"progress": 1.0,
|
|
"result": {"value": 7},
|
|
},
|
|
)
|
|
)
|
|
controller.submit = lambda _job_type, _parameters: {
|
|
"job_id": "job",
|
|
"directory": ".",
|
|
}
|
|
controller.status = lambda _job_id: next(statuses)
|
|
controller.mark_job_complete = lambda _job_id: None
|
|
responses = []
|
|
|
|
result = controller.request(
|
|
"example",
|
|
{"input": 3},
|
|
response=responses.append,
|
|
poll_interval=0.01,
|
|
)
|
|
|
|
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_registers_system_operators_and_status_bar(self):
|
|
self.addon.register()
|
|
system_classes = self.addon.runtime.operator_classes()
|
|
|
|
self.assertTrue(set(system_classes) <= set(self.registered))
|
|
self.assertEqual(len(self.status_bar_draws), 1)
|
|
self.assertEqual(self.status_bar_draws[0]._owner, "snapform")
|
|
|
|
self.addon.unregister()
|
|
self.assertEqual(self.registered, [])
|
|
self.assertEqual(self.status_bar_draws, [])
|
|
|
|
def test_status_bar_reads_runtime_state(self):
|
|
calls = []
|
|
|
|
class Layout:
|
|
def row(self, **_options):
|
|
return self
|
|
|
|
def progress(self, **options):
|
|
calls.append(("progress", options))
|
|
|
|
def operator(self, identifier, **_options):
|
|
calls.append(("operator", identifier))
|
|
|
|
runtime = self.addon.runtime
|
|
runtime.active_job = object()
|
|
runtime.progress = 0.4
|
|
runtime.message = "Working"
|
|
runtime._draw_status_bar(SimpleNamespace(layout=Layout()), None)
|
|
runtime.active_job = None
|
|
|
|
self.assertEqual(calls[0][1]["factor"], 0.4)
|
|
self.assertEqual(calls[0][1]["text"], "Working")
|
|
self.assertEqual(calls[1], ("operator", "snapform.cancel_job"))
|
|
|
|
def test_moge_depth_request_is_written_on_the_operator(self):
|
|
operator = self.addon.operators.GenerateMoge2Depth()
|
|
operator.input_path = __file__
|
|
operator.source_object_name = ""
|
|
operator.model = "MOGE2_VITS_NORMAL"
|
|
operator.resolution_level = 6
|
|
operator.remove_background = False
|
|
settings = SimpleNamespace(video_frames=12)
|
|
context = SimpleNamespace(
|
|
scene=SimpleNamespace(snapform_settings=settings)
|
|
)
|
|
|
|
with patch.object(
|
|
self.addon.operators, "require_environment"
|
|
), patch.object(self.addon.operators, "require_model"):
|
|
parameters = operator.request(context)
|
|
|
|
self.assertEqual(parameters["input"], __file__)
|
|
self.assertEqual(parameters["model"], "MOGE2_VITS_NORMAL")
|
|
self.assertEqual(parameters["resolution_level"], 6)
|
|
self.assertNotIn("job_type", parameters)
|
|
|
|
def test_moge_depth_response_calls_an_ordinary_function(self):
|
|
operator = self.addon.operators.GenerateMoge2Depth()
|
|
operator.source_object_name = "Reference"
|
|
operator.remove_background = False
|
|
operator.delete_source = True
|
|
operator.unshaded = False
|
|
result = object()
|
|
context = object()
|
|
|
|
with patch.object(
|
|
self.addon.operators, "build_moge2_depth_response"
|
|
) as build:
|
|
operator.response(context, result)
|
|
|
|
build.assert_called_once_with(
|
|
context,
|
|
result,
|
|
source_object_name="Reference",
|
|
remove_background=False,
|
|
delete_source=True,
|
|
unshaded=False,
|
|
)
|
|
|
|
def test_entry_operator_dispatches_plain_properties(self):
|
|
operator = self.addon.ConvertToDepthPlane()
|
|
operator.model = "MOGE2_VITS_NORMAL"
|
|
operator.resolution_level = 7
|
|
operator.keep_original = False
|
|
operator.unshaded = True
|
|
operator.remove_background = False
|
|
source_object = SimpleNamespace(
|
|
name="Reference",
|
|
data=SimpleNamespace(source="FILE"),
|
|
)
|
|
settings = SimpleNamespace(video_frames=1)
|
|
context = SimpleNamespace(
|
|
object=source_object,
|
|
scene=SimpleNamespace(snapform_settings=settings),
|
|
)
|
|
calls = []
|
|
self.fake_bpy.ops = SimpleNamespace(
|
|
snapform=SimpleNamespace(
|
|
generate_moge2_depth=lambda **options: calls.append(options)
|
|
or {"RUNNING_MODAL"}
|
|
)
|
|
)
|
|
|
|
with patch.object(
|
|
self.addon.image_depth,
|
|
"prepare_image_input",
|
|
return_value=(Path(__file__), False),
|
|
):
|
|
self.assertEqual(operator.execute(context), {"FINISHED"})
|
|
|
|
self.assertEqual(calls[0]["resolution_level"], 7)
|
|
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()
|