nan
2ad30bf4a3
- Added `ben2_onnx.py` for BEN2 ONNX inference with preprocessing. - Introduced `da3_onnx.py` for Depth Anything 3 inference, including camera projection calculations. - Created `moge_onnx.py` for MoGe-2 ONNX inference and post-processing. - Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`. - Added model synchronization script `sync_models.py` for downloading and verifying ONNX models. - Developed unit tests for BEN2, Depth Anything 3, and MoGe ONNX models to ensure functionality. - Updated `vendor.lock.toml` to remove unused artifacts.
263 lines
8.9 KiB
Python
263 lines
8.9 KiB
Python
import importlib.util
|
|
import sys
|
|
import tempfile
|
|
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_service():
|
|
specification = importlib.util.spec_from_file_location(
|
|
"snapform_service",
|
|
PROJECT_ROOT / "src" / "snapform" / "service.py",
|
|
)
|
|
module = importlib.util.module_from_spec(specification)
|
|
specification.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class FakeCuda:
|
|
class OutOfMemoryError(RuntimeError):
|
|
pass
|
|
|
|
@staticmethod
|
|
def is_available():
|
|
return False
|
|
|
|
@staticmethod
|
|
def empty_cache():
|
|
return None
|
|
|
|
|
|
class ServiceTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.service = load_service()
|
|
|
|
def create_cache(self):
|
|
cache = self.service.ModelCache.__new__(self.service.ModelCache)
|
|
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 test_ben2_model_is_reused_for_matching_runtime_key(self):
|
|
loads = []
|
|
|
|
ben2 = ModuleType("ben2_onnx")
|
|
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, {"ben2_onnx": 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.status()["ben2"], "model")
|
|
self.assertEqual(cached_load_ms, 0.0)
|
|
|
|
def test_da3_model_is_reused_for_matching_runtime_key(self):
|
|
loads = []
|
|
|
|
da3 = ModuleType("da3_onnx")
|
|
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,
|
|
{
|
|
"da3_onnx": 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.status()["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("moge_onnx")
|
|
fake_moge.create_session = lambda _model_dir, _device: moge
|
|
with patch.dict(sys.modules, {"moge_onnx": fake_moge}):
|
|
self.assertIs(cache.get_moge_onnx("moge", "cuda"), moge)
|
|
|
|
self.assertIsNone(cache.da3_model)
|
|
self.assertIsNotNone(cache.ben2_model)
|
|
|
|
def test_service_tracks_clients_and_rejects_shared_shutdown(self):
|
|
fake_models = SimpleNamespace(
|
|
status=lambda: {"ben2": "", "da3": ""},
|
|
close=lambda: None,
|
|
)
|
|
with patch.object(self.service, "ModelCache", return_value=fake_models):
|
|
state = self.service.ServiceState("token", "instance")
|
|
|
|
state.register_client("first", "Blender A")
|
|
state.register_client("second", "Blender B")
|
|
|
|
self.assertEqual(state.snapshot()["clients"], 2)
|
|
self.assertEqual(state.request_shutdown("first"), "shared")
|
|
self.assertTrue(
|
|
state.reserve("job", "first", "reconstruct", "status.json")
|
|
)
|
|
self.assertFalse(
|
|
state.reserve("other", "second", "reconstruct", "other.json")
|
|
)
|
|
|
|
def test_service_accepts_moge_command(self):
|
|
self.assertIn("moge", self.service.JOB_COMMANDS)
|
|
|
|
def test_cancelling_job_keeps_service_and_model_cache_alive(self):
|
|
fake_models = SimpleNamespace(
|
|
status=lambda: {"ben2": "BEN2_BASE", "da3": ""},
|
|
close=lambda: None,
|
|
)
|
|
with patch.object(self.service, "ModelCache", return_value=fake_models):
|
|
state = self.service.ServiceState("token", "instance")
|
|
state.register_client("client", "Blender")
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
status_path = Path(directory) / "status.json"
|
|
state.reserve(
|
|
"job",
|
|
"client",
|
|
"remove-background",
|
|
str(status_path),
|
|
)
|
|
|
|
active = state.cancel_job("job", "client")
|
|
self.assertTrue(active["cancel_requested"])
|
|
self.assertTrue(state.is_job_cancelled("job"))
|
|
state.run_job(
|
|
{
|
|
"job_id": "job",
|
|
"command": "remove-background",
|
|
"parameters": {"status": str(status_path)},
|
|
}
|
|
)
|
|
status = self.service.json.loads(
|
|
status_path.read_text(encoding="utf-8")
|
|
)
|
|
|
|
self.assertEqual(status["state"], "cancelled")
|
|
self.assertFalse(state.shutdown_event.is_set())
|
|
self.assertEqual(state.snapshot()["cached_models"]["ben2"], "BEN2_BASE")
|
|
|
|
def test_expired_clients_are_removed(self):
|
|
fake_models = SimpleNamespace(
|
|
status=lambda: {"ben2": "", "da3": ""},
|
|
close=lambda: None,
|
|
)
|
|
with patch.object(self.service, "ModelCache", return_value=fake_models):
|
|
state = self.service.ServiceState("token", "instance")
|
|
state.register_client("client", "Blender")
|
|
state.clients["client"]["last_seen"] -= (
|
|
self.service.CLIENT_TIMEOUT_SECONDS + 1.0
|
|
)
|
|
|
|
state.expire_clients()
|
|
|
|
self.assertEqual(state.snapshot()["clients"], 0)
|
|
|
|
def test_discovery_is_only_removed_by_matching_service(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "service.json"
|
|
args = SimpleNamespace(
|
|
discovery=str(path),
|
|
port=51234,
|
|
token="token",
|
|
instance_id="owner",
|
|
)
|
|
with patch.object(self.service.os, "getpid", return_value=42):
|
|
self.service.write_discovery(args)
|
|
|
|
payload = self.service.json.loads(path.read_text(encoding="utf-8"))
|
|
self.assertEqual(payload["port"], 51234)
|
|
self.assertEqual(payload["pid"], 42)
|
|
self.service.remove_discovery(path, "other")
|
|
self.assertTrue(path.is_file())
|
|
self.service.remove_discovery(path, "owner")
|
|
self.assertFalse(path.exists())
|
|
|
|
def test_service_dispatches_frame_extraction_without_models(self):
|
|
fake_models = SimpleNamespace(
|
|
status=lambda: {"ben2": "", "da3": ""},
|
|
close=lambda: None,
|
|
)
|
|
with patch.object(self.service, "ModelCache", return_value=fake_models):
|
|
state = self.service.ServiceState("token", "instance")
|
|
state.register_client("client", "Blender")
|
|
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()
|
|
result_directory = root / "result"
|
|
status_path = root / "status.json"
|
|
state.reserve(
|
|
"job",
|
|
"client",
|
|
"extract-frames",
|
|
str(status_path),
|
|
)
|
|
|
|
state.run_job(
|
|
{
|
|
"job_id": "job",
|
|
"command": "extract-frames",
|
|
"parameters": {
|
|
"input": str(input_directory),
|
|
"output_dir": str(result_directory),
|
|
"status": str(status_path),
|
|
"video_frames": 16,
|
|
"start_frame": 1,
|
|
},
|
|
}
|
|
)
|
|
status = self.service.json.loads(
|
|
status_path.read_text(encoding="utf-8")
|
|
)
|
|
self.assertTrue(
|
|
(result_directory / "frame.0001.jpg").is_file()
|
|
)
|
|
|
|
self.assertEqual(status["state"], "succeeded")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|