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:
2739
tests/test_addon.py
2739
tests/test_addon.py
File diff suppressed because it is too large
Load Diff
@@ -11,11 +11,11 @@ SOURCE_DIR = (
|
||||
Path(__file__).resolve().parents[1]
|
||||
/ "src"
|
||||
/ "snapform"
|
||||
/ "ai_service"
|
||||
/ "server"
|
||||
)
|
||||
sys.path.insert(0, str(SOURCE_DIR))
|
||||
|
||||
import ben2_onnx
|
||||
import onnx_ben2
|
||||
|
||||
|
||||
class FakeSession:
|
||||
@@ -35,7 +35,7 @@ class FakeSession:
|
||||
class Ben2OnnxTest(unittest.TestCase):
|
||||
def test_preprocess_matches_model_layout(self):
|
||||
image = Image.new("RGB", (40, 20), (64, 128, 192))
|
||||
result = ben2_onnx.preprocess(image)
|
||||
result = onnx_ben2.preprocess(image)
|
||||
|
||||
self.assertEqual(result.shape, (1, 3, 1024, 1024))
|
||||
self.assertEqual(result.dtype, np.float32)
|
||||
@@ -45,7 +45,7 @@ class Ben2OnnxTest(unittest.TestCase):
|
||||
session = FakeSession()
|
||||
image = Image.new("RGB", (40, 20), (10, 20, 30))
|
||||
|
||||
result = ben2_onnx.infer(session, image)
|
||||
result = onnx_ben2.infer(session, image)
|
||||
|
||||
self.assertEqual(result.mode, "RGBA")
|
||||
self.assertEqual(result.size, image.size)
|
||||
|
||||
@@ -12,11 +12,11 @@ SOURCE_DIR = (
|
||||
Path(__file__).resolve().parents[1]
|
||||
/ "src"
|
||||
/ "snapform"
|
||||
/ "ai_service"
|
||||
/ "server"
|
||||
)
|
||||
sys.path.insert(0, str(SOURCE_DIR))
|
||||
|
||||
import da3_onnx
|
||||
import onnx_da3
|
||||
|
||||
|
||||
class FakeSession:
|
||||
@@ -37,7 +37,7 @@ class FakeSession:
|
||||
class Da3OnnxTest(unittest.TestCase):
|
||||
def test_target_size_uses_patch_multiple(self):
|
||||
self.assertEqual(
|
||||
da3_onnx.target_size((800, 600), 504, "upper_bound_resize"),
|
||||
onnx_da3.target_size((800, 600), 504, "upper_bound_resize"),
|
||||
(504, 378),
|
||||
)
|
||||
|
||||
@@ -47,7 +47,7 @@ class Da3OnnxTest(unittest.TestCase):
|
||||
Image.new("RGB", (80, 60), (10, 20, 30)).save(image_path)
|
||||
session = FakeSession()
|
||||
|
||||
result = da3_onnx.infer(
|
||||
result = onnx_da3.infer(
|
||||
session,
|
||||
[image_path],
|
||||
process_res=56,
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
@@ -10,14 +8,16 @@ from unittest.mock import patch
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def load_worker():
|
||||
specification = importlib.util.spec_from_file_location(
|
||||
"snapform_download_worker",
|
||||
PROJECT_ROOT / "src" / "snapform" / "ai_service" / "jobs.py",
|
||||
)
|
||||
module = importlib.util.module_from_spec(specification)
|
||||
specification.loader.exec_module(module)
|
||||
return module
|
||||
def load_download_module():
|
||||
addon_directory = PROJECT_ROOT / "src" / "snapform"
|
||||
sys.path.insert(0, str(addon_directory))
|
||||
try:
|
||||
return __import__(
|
||||
"server.model_download",
|
||||
fromlist=["model_download"],
|
||||
)
|
||||
finally:
|
||||
sys.path.remove(str(addon_directory))
|
||||
|
||||
|
||||
class FakeTqdm:
|
||||
@@ -36,52 +36,41 @@ class FakeTqdm:
|
||||
class DownloadProgressTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.worker = load_worker()
|
||||
cls.download = load_download_module()
|
||||
|
||||
def create_progress(self, description="Downloading bytes", **options):
|
||||
reports = []
|
||||
progress_class = self.download.create_download_progress_class(
|
||||
lambda value, message: reports.append((value, message)),
|
||||
"owner/model",
|
||||
tqdm_type=options.pop("tqdm_type", FakeTqdm),
|
||||
clock=options.pop("clock", lambda: 1.0),
|
||||
**options,
|
||||
)
|
||||
progress = progress_class(
|
||||
desc=description,
|
||||
total=100,
|
||||
initial=0,
|
||||
unit="B",
|
||||
)
|
||||
return progress, reports
|
||||
|
||||
def test_reports_aggregated_download_bytes(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
progress_class = self.worker.create_download_progress_class(
|
||||
status_path,
|
||||
"owner/model",
|
||||
tqdm_type=FakeTqdm,
|
||||
clock=lambda: 1.0,
|
||||
)
|
||||
progress = progress_class(
|
||||
desc="Downloading bytes",
|
||||
total=100,
|
||||
initial=0,
|
||||
unit="B",
|
||||
)
|
||||
progress, reports = self.create_progress()
|
||||
|
||||
progress.update(25)
|
||||
progress.update(25)
|
||||
|
||||
status = json.loads(status_path.read_text(encoding="utf-8"))
|
||||
self.assertAlmostEqual(status["progress"], 0.275)
|
||||
self.assertEqual(status["message"], "Downloading owner/model")
|
||||
self.assertAlmostEqual(reports[-1][0], 0.275)
|
||||
self.assertEqual(reports[-1][1], "Downloading owner/model")
|
||||
|
||||
def test_reports_byte_progress_for_aggregated_bar_description(self):
|
||||
# huggingface_hub 1.x 用 "Downloading (incomplete total...)" 作为字节进度条描述
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
progress_class = self.worker.create_download_progress_class(
|
||||
status_path,
|
||||
"owner/model",
|
||||
tqdm_type=FakeTqdm,
|
||||
clock=lambda: 1.0,
|
||||
)
|
||||
progress = progress_class(
|
||||
desc="Downloading (incomplete total...)",
|
||||
total=100,
|
||||
initial=0,
|
||||
unit="B",
|
||||
)
|
||||
def test_reports_hugging_face_aggregated_bar(self):
|
||||
progress, reports = self.create_progress(
|
||||
"Downloading (incomplete total...)"
|
||||
)
|
||||
|
||||
progress.update(25)
|
||||
progress.update(25)
|
||||
|
||||
status = json.loads(status_path.read_text(encoding="utf-8"))
|
||||
self.assertAlmostEqual(status["progress"], 0.275)
|
||||
self.assertEqual(status["message"], "Downloading owner/model")
|
||||
self.assertAlmostEqual(reports[-1][0], 0.275)
|
||||
|
||||
def test_reports_download_speed(self):
|
||||
class FakeTqdmWithRate(FakeTqdm):
|
||||
@@ -89,129 +78,48 @@ class DownloadProgressTest(unittest.TestCase):
|
||||
def format_dict(self):
|
||||
return {"rate": 14_000_000}
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
progress_class = self.worker.create_download_progress_class(
|
||||
status_path,
|
||||
"owner/model",
|
||||
tqdm_type=FakeTqdmWithRate,
|
||||
clock=lambda: 1.0,
|
||||
label="da3 small",
|
||||
)
|
||||
progress = progress_class(
|
||||
desc="Downloading bytes",
|
||||
total=100,
|
||||
initial=0,
|
||||
unit="B",
|
||||
)
|
||||
progress, reports = self.create_progress(
|
||||
tqdm_type=FakeTqdmWithRate,
|
||||
label="da3 small",
|
||||
)
|
||||
|
||||
progress.update(50)
|
||||
progress.update(50)
|
||||
|
||||
status = json.loads(status_path.read_text(encoding="utf-8"))
|
||||
self.assertEqual(
|
||||
status["message"],
|
||||
"Downloading da3 small 13.4 MB/s",
|
||||
)
|
||||
self.assertEqual(
|
||||
reports[-1][1],
|
||||
"Downloading da3 small 13.4 MB/s",
|
||||
)
|
||||
|
||||
def test_does_not_report_estimated_time_remaining(self):
|
||||
def test_progress_callback_has_no_eta_or_state_protocol(self):
|
||||
times = iter((0.0, 10.0))
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
progress_class = self.worker.create_download_progress_class(
|
||||
status_path,
|
||||
"owner/model",
|
||||
tqdm_type=FakeTqdm,
|
||||
clock=lambda: next(times),
|
||||
)
|
||||
progress = progress_class(
|
||||
desc="Downloading bytes",
|
||||
total=100,
|
||||
initial=0,
|
||||
unit="B",
|
||||
)
|
||||
progress, reports = self.create_progress(clock=lambda: next(times))
|
||||
|
||||
progress.update(50)
|
||||
progress.update(50)
|
||||
|
||||
status = json.loads(status_path.read_text(encoding="utf-8"))
|
||||
self.assertNotIn("eta_seconds", status)
|
||||
|
||||
def test_staged_download_keeps_model_label(self):
|
||||
class FakeTqdmWithRate(FakeTqdm):
|
||||
@property
|
||||
def format_dict(self):
|
||||
return {"rate": 14_000_000}
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
progress_class = self.worker.create_download_progress_class(
|
||||
status_path,
|
||||
"owner/model",
|
||||
tqdm_type=FakeTqdmWithRate,
|
||||
clock=lambda: 1.0,
|
||||
label="da3 small",
|
||||
stage=6,
|
||||
stages=8,
|
||||
stage_label="da3 small",
|
||||
)
|
||||
progress = progress_class(
|
||||
desc="Downloading bytes",
|
||||
total=100,
|
||||
initial=0,
|
||||
unit="B",
|
||||
)
|
||||
|
||||
progress.update(50)
|
||||
|
||||
status = json.loads(status_path.read_text(encoding="utf-8"))
|
||||
self.assertEqual(
|
||||
status["message"],
|
||||
"Downloading da3 small 13.4 MB/s",
|
||||
)
|
||||
self.assertEqual(status["stage"], 6)
|
||||
self.assertEqual(status["stages"], 8)
|
||||
self.assertEqual(status["stage_label"], "da3 small")
|
||||
self.assertEqual(len(reports[-1]), 2)
|
||||
|
||||
def test_ignores_unrelated_progress_bars(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
status_path = Path(directory) / "status.json"
|
||||
progress_class = self.worker.create_download_progress_class(
|
||||
status_path,
|
||||
"owner/model",
|
||||
tqdm_type=FakeTqdm,
|
||||
clock=lambda: 1.0,
|
||||
)
|
||||
progress = progress_class(
|
||||
desc="Reconstructing",
|
||||
total=100,
|
||||
initial=0,
|
||||
unit="B",
|
||||
)
|
||||
progress, reports = self.create_progress("Reconstructing")
|
||||
|
||||
progress.update(25)
|
||||
progress.update(25)
|
||||
|
||||
self.assertFalse(status_path.exists())
|
||||
self.assertEqual(reports, [])
|
||||
|
||||
def test_download_endpoints_uses_direct_endpoint(self):
|
||||
with patch.dict(os.environ, {"HF_ENDPOINT": ""}):
|
||||
endpoints = self.worker.download_endpoints()
|
||||
endpoints = self.download.download_endpoints()
|
||||
|
||||
self.assertEqual(
|
||||
endpoints,
|
||||
["https://huggingface.co"],
|
||||
)
|
||||
self.assertEqual(endpoints, ["https://huggingface.co"])
|
||||
|
||||
def test_download_endpoints_respects_environment_override(self):
|
||||
with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example"}):
|
||||
endpoints = self.worker.download_endpoints()
|
||||
endpoints = self.download.download_endpoints()
|
||||
|
||||
self.assertEqual(
|
||||
endpoints,
|
||||
["https://hf.example"],
|
||||
)
|
||||
self.assertEqual(endpoints, ["https://hf.example"])
|
||||
|
||||
def test_download_endpoints_strips_trailing_slash(self):
|
||||
with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example/"}):
|
||||
endpoints = self.worker.download_endpoints()
|
||||
endpoints = self.download.download_endpoints()
|
||||
|
||||
self.assertEqual(endpoints, ["https://hf.example"])
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -21,7 +21,7 @@ class ModelTest(unittest.TestCase):
|
||||
def setUpClass(cls):
|
||||
cls.models = load_module(
|
||||
"snapform_models",
|
||||
"src/snapform/models.py",
|
||||
"src/snapform/server/model_catalog.py",
|
||||
)
|
||||
|
||||
def test_base_model_is_default_friendly_apache_model(self):
|
||||
|
||||
@@ -12,8 +12,8 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
def load_module():
|
||||
specification = importlib.util.spec_from_file_location(
|
||||
"snapform_moge_onnx",
|
||||
PROJECT_ROOT / "src" / "snapform" / "ai_service" / "moge_onnx.py",
|
||||
"snapform_onnx_moge2",
|
||||
PROJECT_ROOT / "src" / "snapform" / "server" / "onnx_moge2.py",
|
||||
)
|
||||
module = importlib.util.module_from_spec(specification)
|
||||
specification.loader.exec_module(module)
|
||||
|
||||
@@ -58,8 +58,23 @@ class PackTest(unittest.TestCase):
|
||||
"assets/O_SnapForm.blend",
|
||||
names,
|
||||
)
|
||||
self.assertIn("ai_service/jobs.py", names)
|
||||
self.assertIn("ai_service/server.py", names)
|
||||
self.assertIn("server/media.py", names)
|
||||
self.assertIn("server/geometry.py", names)
|
||||
self.assertIn("server/background.py", names)
|
||||
self.assertIn("server/moge2.py", names)
|
||||
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("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.assertNotIn("install_env.py", names)
|
||||
self.assertNotIn("snapform/__init__.py", names)
|
||||
self.assertNotIn("importer.py", names)
|
||||
self.assertFalse(
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
import importlib.util
|
||||
import importlib
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import ModuleType, SimpleNamespace
|
||||
@@ -10,14 +12,13 @@ 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" / "ai_service" / "server.py",
|
||||
)
|
||||
module = importlib.util.module_from_spec(specification)
|
||||
specification.loader.exec_module(module)
|
||||
return module
|
||||
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:
|
||||
@@ -33,19 +34,19 @@ class FakeCuda:
|
||||
return None
|
||||
|
||||
|
||||
class ServiceTest(unittest.TestCase):
|
||||
class ServerTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.service = load_service()
|
||||
cls.server_module = load_server_module()
|
||||
|
||||
def setUp(self):
|
||||
self.jobs_temp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.jobs_temp.cleanup)
|
||||
self.jobs_directory = Path(self.jobs_temp.name) / "jobs"
|
||||
self.storage_root = Path(self.jobs_temp.name)
|
||||
|
||||
def create_cache(self):
|
||||
cache = self.service.ModelSessionCache.__new__(
|
||||
self.service.ModelSessionCache
|
||||
cache = self.server_module.ModelSessionCache.__new__(
|
||||
self.server_module.ModelSessionCache
|
||||
)
|
||||
cache.ben2_model = None
|
||||
cache.ben2_key = None
|
||||
@@ -55,16 +56,34 @@ class ServiceTest(unittest.TestCase):
|
||||
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("ben2_onnx")
|
||||
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, {"ben2_onnx": ben2}):
|
||||
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)
|
||||
|
||||
@@ -76,7 +95,7 @@ class ServiceTest(unittest.TestCase):
|
||||
def test_da3_model_is_reused_for_matching_runtime_key(self):
|
||||
loads = []
|
||||
|
||||
da3 = ModuleType("da3_onnx")
|
||||
da3 = ModuleType("server.onnx_da3")
|
||||
da3.create_session = lambda model_dir, device: (
|
||||
loads.append((model_dir, device)) or object()
|
||||
)
|
||||
@@ -85,7 +104,7 @@ class ServiceTest(unittest.TestCase):
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"da3_onnx": da3,
|
||||
"server.onnx_da3": da3,
|
||||
},
|
||||
):
|
||||
first = cache.get_da3("model", device)
|
||||
@@ -117,34 +136,44 @@ class ServiceTest(unittest.TestCase):
|
||||
|
||||
moge = object()
|
||||
|
||||
fake_moge = ModuleType("moge_onnx")
|
||||
fake_moge = ModuleType("server.onnx_moge2")
|
||||
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)
|
||||
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_service_runs_only_one_job_at_a_time(self):
|
||||
fake_models = SimpleNamespace(
|
||||
snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}},
|
||||
clear=lambda: None,
|
||||
close=lambda: None,
|
||||
)
|
||||
with patch.object(self.service, "ModelManager", return_value=fake_models):
|
||||
server = self.service.create_server(jobs_directory=self.jobs_directory)
|
||||
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 = []
|
||||
|
||||
parameters = {"model": "BASE"}
|
||||
first = server.submit("debug-da3", parameters, job_id="job")
|
||||
second = server.submit("debug-da3", parameters, job_id="other")
|
||||
@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.assertIsNone(second)
|
||||
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_service_registers_tasks_instead_of_model_names(self):
|
||||
server = self.service.create_server(jobs_directory=self.jobs_directory)
|
||||
def test_server_registers_tasks_instead_of_model_names(self):
|
||||
server = self.create_server()
|
||||
self.assertEqual(
|
||||
set(server.handlers),
|
||||
{
|
||||
@@ -161,7 +190,33 @@ class ServiceTest(unittest.TestCase):
|
||||
},
|
||||
)
|
||||
|
||||
def test_cancelling_job_keeps_service_and_model_cache_alive(self):
|
||||
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": ""}
|
||||
@@ -169,36 +224,55 @@ class ServiceTest(unittest.TestCase):
|
||||
clear=lambda: None,
|
||||
close=lambda: None,
|
||||
)
|
||||
with patch.object(self.service, "ModelManager", return_value=fake_models):
|
||||
server = self.service.create_server(jobs_directory=self.jobs_directory)
|
||||
with patch.object(
|
||||
self.server_module, "ModelManager", return_value=fake_models
|
||||
):
|
||||
server = self.create_server()
|
||||
context = server.submit("remove-background", {}, job_id="job")
|
||||
self.assertTrue(context.request_cancel())
|
||||
self.assertIsNotNone(server.cancel("job"))
|
||||
self.assertTrue(context.is_cancelled())
|
||||
server.run(context)
|
||||
status = context.snapshot()
|
||||
status = self.wait_for_job(context)
|
||||
|
||||
self.assertEqual(status["state"], "cancelled")
|
||||
self.assertFalse(server.shutdown_event.is_set())
|
||||
self.assertEqual(
|
||||
server.snapshot("instance")["resources"]["models"]["loaded"]["ben2"],
|
||||
server.snapshot("instance")["resources"]["model_manager"]["loaded"]["ben2"],
|
||||
"BEN2_BASE",
|
||||
)
|
||||
|
||||
def test_job_progress_is_kept_in_service_memory(self):
|
||||
server = self.service.JobServer("Test", jobs_directory=self.jobs_directory)
|
||||
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.update(0.5, "Half way")
|
||||
context.progress(0.5, "Half way")
|
||||
|
||||
context = server.submit("example", {})
|
||||
server.run(context)
|
||||
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(
|
||||
@@ -206,33 +280,40 @@ class ServiceTest(unittest.TestCase):
|
||||
clear=lambda: events.append("clear"),
|
||||
close=lambda: events.append("close"),
|
||||
)
|
||||
server = self.service.JobServer("Test", jobs_directory=self.jobs_directory)
|
||||
server.add_resource("models", resource)
|
||||
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("models"), resource)
|
||||
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("models"))
|
||||
server.run(context)
|
||||
self.assertTrue(server.clear_resource("models"))
|
||||
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"]["models"]["loaded"],
|
||||
server.snapshot("instance")["resources"]["model_manager"]["loaded"],
|
||||
"model",
|
||||
)
|
||||
self.assertEqual(events, ["clear", "close"])
|
||||
|
||||
def test_service_dispatches_frame_extraction_without_models(self):
|
||||
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.service, "ModelManager", return_value=fake_models):
|
||||
server = self.service.create_server(jobs_directory=self.jobs_directory)
|
||||
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"
|
||||
@@ -248,8 +329,7 @@ class ServiceTest(unittest.TestCase):
|
||||
},
|
||||
job_id="job",
|
||||
)
|
||||
server.run(context)
|
||||
status = context.snapshot()
|
||||
status = self.wait_for_job(context)
|
||||
self.assertTrue(
|
||||
(context.directory / "frame.0001.jpg").is_file()
|
||||
)
|
||||
@@ -1,7 +1,5 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
@@ -16,14 +14,13 @@ import numpy as np
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def load_worker():
|
||||
specification = importlib.util.spec_from_file_location(
|
||||
"snapform_worker",
|
||||
PROJECT_ROOT / "src" / "snapform" / "ai_service" / "jobs.py",
|
||||
)
|
||||
module = importlib.util.module_from_spec(specification)
|
||||
specification.loader.exec_module(module)
|
||||
return module
|
||||
def load_server_module(name):
|
||||
addon_directory = PROJECT_ROOT / "src" / "snapform"
|
||||
sys.path.insert(0, str(addon_directory))
|
||||
try:
|
||||
return __import__(f"server.{name}", fromlist=[name])
|
||||
finally:
|
||||
sys.path.remove(str(addon_directory))
|
||||
|
||||
|
||||
def _make_fake_av(frame_count):
|
||||
@@ -69,10 +66,54 @@ def _make_fake_av(frame_count):
|
||||
return FakeAv()
|
||||
|
||||
|
||||
class FakeModels:
|
||||
def __init__(self, directory="model"):
|
||||
self.model_directory = Path(directory)
|
||||
|
||||
def directory(self, _model):
|
||||
return self.model_directory
|
||||
|
||||
def get_ben2(self, _model_directory, _device):
|
||||
from server.onnx_ben2 import create_session
|
||||
|
||||
return create_session(self.model_directory, "cpu"), 0.0
|
||||
|
||||
def get_da3(self, _model_directory, _device):
|
||||
return object()
|
||||
|
||||
|
||||
class FakeJobContext:
|
||||
def __init__(self, directory, job_type, models=None):
|
||||
self.directory = Path(directory)
|
||||
self.job_type = job_type
|
||||
self.models = models or FakeModels()
|
||||
self.status = {}
|
||||
|
||||
def resource(self, name):
|
||||
if name != "model_manager":
|
||||
raise KeyError(name)
|
||||
return self.models
|
||||
|
||||
def progress(self, progress, message, **details):
|
||||
self.status = {
|
||||
"progress": progress,
|
||||
"message": message,
|
||||
**details,
|
||||
}
|
||||
|
||||
def check_cancelled(self):
|
||||
return None
|
||||
|
||||
|
||||
class WorkerTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.worker = load_worker()
|
||||
cls.media = load_server_module("media")
|
||||
cls.geometry = load_server_module("geometry")
|
||||
cls.background = load_server_module("background")
|
||||
cls.moge2 = load_server_module("moge2")
|
||||
cls.da3 = load_server_module("da3")
|
||||
cls.download = load_server_module("model_download")
|
||||
|
||||
def test_foreground_preserves_source_rgb_and_ben2_alpha(self):
|
||||
alpha = SimpleNamespace(size=(2, 1))
|
||||
@@ -88,7 +129,7 @@ class WorkerTest(unittest.TestCase):
|
||||
)
|
||||
foreground = SimpleNamespace(convert=lambda _mode: foreground_rgba)
|
||||
|
||||
result = self.worker.preserve_source_rgb(source, foreground)
|
||||
result = self.background.preserve_source_rgb(source, foreground)
|
||||
|
||||
self.assertIs(result, source_rgba)
|
||||
self.assertIs(result.applied_alpha, alpha)
|
||||
@@ -100,62 +141,22 @@ class WorkerTest(unittest.TestCase):
|
||||
(root / "a.jpg").touch()
|
||||
(root / "notes.txt").touch()
|
||||
|
||||
images = self.worker.collect_images(root)
|
||||
images = self.media.collect_images(root)
|
||||
|
||||
self.assertEqual([path.name for path in images], ["a.jpg", "b.PNG"])
|
||||
|
||||
def test_ai_tasks_support_cooperative_cancellation(self):
|
||||
with self.assertRaises(self.worker.TaskCancelled):
|
||||
self.worker.check_cancelled(lambda: True)
|
||||
with self.assertRaises(self.media.TaskCancelled):
|
||||
self.media.check_cancelled(lambda: True)
|
||||
|
||||
self.worker.check_cancelled(lambda: False)
|
||||
|
||||
def test_reconstruct_parser_accepts_background_model_directory(self):
|
||||
args = self.worker.parser().parse_args(
|
||||
[
|
||||
"reconstruct",
|
||||
"--input",
|
||||
"input.png",
|
||||
"--model-dir",
|
||||
"da3",
|
||||
"--output-dir",
|
||||
"result",
|
||||
"--export-format",
|
||||
"mini_npz",
|
||||
"--status",
|
||||
"status.json",
|
||||
"--remove-background-model-dir",
|
||||
"ben2",
|
||||
"--process-res",
|
||||
"1024",
|
||||
"--process-res-method",
|
||||
"upper_bound_resize",
|
||||
"--conf-thresh-percentile",
|
||||
"40",
|
||||
"--num-max-points",
|
||||
"1000",
|
||||
"--show-cameras",
|
||||
"0",
|
||||
"--gs-views-interval",
|
||||
"1",
|
||||
"--video-frames",
|
||||
"16",
|
||||
"--device",
|
||||
"auto",
|
||||
"--use-ray-pose",
|
||||
"1",
|
||||
"--ref-view-strategy",
|
||||
"first",
|
||||
]
|
||||
)
|
||||
self.assertEqual(args.remove_background_model_dir, "ben2")
|
||||
self.media.check_cancelled(lambda: False)
|
||||
|
||||
def test_prepares_a_single_image_file(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
image_path = Path(directory) / "input.png"
|
||||
image_path.touch()
|
||||
|
||||
images, temporary_directory = self.worker.prepare_input(
|
||||
images, temporary_directory = self.media.prepare_input(
|
||||
image_path,
|
||||
16,
|
||||
directory,
|
||||
@@ -171,7 +172,7 @@ class WorkerTest(unittest.TestCase):
|
||||
(root / "frame-0002.png").touch()
|
||||
(root / "notes.txt").touch()
|
||||
|
||||
images, temporary_directory = self.worker.prepare_input(
|
||||
images, temporary_directory = self.media.prepare_input(
|
||||
root,
|
||||
16,
|
||||
directory,
|
||||
@@ -189,7 +190,7 @@ class WorkerTest(unittest.TestCase):
|
||||
for index in range(1, 6):
|
||||
(root / f"frame-{index:04d}.jpg").touch()
|
||||
|
||||
images, temporary_directory = self.worker.prepare_input(
|
||||
images, temporary_directory = self.media.prepare_input(
|
||||
root,
|
||||
3,
|
||||
directory,
|
||||
@@ -207,7 +208,7 @@ class WorkerTest(unittest.TestCase):
|
||||
for index in range(1, 6):
|
||||
(root / f"frame-{index:04d}.jpg").touch()
|
||||
|
||||
images, temporary_directory = self.worker.prepare_input(
|
||||
images, temporary_directory = self.media.prepare_input(
|
||||
root,
|
||||
3,
|
||||
directory,
|
||||
@@ -224,7 +225,7 @@ class WorkerTest(unittest.TestCase):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
fake_av = _make_fake_av(100)
|
||||
image_paths = self.worker.extract_video_frames(
|
||||
image_paths = self.media.extract_video_frames(
|
||||
root / "clip.mkv",
|
||||
root / "out",
|
||||
10,
|
||||
@@ -242,7 +243,7 @@ class WorkerTest(unittest.TestCase):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
fake_av = _make_fake_av(100)
|
||||
image_paths = self.worker.extract_video_frames(
|
||||
image_paths = self.media.extract_video_frames(
|
||||
root / "clip.mkv",
|
||||
root / "out",
|
||||
10,
|
||||
@@ -260,7 +261,7 @@ class WorkerTest(unittest.TestCase):
|
||||
def test_extract_video_frames_tolerates_short_video_tail(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
image_paths = self.worker.extract_video_frames(
|
||||
image_paths = self.media.extract_video_frames(
|
||||
root / "clip.mkv",
|
||||
root / "out",
|
||||
10,
|
||||
@@ -280,11 +281,11 @@ class WorkerTest(unittest.TestCase):
|
||||
valid_path,
|
||||
depth=np.zeros((4, 4)),
|
||||
)
|
||||
self.assertTrue(self.worker.npz_is_complete(valid_path))
|
||||
self.assertTrue(self.media.npz_is_complete(valid_path))
|
||||
|
||||
empty_path = root / "empty.npz"
|
||||
empty_path.write_bytes(b"PK\x05\x06" + b"\x00" * 18)
|
||||
self.assertFalse(self.worker.npz_is_complete(empty_path))
|
||||
self.assertFalse(self.media.npz_is_complete(empty_path))
|
||||
|
||||
def test_wait_for_exported_paths_waits_for_complete_npz(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
@@ -293,11 +294,11 @@ class WorkerTest(unittest.TestCase):
|
||||
npz_path.parent.mkdir(parents=True)
|
||||
npz_path.touch()
|
||||
with patch.object(
|
||||
self.worker,
|
||||
self.media,
|
||||
"npz_is_complete",
|
||||
side_effect=[False, True],
|
||||
):
|
||||
result = self.worker.wait_for_exported_paths(
|
||||
result = self.media.wait_for_exported_paths(
|
||||
root,
|
||||
"mini_npz",
|
||||
timeout=2.0,
|
||||
@@ -324,10 +325,10 @@ class WorkerTest(unittest.TestCase):
|
||||
return original_replace(path, target)
|
||||
|
||||
with patch.object(path_type, "replace", new=locked_once), patch.object(
|
||||
self.worker.time,
|
||||
self.media.time,
|
||||
"sleep",
|
||||
) as sleep:
|
||||
result = self.worker.move_exported_file(
|
||||
result = self.media.move_exported_file(
|
||||
source,
|
||||
destination,
|
||||
)
|
||||
@@ -342,7 +343,7 @@ class WorkerTest(unittest.TestCase):
|
||||
input_path.touch()
|
||||
|
||||
with self.assertRaisesRegex(RuntimeError, "Unsupported input type"):
|
||||
self.worker.prepare_input(input_path, 16, directory)
|
||||
self.media.prepare_input(input_path, 16, directory)
|
||||
|
||||
def test_download_model_uses_actual_request_without_preflight(self):
|
||||
calls = []
|
||||
@@ -353,7 +354,7 @@ class WorkerTest(unittest.TestCase):
|
||||
args = SimpleNamespace(
|
||||
repository="depth-anything/DA3-SMALL",
|
||||
model_dir="model",
|
||||
status="status.json",
|
||||
progress=lambda *_args, **_kwargs: None,
|
||||
allow_pattern=[],
|
||||
)
|
||||
|
||||
@@ -361,17 +362,16 @@ class WorkerTest(unittest.TestCase):
|
||||
sys.modules,
|
||||
{"huggingface_hub": fake_hugging_face},
|
||||
), patch.object(
|
||||
self.worker,
|
||||
self.download,
|
||||
"create_download_progress_class",
|
||||
return_value=object,
|
||||
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
||||
self.worker,
|
||||
self.download,
|
||||
"_verify_model_complete",
|
||||
return_value=None,
|
||||
):
|
||||
args.model_dir = str(Path(directory) / "models" / "DA3")
|
||||
args.status = str(Path(directory) / "status.json")
|
||||
self.worker.download_model(args)
|
||||
self.download.download_model(args)
|
||||
|
||||
self.assertEqual(len(calls), 1)
|
||||
|
||||
@@ -381,7 +381,7 @@ class WorkerTest(unittest.TestCase):
|
||||
args = SimpleNamespace(
|
||||
repository="PramaLLC/BEN2",
|
||||
model_dir="model",
|
||||
status="status.json",
|
||||
progress=lambda *_args, **_kwargs: None,
|
||||
allow_pattern=[],
|
||||
)
|
||||
|
||||
@@ -389,47 +389,45 @@ class WorkerTest(unittest.TestCase):
|
||||
sys.modules,
|
||||
{"huggingface_hub": fake_hugging_face},
|
||||
), patch.object(
|
||||
self.worker,
|
||||
self.download,
|
||||
"create_download_progress_class",
|
||||
return_value=object,
|
||||
), patch.dict(os.environ, {"HF_ENDPOINT": ""}):
|
||||
args.model_dir = str(Path(directory) / "model")
|
||||
args.status = str(Path(directory) / "status.json")
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"All download sources failed",
|
||||
):
|
||||
self.worker.download_model(args)
|
||||
self.download.download_model(args)
|
||||
|
||||
def test_download_model_reports_all_failed_sources(self):
|
||||
args = SimpleNamespace(
|
||||
repository="depth-anything/DA3-SMALL",
|
||||
model_dir="model",
|
||||
status="status.json",
|
||||
progress=lambda *_args, **_kwargs: None,
|
||||
allow_pattern=[],
|
||||
r2_directory="DA3Small",
|
||||
r2_file=["config.json|100|abc"],
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory, patch.object(
|
||||
self.worker,
|
||||
self.download,
|
||||
"download_from_hf",
|
||||
side_effect=OSError("direct down"),
|
||||
), patch.object(
|
||||
self.worker,
|
||||
self.download,
|
||||
"download_from_r2",
|
||||
side_effect=OSError("mirror down"),
|
||||
):
|
||||
args.model_dir = str(Path(directory) / "model")
|
||||
args.status = str(Path(directory) / "status.json")
|
||||
with self.assertRaisesRegex(RuntimeError, "models.omoolab.xyz"):
|
||||
self.worker.download_model(args)
|
||||
self.download.download_model(args)
|
||||
|
||||
def test_download_model_prefers_direct_hf_then_r2(self):
|
||||
args = SimpleNamespace(
|
||||
repository="depth-anything/DA3-SMALL",
|
||||
model_dir="model",
|
||||
status="status.json",
|
||||
progress=lambda *_args, **_kwargs: None,
|
||||
allow_pattern=["config.json"],
|
||||
r2_directory="DA3Small",
|
||||
r2_file=["config.json|100|abc"],
|
||||
@@ -444,83 +442,28 @@ class WorkerTest(unittest.TestCase):
|
||||
order.append("r2")
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory, patch.object(
|
||||
self.worker,
|
||||
self.download,
|
||||
"download_from_hf",
|
||||
side_effect=fake_download_from_hf,
|
||||
), patch.object(
|
||||
self.worker,
|
||||
self.download,
|
||||
"download_from_r2",
|
||||
side_effect=fake_download_from_r2,
|
||||
), patch.object(
|
||||
self.worker,
|
||||
self.download,
|
||||
"_verify_model_complete",
|
||||
return_value=None,
|
||||
):
|
||||
args.model_dir = str(Path(directory) / "models" / "DA3")
|
||||
args.status = str(Path(directory) / "status.json")
|
||||
self.worker.download_model(args)
|
||||
self.download.download_model(args)
|
||||
|
||||
self.assertEqual(order, ["hf", "r2"])
|
||||
|
||||
def test_download_declared_skips_hf_when_r2_preferred(self):
|
||||
root = PROJECT_ROOT / ".tmp-models-from-r2"
|
||||
model_dir = root / "models" / "DA3"
|
||||
args = SimpleNamespace(
|
||||
repository="depth-anything/DA3-SMALL",
|
||||
model_dir=str(model_dir),
|
||||
status="status.json",
|
||||
allow_pattern=[],
|
||||
r2_directory="DA3Small",
|
||||
r2_file=["config.json|100|abc"],
|
||||
)
|
||||
order = []
|
||||
try:
|
||||
self.worker._mark_models_from_r2(model_dir)
|
||||
with patch.object(
|
||||
self.worker,
|
||||
"download_from_hf",
|
||||
side_effect=lambda *a, **k: order.append("hf"),
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"download_from_r2",
|
||||
side_effect=lambda *a, **k: order.append("r2"),
|
||||
), patch.object(
|
||||
self.worker,
|
||||
"_verify_model_complete",
|
||||
return_value=None,
|
||||
), patch.object(self.worker, "write_status"):
|
||||
self.worker._download_declared_or_fail(
|
||||
args, model_dir, "DA3 Small", 0.0, 1.0, {}
|
||||
)
|
||||
finally:
|
||||
shutil.rmtree(root, ignore_errors=True)
|
||||
|
||||
self.assertEqual(order, ["r2"])
|
||||
|
||||
def test_models_from_r2_roundtrip_preserves_fields(self):
|
||||
root = PROJECT_ROOT / ".tmp-models-from-r2"
|
||||
model_dir = root / "models" / "DA3"
|
||||
runtime_state = root / "runtime.json"
|
||||
try:
|
||||
runtime_state.parent.mkdir(parents=True, exist_ok=True)
|
||||
runtime_state.write_text(
|
||||
json.dumps({"keep": "value"}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
self.assertFalse(self.worker._models_from_r2(model_dir))
|
||||
self.worker._mark_models_from_r2(model_dir)
|
||||
self.assertTrue(self.worker._models_from_r2(model_dir))
|
||||
payload = json.loads(runtime_state.read_text(encoding="utf-8"))
|
||||
self.assertEqual(payload["keep"], "value")
|
||||
self.assertTrue(payload["models_from_r2"])
|
||||
finally:
|
||||
shutil.rmtree(root, ignore_errors=True)
|
||||
|
||||
def test_download_from_r2_downloads_declared_files(self):
|
||||
args = SimpleNamespace(
|
||||
repository="depth-anything/DA3-SMALL",
|
||||
model_dir="model",
|
||||
status="status.json",
|
||||
progress=lambda *_args, **_kwargs: None,
|
||||
allow_pattern=["config.json"],
|
||||
r2_directory="DA3Small",
|
||||
r2_file=["config.json|100|abc"],
|
||||
@@ -528,13 +471,12 @@ class WorkerTest(unittest.TestCase):
|
||||
downloaded = []
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory, patch.object(
|
||||
self.worker,
|
||||
self.download,
|
||||
"_download_file_from_url",
|
||||
side_effect=lambda *args, **kwargs: downloaded.append(args[1].name),
|
||||
), patch.object(self.worker, "write_status"):
|
||||
):
|
||||
args.model_dir = str(Path(directory) / "model")
|
||||
args.status = str(Path(directory) / "status.json")
|
||||
self.worker.download_from_r2(args, args.model_dir)
|
||||
self.download.download_from_r2(args, args.model_dir)
|
||||
|
||||
self.assertEqual(downloaded, ["config.json"])
|
||||
|
||||
@@ -589,7 +531,7 @@ class WorkerTest(unittest.TestCase):
|
||||
)
|
||||
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
|
||||
fake_torch.inference_mode = nullcontext
|
||||
fake_ben2 = ModuleType("ben2_onnx")
|
||||
fake_ben2 = ModuleType("server.onnx_ben2")
|
||||
fake_ben2.create_session = lambda _model_dir, _device: object()
|
||||
fake_ben2.infer = lambda _model, image: (
|
||||
inference_options.append((image, {})) or FakeImage()
|
||||
@@ -601,22 +543,18 @@ class WorkerTest(unittest.TestCase):
|
||||
sys.modules,
|
||||
{
|
||||
"torch": fake_torch,
|
||||
"ben2_onnx": fake_ben2,
|
||||
"server.onnx_ben2": fake_ben2,
|
||||
"PIL": fake_pil,
|
||||
},
|
||||
):
|
||||
root = Path(directory)
|
||||
input_path = root / "input.png"
|
||||
input_path.touch()
|
||||
args = SimpleNamespace(
|
||||
input=str(input_path),
|
||||
model_dir="model",
|
||||
output_dir=str(root / "result"),
|
||||
status=str(root / "status.json"),
|
||||
device="cpu",
|
||||
context = FakeJobContext(root / "result", "remove-background")
|
||||
self.background.remove_background(
|
||||
context,
|
||||
{"input": str(input_path), "model": "BEN2", "device": "cpu"},
|
||||
)
|
||||
self.worker.remove_background(args)
|
||||
status = json.loads(Path(args.status).read_text(encoding="utf-8"))
|
||||
|
||||
self.assertTrue((root / "result" / "foreground.png").is_file())
|
||||
|
||||
@@ -628,7 +566,7 @@ class WorkerTest(unittest.TestCase):
|
||||
save_options[0],
|
||||
{"format": "PNG", "compress_level": 1},
|
||||
)
|
||||
self.assertIn("Background removed in", status["message"])
|
||||
self.assertIn("Background removed in", context.status["message"])
|
||||
|
||||
def test_ben2_background_removal_processes_each_frame(self):
|
||||
inference_options = []
|
||||
@@ -677,7 +615,7 @@ class WorkerTest(unittest.TestCase):
|
||||
)
|
||||
fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
|
||||
fake_torch.inference_mode = nullcontext
|
||||
fake_ben2 = ModuleType("ben2_onnx")
|
||||
fake_ben2 = ModuleType("server.onnx_ben2")
|
||||
fake_ben2.create_session = lambda _model_dir, _device: object()
|
||||
fake_ben2.infer = lambda _model, image: (
|
||||
inference_options.append((image, {})) or FakeImage()
|
||||
@@ -689,7 +627,7 @@ class WorkerTest(unittest.TestCase):
|
||||
sys.modules,
|
||||
{
|
||||
"torch": fake_torch,
|
||||
"ben2_onnx": fake_ben2,
|
||||
"server.onnx_ben2": fake_ben2,
|
||||
"PIL": fake_pil,
|
||||
},
|
||||
):
|
||||
@@ -698,17 +636,18 @@ class WorkerTest(unittest.TestCase):
|
||||
input_directory.mkdir()
|
||||
(input_directory / "frame-0001.jpg").touch()
|
||||
(input_directory / "frame-0002.jpg").touch()
|
||||
args = SimpleNamespace(
|
||||
input=str(input_directory),
|
||||
model_dir="model",
|
||||
output_dir=str(root / "result"),
|
||||
status=str(root / "status.json"),
|
||||
device="cpu",
|
||||
video_frames=8,
|
||||
context = FakeJobContext(root / "result", "remove-background")
|
||||
self.background.remove_background(
|
||||
context,
|
||||
{
|
||||
"input": str(input_directory),
|
||||
"model": "BEN2",
|
||||
"device": "cpu",
|
||||
"video_frames": 8,
|
||||
},
|
||||
)
|
||||
self.worker.remove_background(args)
|
||||
|
||||
result_directory = Path(args.output_dir)
|
||||
result_directory = context.directory
|
||||
self.assertTrue(
|
||||
(result_directory / "foreground.0001.png").is_file()
|
||||
)
|
||||
@@ -726,49 +665,6 @@ class WorkerTest(unittest.TestCase):
|
||||
for _image, options in inference_options:
|
||||
self.assertEqual(options, {})
|
||||
|
||||
def test_remove_background_parser_accepts_video_frames(self):
|
||||
args = self.worker.parser().parse_args(
|
||||
[
|
||||
"remove-background",
|
||||
"--input",
|
||||
"input.mp4",
|
||||
"--model-dir",
|
||||
"ben2",
|
||||
"--output-dir",
|
||||
"result",
|
||||
"--status",
|
||||
"status.json",
|
||||
"--device",
|
||||
"auto",
|
||||
"--video-frames",
|
||||
"24",
|
||||
]
|
||||
)
|
||||
|
||||
self.assertEqual(args.command, "remove-background")
|
||||
self.assertEqual(args.video_frames, 24)
|
||||
|
||||
def test_extract_frames_parser_accepts_frame_range(self):
|
||||
args = self.worker.parser().parse_args(
|
||||
[
|
||||
"extract-frames",
|
||||
"--input",
|
||||
"frames",
|
||||
"--output-dir",
|
||||
"result",
|
||||
"--status",
|
||||
"status.json",
|
||||
"--video-frames",
|
||||
"24",
|
||||
"--start-frame",
|
||||
"70",
|
||||
]
|
||||
)
|
||||
|
||||
self.assertEqual(args.command, "extract-frames")
|
||||
self.assertEqual(args.video_frames, 24)
|
||||
self.assertEqual(args.start_frame, 70)
|
||||
|
||||
def test_extract_frames_copies_selected_frames(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
@@ -776,17 +672,17 @@ class WorkerTest(unittest.TestCase):
|
||||
input_directory.mkdir()
|
||||
for index in range(1, 6):
|
||||
(input_directory / f"frame-{index:04d}.jpg").touch()
|
||||
args = SimpleNamespace(
|
||||
input=str(input_directory),
|
||||
output_dir=str(root / "result"),
|
||||
status=str(root / "status.json"),
|
||||
video_frames=3,
|
||||
start_frame=2,
|
||||
context = FakeJobContext(root / "result", "extract-frames")
|
||||
self.media.extract_frames(
|
||||
context,
|
||||
{
|
||||
"input": str(input_directory),
|
||||
"video_frames": 3,
|
||||
"start_frame": 2,
|
||||
},
|
||||
)
|
||||
|
||||
self.worker.extract_frames(args)
|
||||
|
||||
result_directory = Path(args.output_dir)
|
||||
result_directory = context.directory
|
||||
self.assertEqual(
|
||||
sorted(
|
||||
path.name
|
||||
@@ -795,11 +691,8 @@ class WorkerTest(unittest.TestCase):
|
||||
),
|
||||
["frame.0001.jpg", "frame.0002.jpg", "frame.0003.jpg"],
|
||||
)
|
||||
status = json.loads(
|
||||
Path(args.status).read_text(encoding="utf-8")
|
||||
)
|
||||
self.assertEqual(status["progress"], 1.0)
|
||||
self.assertIn("3 frame", status["message"])
|
||||
self.assertEqual(context.status["progress"], 1.0)
|
||||
self.assertIn("3 frame", context.status["message"])
|
||||
|
||||
def test_uses_da3_export_parameters_for_formats_and_debug_mesh(self):
|
||||
inference_options = []
|
||||
@@ -812,7 +705,7 @@ class WorkerTest(unittest.TestCase):
|
||||
intrinsics=np.eye(3, dtype=np.float32)[None],
|
||||
processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8),
|
||||
)
|
||||
fake_da3 = ModuleType("da3_onnx")
|
||||
fake_da3 = ModuleType("server.onnx_da3")
|
||||
fake_da3.create_session = lambda _model_dir, _device: object()
|
||||
|
||||
def fake_infer(_session, images, **options):
|
||||
@@ -820,7 +713,7 @@ class WorkerTest(unittest.TestCase):
|
||||
return prediction
|
||||
|
||||
fake_da3.infer = fake_infer
|
||||
fake_runtime = ModuleType("onnx_runtime")
|
||||
fake_runtime = ModuleType("server.onnx_runtime")
|
||||
fake_runtime.provider_label = lambda _session: "CPU"
|
||||
|
||||
def fake_mesh_export(
|
||||
@@ -865,23 +758,23 @@ class WorkerTest(unittest.TestCase):
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"da3_onnx": fake_da3,
|
||||
"onnx_runtime": fake_runtime,
|
||||
"server.onnx_da3": fake_da3,
|
||||
"server.onnx_runtime": fake_runtime,
|
||||
},
|
||||
), patch.object(
|
||||
self.worker,
|
||||
self.da3,
|
||||
"export_textured_mesh_glb",
|
||||
side_effect=fake_mesh_export,
|
||||
), patch.object(
|
||||
self.worker,
|
||||
self.da3,
|
||||
"export_da3_point_cloud",
|
||||
side_effect=fake_point_cloud,
|
||||
), patch.object(
|
||||
self.worker,
|
||||
self.da3,
|
||||
"export_da3_depth_preview",
|
||||
side_effect=fake_depth_preview,
|
||||
), patch.object(
|
||||
self.worker,
|
||||
self.da3,
|
||||
"export_da3_mini_npz",
|
||||
side_effect=fake_mini_npz,
|
||||
):
|
||||
@@ -892,32 +785,31 @@ class WorkerTest(unittest.TestCase):
|
||||
"mesh",
|
||||
)
|
||||
for export_format in debug_formats:
|
||||
args = SimpleNamespace(
|
||||
input=str(input_path),
|
||||
video_frames=16,
|
||||
output_dir=str(root / export_format),
|
||||
export_format=export_format,
|
||||
status=str(root / f"{export_format}.json"),
|
||||
device="cpu",
|
||||
model_dir="model",
|
||||
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,
|
||||
show_cameras=0,
|
||||
gs_views_interval=2,
|
||||
mesh_filter_edges=True,
|
||||
mesh_max_edge_length=0.2,
|
||||
context = FakeJobContext(
|
||||
root / export_format,
|
||||
"debug-da3",
|
||||
)
|
||||
self.worker.reconstruct(args)
|
||||
status = json.loads(
|
||||
Path(args.status).read_text(encoding="utf-8")
|
||||
self.da3.generate_da3(
|
||||
context,
|
||||
{
|
||||
"input": str(input_path),
|
||||
"model": "DA3",
|
||||
"video_frames": 16,
|
||||
"export_format": export_format,
|
||||
"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,
|
||||
"mesh_filter_edges": True,
|
||||
"mesh_max_edge_length": 0.2,
|
||||
},
|
||||
)
|
||||
self.assertEqual(status["progress"], 1.0)
|
||||
self.assertIn(export_format, status["message"])
|
||||
result_directory = Path(args.output_dir)
|
||||
self.assertEqual(context.status["progress"], 1.0)
|
||||
self.assertIn(export_format, context.status["message"])
|
||||
result_directory = context.directory
|
||||
if export_format == "depth_vis":
|
||||
self.assertTrue((result_directory / "0000.jpg").is_file())
|
||||
self.assertFalse((result_directory / "depth_vis").exists())
|
||||
@@ -938,7 +830,7 @@ class WorkerTest(unittest.TestCase):
|
||||
|
||||
def test_reconstruct_copies_frames_for_multi_frame_input(self):
|
||||
inference_inputs = []
|
||||
fake_da3 = ModuleType("da3_onnx")
|
||||
fake_da3 = ModuleType("server.onnx_da3")
|
||||
fake_da3.create_session = lambda _model_dir, _device: object()
|
||||
|
||||
def fake_infer(_session, images, **_options):
|
||||
@@ -952,7 +844,7 @@ class WorkerTest(unittest.TestCase):
|
||||
)
|
||||
|
||||
fake_da3.infer = fake_infer
|
||||
fake_runtime = ModuleType("onnx_runtime")
|
||||
fake_runtime = ModuleType("server.onnx_runtime")
|
||||
fake_runtime.provider_label = lambda _session: "CPU"
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
@@ -964,32 +856,33 @@ class WorkerTest(unittest.TestCase):
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"da3_onnx": fake_da3,
|
||||
"onnx_runtime": fake_runtime,
|
||||
"server.onnx_da3": fake_da3,
|
||||
"server.onnx_runtime": fake_runtime,
|
||||
},
|
||||
):
|
||||
args = SimpleNamespace(
|
||||
input=str(input_directory),
|
||||
video_frames=16,
|
||||
output_dir=str(root / "result"),
|
||||
export_format="mini_npz",
|
||||
status=str(root / "status.json"),
|
||||
device="cpu",
|
||||
model_dir="model",
|
||||
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,
|
||||
show_cameras=0,
|
||||
gs_views_interval=1,
|
||||
remove_background_model_dir="",
|
||||
inference_mode="per_frame",
|
||||
context = FakeJobContext(
|
||||
root / "result",
|
||||
"generate-da3-depth",
|
||||
)
|
||||
self.da3.generate_da3(
|
||||
context,
|
||||
{
|
||||
"input": str(input_directory),
|
||||
"model": "DA3",
|
||||
"video_frames": 16,
|
||||
"export_format": "mini_npz",
|
||||
"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,
|
||||
"inference_mode": "per_frame",
|
||||
},
|
||||
)
|
||||
self.worker.reconstruct(args)
|
||||
|
||||
result_directory = Path(args.output_dir)
|
||||
result_directory = context.directory
|
||||
self.assertTrue(
|
||||
(result_directory / "frame.0001.jpg").is_file()
|
||||
)
|
||||
@@ -1006,14 +899,13 @@ class WorkerTest(unittest.TestCase):
|
||||
self.assertFalse(
|
||||
(result_directory / "exports").exists()
|
||||
)
|
||||
status = json.loads(Path(args.status).read_text(encoding="utf-8"))
|
||||
self.assertEqual(status["progress"], 1.0)
|
||||
self.assertEqual(context.status["progress"], 1.0)
|
||||
|
||||
def test_mesh_confidence_uses_percentile(self):
|
||||
depth = np.ones((1, 2, 2), dtype=np.float32)
|
||||
confidence = np.array([[[1.0, 2.0], [3.0, 4.0]]])
|
||||
|
||||
threshold = self.worker.confidence_percentile_threshold(
|
||||
threshold = self.geometry.confidence_percentile_threshold(
|
||||
depth,
|
||||
confidence,
|
||||
50.0,
|
||||
@@ -1027,12 +919,12 @@ class WorkerTest(unittest.TestCase):
|
||||
x, y = np.meshgrid([0.0, 1.0], [0.0, 1.0])
|
||||
points = np.stack((x, y, depth), axis=-1)
|
||||
|
||||
threshold = self.worker.confidence_percentile_threshold(
|
||||
threshold = self.geometry.confidence_percentile_threshold(
|
||||
depth,
|
||||
confidence,
|
||||
0.0,
|
||||
)
|
||||
faces = self.worker.mesh_triangle_faces(
|
||||
faces = self.geometry.mesh_triangle_faces(
|
||||
points,
|
||||
depth,
|
||||
confidence,
|
||||
@@ -1054,7 +946,7 @@ class WorkerTest(unittest.TestCase):
|
||||
)
|
||||
points = np.stack((x, y, depth), axis=-1)
|
||||
|
||||
faces = self.worker.mesh_triangle_faces(
|
||||
faces = self.geometry.mesh_triangle_faces(
|
||||
points,
|
||||
depth,
|
||||
confidence,
|
||||
@@ -1076,7 +968,7 @@ class WorkerTest(unittest.TestCase):
|
||||
foreground = np.ones((2, 3), dtype=bool)
|
||||
foreground[:, 0] = False
|
||||
|
||||
faces = self.worker.mesh_triangle_faces(
|
||||
faces = self.geometry.mesh_triangle_faces(
|
||||
points,
|
||||
depth,
|
||||
confidence,
|
||||
@@ -1096,12 +988,12 @@ class WorkerTest(unittest.TestCase):
|
||||
discontinuous = smooth.copy()
|
||||
discontinuous[:, 3:] += 1.0
|
||||
|
||||
self.assertFalse(self.worker.depth_edge_mask(smooth).any())
|
||||
self.assertTrue(self.worker.depth_edge_mask(discontinuous)[:, 2:4].all())
|
||||
self.assertFalse(self.geometry.depth_edge_mask(smooth).any())
|
||||
self.assertTrue(self.geometry.depth_edge_mask(discontinuous)[:, 2:4].all())
|
||||
|
||||
def test_higher_edge_culling_allows_less_stretch(self):
|
||||
limits = [
|
||||
self.worker.edge_culling_stretch_limit(value)
|
||||
self.geometry.edge_culling_stretch_limit(value)
|
||||
for value in (0.0, 50.0, 100.0)
|
||||
]
|
||||
|
||||
@@ -1119,17 +1011,17 @@ class WorkerTest(unittest.TestCase):
|
||||
stretched_points[0, 1] = (2.0, 0.0, 2.0)
|
||||
triangles = np.asarray([[0, 2, 3], [0, 3, 1]], dtype=np.int32)
|
||||
|
||||
flat_mask = self.worker.triangle_edge_stretch_mask(
|
||||
flat_mask = self.geometry.triangle_edge_stretch_mask(
|
||||
flat_points,
|
||||
triangles,
|
||||
1.0,
|
||||
)
|
||||
loose_mask = self.worker.triangle_edge_stretch_mask(
|
||||
loose_mask = self.geometry.triangle_edge_stretch_mask(
|
||||
stretched_points,
|
||||
triangles,
|
||||
2.0,
|
||||
)
|
||||
strict_mask = self.worker.triangle_edge_stretch_mask(
|
||||
strict_mask = self.geometry.triangle_edge_stretch_mask(
|
||||
stretched_points,
|
||||
triangles,
|
||||
1.0,
|
||||
@@ -1151,7 +1043,7 @@ class WorkerTest(unittest.TestCase):
|
||||
confidence = np.ones((2, 3), dtype=np.float32)
|
||||
confidence[:, 0] = 0.0
|
||||
|
||||
faces = self.worker.mesh_triangle_faces(
|
||||
faces = self.geometry.mesh_triangle_faces(
|
||||
points,
|
||||
depth,
|
||||
confidence,
|
||||
@@ -1178,12 +1070,12 @@ class WorkerTest(unittest.TestCase):
|
||||
points = np.stack((x, y, depth), axis=-1)
|
||||
face_counts = []
|
||||
for percentile in (0.0, 50.0, 100.0):
|
||||
threshold = self.worker.confidence_percentile_threshold(
|
||||
threshold = self.geometry.confidence_percentile_threshold(
|
||||
depth,
|
||||
confidence,
|
||||
percentile,
|
||||
)
|
||||
faces = self.worker.mesh_triangle_faces(
|
||||
faces = self.geometry.mesh_triangle_faces(
|
||||
points,
|
||||
depth,
|
||||
confidence,
|
||||
@@ -1199,7 +1091,7 @@ class WorkerTest(unittest.TestCase):
|
||||
vertices = np.arange(18, dtype=np.float32).reshape(6, 3)
|
||||
faces = np.asarray([[1, 3, 5]], dtype=np.int32)
|
||||
|
||||
compact_vertices, compact_faces = self.worker.compact_mesh(
|
||||
compact_vertices, compact_faces = self.geometry.compact_mesh(
|
||||
vertices,
|
||||
faces,
|
||||
)
|
||||
@@ -1220,7 +1112,7 @@ class WorkerTest(unittest.TestCase):
|
||||
dtype=bool,
|
||||
)
|
||||
|
||||
filtered = self.worker.remove_background_faces(faces, foreground)
|
||||
filtered = self.geometry.remove_background_faces(faces, foreground)
|
||||
|
||||
np.testing.assert_array_equal(filtered, faces[2:])
|
||||
|
||||
@@ -1258,14 +1150,14 @@ class WorkerTest(unittest.TestCase):
|
||||
sys.modules,
|
||||
{"PIL": fake_pil},
|
||||
), patch.object(
|
||||
self.worker,
|
||||
self.geometry,
|
||||
"compact_mesh",
|
||||
side_effect=lambda vertices, faces: (
|
||||
vertices,
|
||||
faces,
|
||||
),
|
||||
):
|
||||
result_path, texture_path = self.worker.export_image_mesh_data(
|
||||
result_path, texture_path = self.geometry.export_image_mesh_data(
|
||||
prediction,
|
||||
Path(directory) / "input.png",
|
||||
Path(directory) / "result",
|
||||
@@ -1313,7 +1205,7 @@ class WorkerTest(unittest.TestCase):
|
||||
)
|
||||
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
output_path = self.worker.export_prediction_attributes(
|
||||
output_path = self.geometry.export_prediction_attributes(
|
||||
prediction,
|
||||
directory,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user