import importlib.util import json 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(): server_directory = PROJECT_ROOT / "src" / "anyimage" / "server" specification = importlib.util.spec_from_file_location( "server", server_directory / "__init__.py", submodule_search_locations=[str(server_directory)], ) module = importlib.util.module_from_spec(specification) sys.modules["server"] = module 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 ServerRuntimeTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.server_module = load_server_module() importlib.import_module("server.geometry.prediction_artifacts") importlib.import_module("server.jobs.cutout") from blendjob import JobServer from server.model_manager import ModelManager, ModelSessionCache from server.models import onnx_runtime cls.JobServer = JobServer cls.ModelManager, cls.ModelSessionCache = ModelManager, ModelSessionCache cls.onnx_runtime = onnx_runtime cls.server = cls.server_module.server cls.app_module = sys.modules["server.app"] def setUp(self): self.jobs_temp = tempfile.TemporaryDirectory() self.addCleanup(self.jobs_temp.cleanup) self.storage_root = Path(self.jobs_temp.name) def test_required_model_download_skips_ready_model(self): manager = self.ModelManager(Path(".")) downloads = [] context = SimpleNamespace(progress=lambda *_args: None) with ( patch.object( manager, "ready", side_effect=lambda key: key == "BEN2_BASE", ), patch.object( manager, "_download", side_effect=lambda _context, key, _progress: downloads.append(key), ), ): result = manager.download_missing( context, ("MOGE2_VITS_NORMAL", "BEN2_BASE"), ) self.assertEqual(downloads, ["MOGE2_VITS_NORMAL"]) self.assertEqual(result, {"downloaded": ["MOGE2_VITS_NORMAL"]}) def test_models_are_owned_by_the_server_resource(self): with tempfile.TemporaryDirectory() as directory: manager = self.ModelManager(Path(directory)) snapshot = manager.snapshot() self.assertEqual( snapshot, { "loaded": { "ben2": "", "geometry": "", "upscale": "", } }, ) self.assertIn("server.model_catalog", sys.modules) self.assertNotIn("model_catalog", sys.modules) def create_cache(self): cache = self.ModelSessionCache.__new__(self.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 cache.upscale_model = None cache.upscale_key = None return cache def create_server(self): source = self.server_module.server server = self.JobServer("AnyImage 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.models.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.models.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.models.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.models.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()["geometry"], "model") def test_model_cache_keeps_three_categories(self): cache = self.create_cache() ben2 = object() old_moge = object() upscale = object() cache.ben2_model = ben2 cache.ben2_key = ("ben2", "cuda") cache.moge_model = old_moge cache.moge_key = ("old-moge", "cuda") cache.upscale_model = upscale cache.upscale_key = ("upscale", "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.assertIs(cache.ben2_model, ben2) self.assertIsNone(cache.moge_model) self.assertIs(cache.upscale_model, upscale) self.assertEqual(cache.snapshot()["geometry"], "da3") new_moge = object() fake_moge = ModuleType("server.models.onnx_moge2") fake_moge.create_session = lambda _model_dir, _device: new_moge with patch.dict(sys.modules, {"server.models.onnx_moge2": fake_moge}): self.assertIs(cache.get_moge2("new-moge", "cuda"), new_moge) self.assertIs(cache.ben2_model, ben2) self.assertIsNone(cache.da3_model) self.assertIs(cache.upscale_model, upscale) self.assertEqual(cache.snapshot()["geometry"], "new-moge") def test_upscale_model_is_reused_for_matching_runtime_key(self): loads = [] upscale = ModuleType("server.models.onnx_upscale") upscale.create_session = lambda model_dir, device: ( loads.append((model_dir, device)) or object() ) cache = self.create_cache() with patch.dict( sys.modules, {"server.models.onnx_upscale": upscale}, ): first = cache.get_upscale("model", "cpu") second = cache.get_upscale("model", "cpu") self.assertIs(first, second) self.assertEqual(loads, [("model", "cpu")]) self.assertEqual(cache.snapshot()["upscale"], "model") def test_inference_cleanup_keeps_every_cached_session(self): class Session: def get_providers(self): return ["CUDAExecutionProvider"] def run(self, _names, _feeds, run_options=None): self.run_options = run_options return [object()] class RunOptions: def add_run_config_entry(self, _key, _value): return None runtime = ModuleType("onnxruntime") runtime.RunOptions = RunOptions sessions = { "ben2": Session(), "da3": Session(), "moge2": Session(), "upscale": Session(), } ben2 = ModuleType("server.models.onnx_ben2") ben2.create_session = lambda _directory, _device: sessions["ben2"] da3 = ModuleType("server.models.onnx_da3") da3.create_session = lambda _directory, _device: sessions["da3"] moge2 = ModuleType("server.models.onnx_moge2") moge2.create_session = lambda _directory, _device: sessions["moge2"] upscale = ModuleType("server.models.onnx_upscale") upscale.create_session = lambda _directory, _device: sessions["upscale"] cache = self.create_cache() def assert_reused(first, getter, release_name): with ( patch.dict(sys.modules, {"onnxruntime": runtime}), patch.object( cache, release_name, wraps=getattr(cache, release_name), ) as release, ): self.onnx_runtime.run_session(first, ["output"], {"input": 1}) second = getter() self.assertIs(first, second) release.assert_not_called() with patch.dict( sys.modules, { "server.models.onnx_ben2": ben2, "server.models.onnx_da3": da3, "server.models.onnx_moge2": moge2, "server.models.onnx_upscale": upscale, }, ): first_ben2, _load_ms = cache.get_ben2("ben2", "cuda") assert_reused( first_ben2, lambda: cache.get_ben2("ben2", "cuda")[0], "release_ben2", ) first_da3 = cache.get_da3("da3", "cuda") assert_reused( first_da3, lambda: cache.get_da3("da3", "cuda"), "release_geometry", ) first_moge2 = cache.get_moge2("moge2", "cuda") assert_reused( first_moge2, lambda: cache.get_moge2("moge2", "cuda"), "release_geometry", ) first_upscale = cache.get_upscale("upscale", "cuda") assert_reused( first_upscale, lambda: cache.get_upscale("upscale", "cuda"), "release_upscale", ) def test_dedicated_server_runs_only_one_job_at_a_time(self): server = self.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", "download-required-models", "remove-background", "upscale-image", "generate-depth-plane-geometry", "generate-cutout-artifacts", "debug-da3", "debug-moge2", "debug-ben2", "debug-upscale", }, ) def test_queued_job_can_be_cancelled_before_it_runs(self): server = self.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.app_module, "ModelManager", return_value=fake_models): server = self.create_server() started = threading.Event() @server.job("cancellable") def cancellable(context, _parameters): started.set() while not context._is_cancelled(): time.sleep(0.001) context.check_cancelled() context = server.submit("cancellable", {}, job_id="job") self.assertTrue(started.wait(1.0)) 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.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.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.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_depth_artifacts_store_texture_and_compact_metadata(self): import numpy as np depth = importlib.import_module("server.geometry.depth_texture") GeometryFrame = importlib.import_module( "server.models.geometry" ).GeometryFrame texture_path = self.storage_root / "z-depth.exr" metadata_path = self.storage_root / "depth.json" validity = np.asarray( ((1.0, 1.0, 0.0), (1.0, 1.0, 1.0)), dtype=np.float32, ) calibration = np.asarray(((False, True, True), (False, False, True))) frame = GeometryFrame( depth=np.asarray( ((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)), dtype=np.float32, ), validity=validity, points=np.ones((2, 3, 3), dtype=np.float32), intrinsics=np.eye(3, dtype=np.float32), ) depth.write_z_depth_texture(frame, texture_path) depth.write_depth_metadata(frame, calibration, metadata_path) self.assertTrue(texture_path.is_file()) self.assertEqual(texture_path.read_bytes()[:4], b"v/1\x01") metadata = json.loads(metadata_path.read_text(encoding="utf-8")) self.assertEqual(metadata["image_size"], [3, 2]) self.assertEqual(metadata["intrinsics"], np.eye(3).tolist()) self.assertAlmostEqual(metadata["reference_depth"], 4.0) def test_object_normal_texture_matches_depth_surface_axes(self): import numpy as np from PIL import Image normal_texture = importlib.import_module("server.geometry.normal_texture") GeometryFrame = importlib.import_module( "server.models.geometry" ).GeometryFrame output_path = self.storage_root / "object-normal.png" camera_normals = np.asarray( [[ (0.0, 0.0, -1.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, -1.0), ]], dtype=np.float32, ) valid = np.asarray([[True, True, True, False]], dtype=bool) frame = GeometryFrame( depth=np.ones((1, 4), dtype=np.float32), validity=np.ones((1, 4), dtype=np.float32), intrinsics=np.eye(3, dtype=np.float32), normal=camera_normals, ) normal_texture.write_camera_normal_texture( frame, valid, (0, 0), (4, 1), (4, 1), output_path, normal_mode="OBJECT", ) with Image.open(output_path) as opened: encoded = np.asarray(opened, dtype=np.uint8) np.testing.assert_array_equal( encoded, np.asarray( [[ (127, 127, 255), (255, 127, 127), (127, 0, 127), (128, 128, 255), ]], dtype=np.uint8, ), ) def test_moge_model_input_composites_rgba_over_black(self): import numpy as np from PIL import Image artifacts = sys.modules["server.geometry.prediction_artifacts"] first_path = self.storage_root / "first.png" second_path = self.storage_root / "second.png" first = np.asarray( [[(240, 120, 60, 0), (200, 100, 50, 128), (9, 8, 7, 255)]], dtype=np.uint8, ) second = first.copy() second[0, 0, :3] = (1, 2, 3) Image.fromarray(first, mode="RGBA").save(first_path) Image.fromarray(second, mode="RGBA").save(second_path) first_output = artifacts._limited_model_input( first_path, self.storage_root / "first-model.png", 2048, ) second_output = artifacts._limited_model_input( second_path, self.storage_root / "second-model.png", 2048, ) with Image.open(first_output) as opened: first_rgb = np.asarray(opened, dtype=np.uint8) with Image.open(second_output) as opened: second_rgb = np.asarray(opened, dtype=np.uint8) np.testing.assert_array_equal(first_rgb, second_rgb) np.testing.assert_array_equal( first_rgb, np.asarray([[(0, 0, 0), (100, 50, 25), (9, 8, 7)]], dtype=np.uint8), ) def test_geometry_writers_reject_missing_or_non_finite_fields(self): import numpy as np depth = importlib.import_module("server.geometry.depth_texture") normal = importlib.import_module("server.geometry.normal_texture") GeometryFrame = importlib.import_module( "server.models.geometry" ).GeometryFrame base = { "depth": np.ones((1, 2), dtype=np.float32), "validity": np.ones((1, 2), dtype=np.float32), "intrinsics": np.eye(3, dtype=np.float32), } with self.assertRaisesRegex(ValueError, "point field"): depth.write_vector_depth_texture( GeometryFrame(**base), self.storage_root / "missing-points.exr", ) point_frame = GeometryFrame( **base, points=np.asarray([[(0.0, 0.0, 1.0), (np.nan, 0.0, 1.0)]], dtype=np.float32), ) with self.assertRaisesRegex(ValueError, "Point field"): depth.write_vector_depth_texture( point_frame, self.storage_root / "nan-points.exr", ) depth_frame = GeometryFrame( **{**base, "depth": np.asarray(((1.0, np.inf),), dtype=np.float32)} ) with self.assertRaisesRegex(ValueError, "Depth field"): depth.write_z_depth_texture( depth_frame, self.storage_root / "inf-depth.exr", ) with self.assertRaisesRegex(ValueError, "normal field"): normal.write_camera_normal_texture( GeometryFrame(**base), np.ones((1, 2), dtype=bool), (0, 0), (2, 1), (2, 1), self.storage_root / "missing-normal.png", ) normal_frame = GeometryFrame( **base, normal=np.asarray( [[(0.0, 0.0, 1.0), (0.0, np.nan, 1.0)]], dtype=np.float32, ), ) with self.assertRaisesRegex(ValueError, "Normal field"): normal.write_camera_normal_texture( normal_frame, np.ones((1, 2), dtype=bool), (0, 0), (2, 1), (2, 1), self.storage_root / "nan-normal.png", ) self.assertFalse(any(self.storage_root.glob("*.exr"))) self.assertFalse(any(self.storage_root.glob("*-normal.png"))) def test_moge_artifact_orchestration_uses_one_geometry_frame(self): import numpy as np from PIL import Image artifacts = sys.modules["server.geometry.prediction_artifacts"] GeometryFrame = importlib.import_module( "server.models.geometry" ).GeometryFrame input_path = self.storage_root / "selection.png" Image.new("RGBA", (2, 2), (255, 128, 64, 255)).save(input_path) output_directory = self.storage_root / "job" output_directory.mkdir() context = SimpleNamespace( directory=output_directory, resource=lambda _name: object(), progress=lambda *_args: None, check_cancelled=lambda: None, ) frame = GeometryFrame( depth=np.ones((2, 2), dtype=np.float32), validity=np.ones((2, 2), dtype=np.float32), intrinsics=np.eye(3, dtype=np.float32), normal=np.zeros((2, 2, 3), dtype=np.float32), points=np.ones((2, 2, 3), dtype=np.float32), ) with patch.object(artifacts.moge2, "infer", return_value=[frame]) as infer: result = artifacts.generate_moge2_artifacts( context, {"max_input_size": 2048}, input_path, depth_mode="VECTOR", normal_mode="OBJECT", ) infer.assert_called_once() self.assertEqual( set(result), {"vector_depth", "depth_metadata", "object_normal"}, ) self.assertFalse((output_directory / "model-input.png").exists()) self.assertTrue(all(path.is_file() for path in result.values())) self.assertTrue(infer.call_args.kwargs["include_points"]) def test_generate_cutout_artifacts_only_returns_requested_outputs(self): job = sys.modules["server.jobs.cutout"] input_path = self.storage_root / "input.png" input_path.write_bytes(b"selection") context = SimpleNamespace( directory=self.storage_root / "job", progress=lambda *_args: None, check_cancelled=lambda: None, ) vector_depth = context.directory / "vector-depth.exr" metadata = context.directory / "depth.json" normal = context.directory / "object-normal.png" with patch.object( job, "generate_moge2_artifacts", return_value={ "vector_depth": vector_depth, "depth_metadata": metadata, "object_normal": normal, }, ) as generate: result = job.run( context, { "input": str(input_path), "depth_mode": "VECTOR", "normal_mode": "OBJECT", }, ) self.assertEqual( result, { "vector_depth": "vector-depth.exr", "depth_metadata": "depth.json", "object_normal": "object-normal.png", }, ) generate.assert_called_once_with( context, { "input": str(input_path), "depth_mode": "VECTOR", "normal_mode": "OBJECT", }, input_path, depth_mode="VECTOR", normal_mode="OBJECT", correct_edge_depth=True, alpha_threshold=0.9, )