Add ONNX model inference implementations for MoGe-2, BEN2, and DA3
- 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.
This commit is contained in:
341
tests/test_server.py
Normal file
341
tests/test_server.py
Normal file
@@ -0,0 +1,341 @@
|
||||
import importlib
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def load_server_module():
|
||||
addon_directory = PROJECT_ROOT / "src" / "snapform"
|
||||
sys.path.insert(0, str(addon_directory))
|
||||
try:
|
||||
return importlib.import_module("server")
|
||||
finally:
|
||||
sys.path.remove(str(addon_directory))
|
||||
|
||||
|
||||
class FakeCuda:
|
||||
class OutOfMemoryError(RuntimeError):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def is_available():
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def empty_cache():
|
||||
return None
|
||||
|
||||
|
||||
class ServerTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.server_module = load_server_module()
|
||||
|
||||
def setUp(self):
|
||||
self.jobs_temp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.jobs_temp.cleanup)
|
||||
self.storage_root = Path(self.jobs_temp.name)
|
||||
|
||||
def create_cache(self):
|
||||
cache = self.server_module.ModelSessionCache.__new__(
|
||||
self.server_module.ModelSessionCache
|
||||
)
|
||||
cache.ben2_model = None
|
||||
cache.ben2_key = None
|
||||
cache.da3_model = None
|
||||
cache.da3_key = None
|
||||
cache.moge_model = None
|
||||
cache.moge_key = None
|
||||
return cache
|
||||
|
||||
def create_server(self):
|
||||
source = self.server_module.server
|
||||
server = self.server_module.JobServer(
|
||||
"SnapForm Job Server", storage_root=self.storage_root
|
||||
)
|
||||
server.handlers.update(source.handlers)
|
||||
server.resource_factories.update(source.resource_factories)
|
||||
server.bind(self.storage_root)
|
||||
return server
|
||||
|
||||
def wait_for_job(self, context, timeout=2.0):
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
if context.snapshot()["state"] in {"succeeded", "failed", "cancelled"}:
|
||||
return context.snapshot()
|
||||
time.sleep(0.01)
|
||||
self.fail(f"Job did not finish: {context.snapshot()}")
|
||||
|
||||
def test_ben2_model_is_reused_for_matching_runtime_key(self):
|
||||
loads = []
|
||||
|
||||
ben2 = ModuleType("server.onnx_ben2")
|
||||
ben2.create_session = lambda model_dir, device: (
|
||||
loads.append((model_dir, device)) or object()
|
||||
)
|
||||
cache = self.create_cache()
|
||||
device = SimpleNamespace(type="cuda")
|
||||
with patch.dict(sys.modules, {"server.onnx_ben2": ben2}):
|
||||
first, _load_ms = cache.get_ben2("model", device)
|
||||
second, cached_load_ms = cache.get_ben2("model", device)
|
||||
|
||||
self.assertIs(first, second)
|
||||
self.assertEqual(loads, [("model", "cuda")])
|
||||
self.assertEqual(cache.snapshot()["ben2"], "model")
|
||||
self.assertEqual(cached_load_ms, 0.0)
|
||||
|
||||
def test_da3_model_is_reused_for_matching_runtime_key(self):
|
||||
loads = []
|
||||
|
||||
da3 = ModuleType("server.onnx_da3")
|
||||
da3.create_session = lambda model_dir, device: (
|
||||
loads.append((model_dir, device)) or object()
|
||||
)
|
||||
cache = self.create_cache()
|
||||
device = SimpleNamespace(type="cuda")
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"server.onnx_da3": da3,
|
||||
},
|
||||
):
|
||||
first = cache.get_da3("model", device)
|
||||
second = cache.get_da3("model", device)
|
||||
|
||||
self.assertIs(first, second)
|
||||
self.assertEqual(loads, [("model", "cuda")])
|
||||
self.assertEqual(cache.snapshot()["da3"], "model")
|
||||
|
||||
def test_depth_model_cache_keeps_only_da3_or_moge(self):
|
||||
cache = self.create_cache()
|
||||
cache.ben2_model = object()
|
||||
cache.ben2_key = ("ben2", "cuda")
|
||||
cache.moge_model = object()
|
||||
cache.moge_key = ("moge", "cuda")
|
||||
device = SimpleNamespace(type="cuda")
|
||||
da3 = object()
|
||||
|
||||
def load_da3(_model_dir, _device, key):
|
||||
cache.da3_model = da3
|
||||
cache.da3_key = key
|
||||
return da3
|
||||
|
||||
with patch.object(cache, "_load_da3", side_effect=load_da3):
|
||||
self.assertIs(cache.get_da3("da3", device), da3)
|
||||
|
||||
self.assertIsNone(cache.moge_model)
|
||||
self.assertIsNotNone(cache.ben2_model)
|
||||
|
||||
moge = object()
|
||||
|
||||
fake_moge = ModuleType("server.onnx_moge2")
|
||||
fake_moge.create_session = lambda _model_dir, _device: moge
|
||||
with patch.dict(sys.modules, {"server.onnx_moge2": fake_moge}):
|
||||
self.assertIs(cache.get_moge2("moge", "cuda"), moge)
|
||||
|
||||
self.assertIsNone(cache.da3_model)
|
||||
self.assertIsNotNone(cache.ben2_model)
|
||||
|
||||
def test_dedicated_server_runs_only_one_job_at_a_time(self):
|
||||
server = self.server_module.JobServer("Test", storage_root=self.storage_root)
|
||||
first_started = threading.Event()
|
||||
release_first = threading.Event()
|
||||
order = []
|
||||
|
||||
@server.job("example")
|
||||
def example(_context, parameters):
|
||||
order.append(parameters["name"])
|
||||
if parameters["name"] == "first":
|
||||
first_started.set()
|
||||
release_first.wait(1.0)
|
||||
|
||||
first = server.submit("example", {"name": "first"}, job_id="job")
|
||||
self.assertTrue(first_started.wait(1.0))
|
||||
second = server.submit("example", {"name": "second"}, job_id="other")
|
||||
|
||||
self.assertIsNotNone(first)
|
||||
self.assertIsNotNone(second)
|
||||
self.assertTrue(server.snapshot("instance")["busy"])
|
||||
self.assertEqual(server.snapshot("instance")["queued_jobs"], 1)
|
||||
self.assertNotIn("clients", server.snapshot("instance"))
|
||||
release_first.set()
|
||||
self.wait_for_job(first)
|
||||
self.wait_for_job(second)
|
||||
self.assertEqual(order, ["first", "second"])
|
||||
server.close()
|
||||
|
||||
def test_server_registers_tasks_instead_of_model_names(self):
|
||||
server = self.create_server()
|
||||
self.assertEqual(
|
||||
set(server.handlers),
|
||||
{
|
||||
"download-model",
|
||||
"remove-background",
|
||||
"extract-frames",
|
||||
"generate-da3-depth",
|
||||
"generate-moge2-depth",
|
||||
"generate-da3-surface",
|
||||
"generate-moge2-surface",
|
||||
"debug-da3",
|
||||
"debug-moge2",
|
||||
"debug-ben2",
|
||||
},
|
||||
)
|
||||
|
||||
def test_queued_job_can_be_cancelled_before_it_runs(self):
|
||||
server = self.server_module.JobServer("Test", storage_root=self.storage_root)
|
||||
first_started = threading.Event()
|
||||
release_first = threading.Event()
|
||||
executed = []
|
||||
|
||||
@server.job("example")
|
||||
def example(_context, parameters):
|
||||
executed.append(parameters["name"])
|
||||
if parameters["name"] == "first":
|
||||
first_started.set()
|
||||
release_first.wait(1.0)
|
||||
|
||||
first = server.submit("example", {"name": "first"})
|
||||
self.assertTrue(first_started.wait(1.0))
|
||||
second = server.submit("example", {"name": "second"})
|
||||
|
||||
self.assertIs(server.cancel(second.job_id), second)
|
||||
release_first.set()
|
||||
self.wait_for_job(first)
|
||||
second_status = self.wait_for_job(second)
|
||||
|
||||
self.assertEqual(second_status["state"], "cancelled")
|
||||
self.assertEqual(executed, ["first"])
|
||||
server.close()
|
||||
|
||||
def test_cancelling_job_keeps_server_and_model_cache_alive(self):
|
||||
fake_models = SimpleNamespace(
|
||||
snapshot=lambda: {
|
||||
"loaded": {"ben2": "BEN2_BASE", "da3": ""}
|
||||
},
|
||||
clear=lambda: None,
|
||||
close=lambda: None,
|
||||
)
|
||||
with patch.object(
|
||||
self.server_module, "ModelManager", return_value=fake_models
|
||||
):
|
||||
server = self.create_server()
|
||||
context = server.submit("remove-background", {}, job_id="job")
|
||||
self.assertIsNotNone(server.cancel("job"))
|
||||
self.assertTrue(context.is_cancelled())
|
||||
status = self.wait_for_job(context)
|
||||
|
||||
self.assertEqual(status["state"], "cancelled")
|
||||
self.assertFalse(server.shutdown_event.is_set())
|
||||
self.assertEqual(
|
||||
server.snapshot("instance")["resources"]["model_manager"]["loaded"]["ben2"],
|
||||
"BEN2_BASE",
|
||||
)
|
||||
|
||||
def test_job_progress_is_kept_in_server_memory(self):
|
||||
server = self.server_module.JobServer(
|
||||
"Test", storage_root=self.storage_root
|
||||
)
|
||||
|
||||
@server.job("example")
|
||||
def example(context, _parameters):
|
||||
context.progress(0.5, "Half way")
|
||||
|
||||
context = server.submit("example", {})
|
||||
self.wait_for_job(context)
|
||||
|
||||
self.assertRegex(context.job_id, r"^[0-9a-f]{32}$")
|
||||
self.assertEqual(context.directory.name, context.job_id)
|
||||
self.assertEqual(context.snapshot()["state"], "succeeded")
|
||||
self.assertEqual(context.snapshot()["progress"], 1.0)
|
||||
|
||||
def test_runtime_storage_root_owns_the_jobs_directory(self):
|
||||
storage_root = self.storage_root
|
||||
server = self.server_module.JobServer("Test")
|
||||
server.bind(storage_root)
|
||||
|
||||
@server.job("example")
|
||||
def example(_context, _parameters):
|
||||
return {"value": 42}
|
||||
|
||||
context = server.submit("example", {})
|
||||
self.wait_for_job(context)
|
||||
|
||||
self.assertEqual(server.jobs_directory, storage_root / "jobs")
|
||||
self.assertEqual(context.storage_root, storage_root.resolve())
|
||||
self.assertEqual(context.snapshot()["result"], {"value": 42})
|
||||
|
||||
def test_job_server_owns_and_clears_shared_resources(self):
|
||||
events = []
|
||||
resource = SimpleNamespace(
|
||||
snapshot=lambda: {"loaded": "model"},
|
||||
clear=lambda: events.append("clear"),
|
||||
close=lambda: events.append("close"),
|
||||
)
|
||||
server = self.server_module.JobServer(
|
||||
"Test", storage_root=self.storage_root
|
||||
)
|
||||
server.add_resource("model_manager", resource)
|
||||
|
||||
@server.job("example")
|
||||
def example(context, _parameters):
|
||||
self.assertIs(context.resource("model_manager"), resource)
|
||||
release.wait(1.0)
|
||||
|
||||
release = threading.Event()
|
||||
context = server.submit("example", {}, job_id="job")
|
||||
self.assertFalse(server.clear_resource("model_manager"))
|
||||
release.set()
|
||||
self.wait_for_job(context)
|
||||
self.assertTrue(server.clear_resource("model_manager"))
|
||||
server.close()
|
||||
|
||||
self.assertEqual(
|
||||
server.snapshot("instance")["resources"]["model_manager"]["loaded"],
|
||||
"model",
|
||||
)
|
||||
self.assertEqual(events, ["clear", "close"])
|
||||
|
||||
def test_server_dispatches_frame_extraction_without_models(self):
|
||||
fake_models = SimpleNamespace(
|
||||
snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}},
|
||||
clear=lambda: None,
|
||||
close=lambda: None,
|
||||
)
|
||||
with patch.object(
|
||||
self.server_module, "ModelManager", return_value=fake_models
|
||||
):
|
||||
server = self.create_server()
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
input_directory = root / "frames"
|
||||
input_directory.mkdir()
|
||||
(input_directory / "frame-0001.jpg").touch()
|
||||
(input_directory / "frame-0002.jpg").touch()
|
||||
context = server.submit(
|
||||
"extract-frames",
|
||||
{
|
||||
"input": str(input_directory),
|
||||
"video_frames": 16,
|
||||
"start_frame": 1,
|
||||
},
|
||||
job_id="job",
|
||||
)
|
||||
status = self.wait_for_job(context)
|
||||
self.assertTrue(
|
||||
(context.directory / "frame.0001.jpg").is_file()
|
||||
)
|
||||
|
||||
self.assertEqual(status["state"], "succeeded")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user