Add comprehensive tests for ONNX models and server jobs

- Introduced tests for the Moge3 ONNX model, including input validation, inference, and caching mechanisms.
- Added tests for ONNX runtime configurations and session management.
- Implemented tests for image upscaling functionality using ONNX models, ensuring proper handling of alpha channels and memory management.
- Created tests for panorama generation, verifying depth and visibility metrics.
- Developed tests for prediction artifacts, ensuring correct depth and normal data handling.
- Added server job tests for background removal and depth plane geometry generation, validating input handling and output correctness.
- Included environment tests to ensure proper dependency management and CI configurations.

tests/ai/__init__.py

@@ -0,0 +1 @@
+"""AI inference tests run locally."""

tests/ai/test_background_alpha.py

@@ -0,0 +1,36 @@
+from unittest.mock import patch
+
+import numpy as np
+from PIL import Image
+import pytest
+
+from server.jobs import remove_background
+from server.models import onnx_ben2
+from tests.anyimage.server.support import FakeJobContext
+
+
+@pytest.mark.parametrize("failure", ["cancel", "prediction", "sequence", "output_kind"])
+def test_alpha_job_failure_does_not_write_result(tmp_path, failure):
+    source = tmp_path / "input.png"
+    Image.new("RGB", (2, 2)).save(source)
+    context = FakeJobContext(tmp_path / "job", "remove-background")
+    parameters = {"input": str(source), "model": "BEN2_BASE", "device": "cpu", "output_kind": "alpha"}
+    if failure == "sequence":
+        frames = tmp_path / "frames"
+        frames.mkdir()
+        for i in (1, 2):
+            Image.new("RGB", (2, 2)).save(frames / f"{i}.png")
+        parameters["input"] = str(frames)
+    if failure == "output_kind":
+        parameters["output_kind"] = "invalid"
+
+    def infer(*args, **kwargs):
+        if failure == "prediction":
+            raise RuntimeError("Prediction failed")
+        context.check_cancelled = lambda: (_ for _ in ()).throw(RuntimeError("Cancelled"))
+        return np.ones((2, 2), np.float32)
+
+    with patch.object(onnx_ben2, "create_session", return_value=object()), patch.object(onnx_ben2, "infer_alpha", side_effect=infer):
+        with pytest.raises((RuntimeError, ValueError)):
+            remove_background.run(context, parameters)
+    assert not (context.directory / "alpha.npy").exists()

tests/ai/test_background_models.py

@@ -0,0 +1,60 @@
+import hashlib
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from server import model_catalog
+from server.jobs import remove_background
+from server.model_manager import ModelManager
+from server.models import background
+from tests.anyimage.server.support import FakeJobContext
+
+
+@pytest.mark.parametrize("key", model_catalog.BACKGROUND_MODELS)
+@pytest.mark.parametrize("kind", ["foreground", "alpha"])
+def test_models_dispatch_keep_alpha(tmp_path, key, kind):
+    source = tmp_path / "source.png"
+    Image.new("RGBA", (7, 5), (60, 120, 180, 128)).save(source)
+    manager = ModelManager(tmp_path / "models")
+    context = FakeJobContext(tmp_path / "job", "remove-background", manager)
+    adapter = background.model_adapter(manager.directory(key))
+    parameters = {"input": str(source), "model": key, "device": "directml"}
+    if kind == "alpha":
+        parameters["output_kind"] = kind
+    with patch.object(adapter, "create_session", return_value=object()) as create, patch.object(
+        adapter, "infer_alpha", return_value=np.full((5, 7), .4, np.float32)
+    ) as infer:
+        remove_background.run(context, parameters)
+    create.assert_called_once_with(manager.directory(key), "cpu" if key == "BIREFNET_HR_MATTING" else "directml")
+    infer.assert_called_once()
+    if kind == "alpha":
+        np.testing.assert_allclose(np.load(context.directory / "alpha.npy"), .4)
+    else:
+        outputs = sorted(context.directory.glob("*.png"))
+        assert len(outputs) == infer.call_count
+        with Image.open(outputs[0]) as result:
+            assert result.size == (7, 5)
+            assert result.getpixel((0, 0)) == (60, 120, 180, 51)
+    manager.close()
+    assert manager.snapshot()["loaded"]["background"] == ""
+
+
+@pytest.mark.parametrize("key", ["BIREFNET_LITE", "BIREFNET_HR_MATTING", "HAT_GAN_X4_SHARPER"])
+def test_readiness_checks_model_and_license_corruption(tmp_path, key):
+    from dataclasses import replace
+
+    content = b"verified"
+    files = tuple((name, len(content), hashlib.sha256(content).hexdigest()) for name in ("model.onnx", "LICENSE.txt"))
+    spec = replace(model_catalog.get_downloadable_model(key), r2_files=files)
+    directory = tmp_path / spec.directory_name
+    directory.mkdir()
+    manager = ModelManager(tmp_path)
+    with patch.object(model_catalog, "get_downloadable_model", return_value=spec):
+        assert not manager.ready(key)
+        for name, _, _ in files:
+            (directory / name).write_bytes(content)
+        assert manager.ready(key)
+        (directory / "LICENSE.txt").write_bytes(b"damaged!")
+        assert not manager.ready(key)

tests/ai/test_geometry_artifacts.py

@@ -0,0 +1,255 @@
+import importlib.util
+import json
+import sys
+from types import SimpleNamespace
+from unittest.mock import patch
+from tests.anyimage.server.runtime_support import ServerTestCase
+
+class GeometryArtifactsTest(ServerTestCase):
+    def test_depth_alpha_multiplies_coverage_without_changing_geometry(self):
+        import numpy as np
+
+        depth = importlib.import_module("server.geometry.depth_texture")
+        GeometryFrame = importlib.import_module("server.models.geometry").GeometryFrame
+        validity = np.asarray([[0.75, 0, 1, 1]], dtype=np.float32)
+        alpha = np.asarray([[0.8, 1, 0, 1]], dtype=np.float32)
+        points = np.arange(12, dtype=np.float32).reshape(1, 4, 3)
+        frame = GeometryFrame(
+            depth=points[..., 2], points=points,
+            validity=validity.copy(), intrinsics=np.eye(3, dtype=np.float32),
+        )
+        output = self.storage_root / "depth.exr"
+        with patch.object(depth, "write_float_exr") as write:
+            depth.write_depth_texture(frame, output, alpha=alpha)
+            packed = write.call_args.args[0]
+            np.testing.assert_allclose(packed[..., 3], [[0.6, 0, 0, 1]])
+            np.testing.assert_array_equal(packed[..., :3], points)
+            np.testing.assert_array_equal(frame.validity, validity)
+            self.assertEqual(packed.dtype, np.float32)
+            with self.assertRaisesRegex(ValueError, "Alpha.*dimensions"):
+                depth.write_depth_texture(frame, output, alpha=np.ones((4, 1)))
+            self.assertEqual(write.call_count, 1)
+
+    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 / "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,
+        )
+        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_depth_texture(frame, texture_path, alpha=np.ones_like(validity))
+        depth.write_depth_metadata(frame, 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.assertEqual(set(metadata), {"image_size", "intrinsics"})
+
+
+    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_preserves_rgb_independently_of_alpha(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[..., 3] = 255
+        Image.fromarray(first, mode="RGBA").save(first_path)
+        Image.fromarray(second, mode="RGBA").save(second_path)
+
+        for max_size in (2048, 2):
+            with self.subTest(max_size=max_size):
+                outputs = []
+                for source in (first_path, second_path):
+                    output = artifacts._limited_model_input(
+                        source, source.with_name(f"model-{source.name}"), max_size,
+                    )
+                    with Image.open(output) as opened:
+                        self.assertEqual(opened.mode, "RGB")
+                        outputs.append(np.asarray(opened))
+                np.testing.assert_array_equal(*outputs)
+                self.assertEqual(outputs[0].shape[1], min(3, max_size))
+                if max_size == 2048:
+                    np.testing.assert_array_equal(outputs[0], first[..., :3])
+
+
+    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"
+        source = Image.new("RGBA", (4, 4), (255, 128, 64, 255))
+        source.paste((255, 128, 64, 0), (2, 0, 4, 2))
+        source.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),
+        )
+
+        texture = importlib.import_module("server.geometry.depth_texture")
+        with (
+            patch.object(artifacts.moge, "infer", return_value=frame) as infer,
+            patch.object(texture, "write_float_exr", wraps=texture.write_float_exr) as write,
+        ):
+            result = artifacts.generate_moge_artifacts(
+                context,
+                {"max_input_size": 2048},
+                input_path,
+                generate_depth=True,
+                normal_mode="OBJECT",
+            )
+
+        np.testing.assert_array_equal(write.call_args.args[0][..., 3], [[1, 0], [1, 1]])
+        np.testing.assert_array_equal(frame.validity, np.ones((2, 2)))
+        metadata = json.loads(result["depth_metadata"].read_text())
+        self.assertEqual(metadata["image_size"], [2, 2])
+        self.assertNotIn("reference_depth", metadata)
+        infer.assert_called_once()
+        self.assertEqual(
+            set(result),
+            {"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,
+        )
+        depth_texture = context.directory / "depth.exr"
+        metadata = context.directory / "depth.json"
+        normal = context.directory / "object-normal.png"
+        with patch.object(
+            job,
+            "generate_moge_artifacts",
+            return_value={
+                "depth": depth_texture,
+                "depth_metadata": metadata,
+                "object_normal": normal,
+            },
+        ) as generate:
+            result = job.run(
+                context,
+                {
+                    "input": str(input_path),
+                    "generate_depth": True,
+                    "normal_mode": "OBJECT",
+                },
+            )
+
+        self.assertEqual(
+            result,
+            {
+                "depth": "depth.exr",
+                "depth_metadata": "depth.json",
+                "object_normal": "object-normal.png",
+            },
+        )
+        generate.assert_called_once_with(
+            context,
+            {
+                "input": str(input_path),
+                "generate_depth": True,
+                "normal_mode": "OBJECT",
+            },
+            input_path,
+            generate_depth=True,
+            normal_mode="OBJECT",
+        )

tests/ai/test_model_ben2_onnx.py

@@ -0,0 +1,75 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+import numpy as np
+from PIL import Image
+
+
+from server.models import onnx_ben2
+
+
+class FakeSession:
+    def get_providers(self):
+        return ["CPUExecutionProvider"]
+
+    def get_inputs(self):
+        return [SimpleNamespace(name="input")]
+
+    def get_outputs(self):
+        return [SimpleNamespace(name="alpha")]
+
+    def run(self, output_names, feeds):
+        self.output_names = output_names
+        self.feeds = feeds
+        gradient = np.linspace(0.0, 1.0, 1024, dtype=np.float32)
+        return [np.tile(gradient, (1024, 1))[None, None]]
+
+
+class Ben2OnnxTest(unittest.TestCase):
+    def test_infer_forwards_memory_release_boundary(self):
+        session = FakeSession()
+
+        with patch.object(
+            onnx_ben2,
+            "run_session",
+            side_effect=lambda model, names, feeds, **_options: model.run(
+                names, feeds
+            ),
+        ) as run:
+            result = onnx_ben2.infer_alpha(
+                session,
+                Image.new("RGB", (40, 20)),
+                release_memory=False,
+            )
+
+        self.assertEqual(result.shape, (20, 40))
+        self.assertFalse(run.call_args.kwargs["release_memory"])
+
+    def test_preprocess_matches_model_layout(self):
+        image = Image.new("RGB", (40, 20), (64, 128, 192))
+        result = onnx_ben2.preprocess(image)
+
+        self.assertEqual(result.shape, (1, 3, 1024, 1024))
+        self.assertEqual(result.dtype, np.float32)
+        self.assertTrue(result.flags.c_contiguous)
+
+    def test_infer_restores_original_size_and_alpha(self):
+        session = FakeSession()
+        image = Image.new("RGB", (40, 20), (10, 20, 30))
+
+        result = onnx_ben2.infer_alpha(session, image)
+
+        self.assertEqual(result.dtype, np.float32)
+        self.assertEqual(result.shape, (20, 40))
+        self.assertEqual(session.feeds["input"].shape, (1, 3, 1024, 1024))
+        self.assertLess(result[:, 0].max(), result[:, -1].min())
+        self.assertTrue(np.any(np.abs(result * 255 - np.rint(result * 255)) > 0.01))
+
+    def test_infer_rejects_nonfinite_mask(self):
+        with patch.object(onnx_ben2, "run_session", return_value=[np.full((1, 1, 2, 2), np.nan)]):
+            with self.assertRaisesRegex(RuntimeError, "invalid alpha"):
+                onnx_ben2.infer_alpha(FakeSession(), Image.new("RGB", (2, 2)))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ai/test_model_birefnet_onnx.py

@@ -0,0 +1,25 @@
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from server.models import onnx_birefnet
+
+
+@pytest.mark.parametrize("size", [1024, 2048])
+def test_probability_alpha_keeps_transparency_and_source_dimensions(size):
+    session = SimpleNamespace(
+        get_inputs=lambda: [SimpleNamespace(name="image", shape=[1, 3, size, size])],
+        get_outputs=lambda: [SimpleNamespace(name="alpha")],
+    )
+    mask = np.full((1, 1, size, size), 0.37, dtype=np.float32)
+    with patch.object(onnx_birefnet, "run_session", return_value=[mask]) as run:
+        result = onnx_birefnet.infer_alpha(session, Image.new("RGB", (31, 17), (64, 128, 192)), release_memory=False)
+    assert result.shape == (17, 31)
+    assert result.dtype == np.float32
+    np.testing.assert_allclose(result, 0.37, atol=1e-6)
+    assert run.call_args.kwargs["release_memory"] is False
+    feed = run.call_args.args[2]["image"]
+    np.testing.assert_allclose(feed[0, :, 0, 0], (np.array([64, 128, 192]) / 255 - [.485, .456, .406]) / [.229, .224, .225], atol=1e-6)

tests/ai/test_model_geometry.py

@@ -0,0 +1,46 @@
+from unittest.mock import patch
+import numpy as np
+
+
+from anyimage.server.models import moge as moge2_model
+
+
+def test_moge_infer_returns_geometry_frame_with_model_validity():
+    mask = np.asarray(((1.0, 0.2), (0.7, 0.0)), dtype=np.float32)
+    prediction = {
+        "depth": np.ones((2, 2), dtype=np.float32),
+        "normal": np.zeros((2, 2, 3), dtype=np.float32),
+        "mask": mask,
+        "intrinsics": np.asarray(
+            ((1.0, 0.0, 0.5), (0.0, 1.0, 0.5), (0.0, 0.0, 1.0)),
+            dtype=np.float32,
+        ),
+        "points": np.ones((2, 2, 3), dtype=np.float32),
+    }
+
+    with patch.object(moge2_model, "infer_one", return_value=prediction):
+        frame = moge2_model.infer(None, {}, "image", include_points=True)
+
+    assert np.array_equal(frame.validity, mask.astype(np.float32))
+    assert np.array_equal(frame.normal, prediction["normal"])
+    assert np.array_equal(frame.points, prediction["points"])
+    assert np.array_equal(
+        frame.intrinsics,
+        np.asarray(((2.0, 0.0, 1.0), (0.0, 2.0, 1.0), (0.0, 0.0, 1.0))),
+    )
+
+
+def test_moge_frame_omits_points_when_not_requested():
+    prediction = {
+        "depth": np.ones((2, 2), dtype=np.float32),
+        "normal": np.zeros((2, 2, 3), dtype=np.float32),
+        "mask": np.ones((2, 2), dtype=bool),
+        "intrinsics": np.eye(3, dtype=np.float32),
+        "points": np.ones((2, 2, 3), dtype=np.float32),
+    }
+
+    with patch.object(moge2_model, "infer_one", return_value=prediction):
+        frame = moge2_model.infer(None, {}, "image", include_points=False)
+
+    assert frame.points is None
+    assert frame.validity.all()

tests/ai/test_model_inference.py

@@ -0,0 +1,105 @@
+import tempfile
+import unittest
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import patch
+import numpy as np
+import pytest
+from PIL import Image
+from tests.anyimage.server.support import load_server_module, FakeModels
+
+
+class ModelInferenceTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.moge2_model = load_server_module("models.moge")
+
+        cls.upscale_model = load_server_module("models.upscale")
+
+        cls.onnx_upscale = load_server_module("models.onnx_upscale")
+
+
+    def test_moge2_infers_one_geometry_frame(self):
+        prediction = {
+            "points": np.ones((2, 2, 3), dtype=np.float32),
+            "depth": np.ones((2, 2), dtype=np.float32),
+            "normal": np.ones((2, 2, 3), dtype=np.float32),
+            "mask": np.ones((2, 2), dtype=bool),
+            "intrinsics": np.eye(3, dtype=np.float32),
+        }
+        with patch.object(
+            self.moge2_model,
+            "infer_one",
+            return_value=prediction,
+        ) as infer_one:
+            frame = self.moge2_model.infer(
+                FakeModels(),
+                {"model": "MOGE2_VITS_NORMAL", "device": "cpu"},
+                Path("input.png"),
+                include_points=True,
+            )
+
+        self.assertEqual(frame.depth.shape, (2, 2))
+        infer_one.assert_called_once()
+
+
+    def test_upscale_resource_error_releases_session_and_is_readable(self):
+        from PIL import Image
+
+        class Models:
+            def __init__(self, directory):
+                self.model_directory = directory
+                self.released = False
+
+            def directory(self, _model):
+                return self.model_directory
+
+            def get_upscale(self, _directory, _device):
+                return object()
+
+            def release_upscale(self):
+                self.released = True
+
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "input.png"
+            Image.new("RGB", (8, 8)).save(path)
+            models = Models(Path(directory))
+            with patch.object(
+                self.onnx_upscale,
+                "infer",
+                side_effect=self.onnx_upscale.UpscaleResourceError(
+                    "provider memory exhausted"
+                ),
+            ):
+                with self.assertRaisesRegex(
+                    RuntimeError,
+                    "Reduce Maximum AI Input Size or select CPU",
+                ):
+                    self.upscale_model.infer_one(
+                        models,
+                        {
+                            "model": "REALESRGAN_GENERAL_WDN_X4V3",
+                            "device": "directml",
+                            "max_input_size": 2048,
+                        },
+                        path,
+                    )
+
+            self.assertTrue(models.released)
+
+
+upscale = load_server_module("models.upscale")
+
+@pytest.mark.parametrize("key", ["REALESRGAN_GENERAL_WDN_X4V3", "HAT_GAN_X4_SHARPER", "REALESRGAN_X4PLUS"])
+def test_upscale_models_resize_prediction_then_restore_alpha(tmp_path, key):
+    rng = np.random.default_rng(1)
+    source = Image.fromarray(rng.integers(0,256,(5,7,4),dtype=np.uint8))
+    path = tmp_path / "source.png"
+    source.save(path)
+    prediction = Image.fromarray(rng.integers(0,256,(20,28,3),dtype=np.uint8))
+    manager = SimpleNamespace(directory=lambda _: tmp_path, get_upscale=lambda *args: object())
+    with patch("server.models.onnx_upscale.infer", return_value=prediction) as infer:
+        output = upscale.infer_one(manager, {"model":key,"device":"cpu"}, path)
+    assert infer.call_args.kwargs["tile_border"] == (48 if key == "HAT_GAN_X4_SHARPER" else 16)
+    np.testing.assert_array_equal(output.convert("RGB"), prediction.resize((14,10), Image.Resampling.LANCZOS))
+    np.testing.assert_array_equal(output.getchannel("A"), source.getchannel("A").resize((14,10), Image.Resampling.LANCZOS))

tests/ai/test_model_manager.py

@@ -0,0 +1,208 @@
+import sys
+import tempfile
+from pathlib import Path
+from types import ModuleType, SimpleNamespace
+from unittest.mock import patch
+from tests.anyimage.server.runtime_support import ServerTestCase
+
+class ModelManagerTest(ServerTestCase):
+    def test_background_cache_switches_model_directory_and_device_and_closes(self):
+        from server.models import background
+
+        cache = self.create_cache()
+        loads = []
+
+        def create(directory, device):
+            result = object()
+            loads.append((str(directory), device, result))
+            return result
+
+        with patch.object(background, "model_adapter", return_value=SimpleNamespace(create_session=create)):
+            first, _ = cache.get_background("a/birefnet-lite", "cpu")
+            assert cache.get_background("a/birefnet-lite", "cpu") == (first, 0.0)
+            second, _ = cache.get_background("a/BEN2-ONNX", "cpu")
+            third, _ = cache.get_background("a/birefnet-hr-matting", "cpu")
+            fourth, _ = cache.get_background("b/birefnet-hr-matting", "cpu")
+            assert cache.get_background("b/birefnet-hr-matting", "directml") == (fourth, 0.0)
+            fifth, _ = cache.get_background("b/birefnet-lite", "directml")
+        assert len(loads) == 5
+        assert len({id(item) for item in (first, second, third, fourth, fifth)}) == 5
+        assert cache.snapshot()["background"] == "birefnet-lite"
+        assert loads[2][1] == loads[3][1] == "cpu"
+        cache.close()
+        assert cache.background_model is None
+        assert cache.snapshot()["background"] == ""
+
+    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": {
+                    "background": "",
+                    "geometry": "",
+                    "upscale": "",
+                }
+            },
+        )
+        self.assertIn("server.model_catalog", sys.modules)
+
+
+    def test_background_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_background("BEN2-ONNX", device)
+            second, cached_load_ms = cache.get_background("BEN2-ONNX", device)
+
+        self.assertIs(first, second)
+        self.assertEqual(loads, [("BEN2-ONNX", "cuda")])
+        self.assertEqual(cache.snapshot()["background"], "BEN2-ONNX")
+        self.assertEqual(cached_load_ms, 0.0)
+
+
+    def test_model_cache_keeps_three_categories(self):
+        cache = self.create_cache()
+        ben2 = object()
+        old_moge = object()
+        upscale = object()
+        cache.background_model = ben2
+        cache.background_key = ("ben2", "cuda")
+        cache.moge_model = old_moge
+        cache.moge_key = ("old-moge", "cuda")
+        cache.upscale_model = upscale
+        cache.upscale_key = ("upscale", "cuda")
+        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_moge("new-moge", "cuda"), new_moge)
+
+        self.assertIs(cache.background_model, ben2)
+        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(),
+            "moge2": Session(),
+            "upscale": Session(),
+        }
+        ben2 = ModuleType("server.models.onnx_ben2")
+        ben2.create_session = lambda _directory, _device: sessions["ben2"]
+        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_moge2": moge2,
+                "server.models.onnx_upscale": upscale,
+            },
+        ):
+            first_ben2, _load_ms = cache.get_background("BEN2-ONNX", "cuda")
+            assert_reused(
+                first_ben2,
+                lambda: cache.get_background("BEN2-ONNX", "cuda")[0],
+                "release_background",
+            )
+
+            first_moge2 = cache.get_moge("moge2", "cuda")
+            assert_reused(
+                first_moge2,
+                lambda: cache.get_moge("moge2", "cuda"),
+                "release_geometry",
+            )
+
+            first_upscale = cache.get_upscale("upscale", "cuda")
+            assert_reused(
+                first_upscale,
+                lambda: cache.get_upscale("upscale", "cuda"),
+                "release_upscale",
+            )

tests/ai/test_model_moge2_onnx.py

@@ -0,0 +1,208 @@
+from tests.support.paths import PROJECT_ROOT
+import importlib.util
+import sys
+import unittest
+from pathlib import Path
+from types import ModuleType
+from unittest.mock import patch
+import numpy as np
+
+
+def load_module():
+    anyimage = ModuleType("anyimage")
+    anyimage.__path__ = [str(PROJECT_ROOT / "src" / "anyimage")]
+    server = ModuleType("anyimage.server")
+    server.__path__ = [str(PROJECT_ROOT / "src" / "anyimage" / "server")]
+    models = ModuleType("anyimage.server.models")
+    models.__path__ = [
+        str(PROJECT_ROOT / "src" / "anyimage" / "server" / "models")
+    ]
+    specification = importlib.util.spec_from_file_location(
+        "anyimage.server.models.onnx_moge2",
+        PROJECT_ROOT
+        / "src"
+        / "anyimage"
+        / "server"
+        / "models"
+        / "onnx_moge2.py",
+    )
+    module = importlib.util.module_from_spec(specification)
+    with patch.dict(
+        sys.modules,
+        {
+            "anyimage": anyimage,
+            "anyimage.server": server,
+            "anyimage.server.models": models,
+        },
+    ):
+        specification.loader.exec_module(module)
+    return module
+
+
+class FakeInput:
+    def __init__(self, name):
+        self.name = name
+
+
+class FakeSession:
+    def __init__(self, outputs):
+        self.outputs = outputs
+        self.calls = []
+
+    def run(self, names, inputs):
+        self.calls.append((names, inputs))
+        return [self.outputs[name] for name in names]
+
+    def get_providers(self):
+        return ["CPUExecutionProvider"]
+
+    def get_outputs(self):
+        return [FakeInput(name) for name in self.outputs]
+
+
+class Moge2OnnxTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.moge = load_module()
+
+
+    def test_coreml_session_requires_static_input_partitions(self):
+        calls = []
+        fake_runtime = ModuleType("onnxruntime")
+        fake_runtime.get_available_providers = lambda: [
+            "CoreMLExecutionProvider",
+            "CPUExecutionProvider",
+        ]
+        fake_runtime.InferenceSession = lambda path, providers: calls.append(
+            (path, providers)
+        )
+
+        with patch.dict(sys.modules, {"onnxruntime": fake_runtime}):
+            self.moge.create_session("moge-model", "coreml")
+
+        self.assertEqual(
+            calls[0][0],
+            str(Path("moge-model") / "model.onnx"),
+        )
+        self.assertEqual(
+            calls[0][1],
+            [
+                (
+                    "CoreMLExecutionProvider",
+                    {"RequireStaticInputShapes": "1"},
+                ),
+                "CPUExecutionProvider",
+            ],
+        )
+
+
+    def test_infer_uses_dynamic_token_input_and_numpy_postprocess(self):
+        height, width = 8, 12
+        raw = {
+            "points": np.zeros((1, height, width, 3), dtype=np.float32),
+            "normal": np.zeros((1, height, width, 3), dtype=np.float32),
+            "mask": np.ones((1, height, width), dtype=np.float32),
+            "scale": np.ones((1,), dtype=np.float32),
+        }
+        raw["points"][..., 2] = 1.0
+        raw["normal"][..., 2] = 1.0
+        session = FakeSession(raw)
+        expected = {
+            "points": np.zeros((height, width, 3), dtype=np.float32),
+            "depth": np.ones((height, width), dtype=np.float32),
+            "normal": raw["normal"][0],
+            "mask": np.ones((height, width), dtype=bool),
+            "intrinsics": np.eye(3, dtype=np.float32),
+        }
+
+        with patch.object(self.moge, "postprocess", return_value=expected):
+            result = self.moge.infer(
+                session,
+                np.zeros((height, width, 3), dtype=np.float32),
+                resolution_level=9,
+            )
+
+        self.assertIs(result, expected)
+        names, inputs = session.calls[0]
+        self.assertEqual(names, ["points", "normal", "mask", "scale"])
+        self.assertEqual(inputs["image"].shape, (1, 3, height, width))
+        self.assertEqual(inputs["num_tokens"].dtype, np.int64)
+        self.assertEqual(inputs["num_tokens"], 3600)
+
+    def test_infer_forwards_memory_release_boundary(self):
+        height, width = 8, 12
+        raw = {
+            "points": np.zeros((1, height, width, 3), dtype=np.float32),
+            "normal": np.zeros((1, height, width, 3), dtype=np.float32),
+            "mask": np.ones((1, height, width), dtype=np.float32),
+            "scale": np.ones((1,), dtype=np.float32),
+        }
+        session = FakeSession(raw)
+
+        with (
+            patch.object(self.moge, "postprocess", return_value={}),
+            patch.object(
+                self.moge,
+                "run_session",
+                side_effect=lambda model, names, feeds, **_options: model.run(
+                    names, feeds
+                ),
+            ) as run,
+        ):
+            self.moge.infer(
+                session,
+                np.zeros((height, width, 3), dtype=np.float32),
+                0,
+                release_memory=False,
+            )
+
+        self.assertFalse(run.call_args.kwargs["release_memory"])
+
+    def test_postprocess_recovers_camera_projection_and_metric_scale(self):
+        height, width = 64, 80
+        focal = np.float32(0.9)
+        shift = np.float32(0.3)
+        metric_scale = np.float32(2.0)
+        uv = self.moge.normalized_view_plane_uv(width, height, np.float32)
+        depth = np.linspace(1.0, 2.4, height * width, dtype=np.float32).reshape(
+            height,
+            width,
+        )
+        raw_points = np.concatenate(
+            (uv * depth[..., None] / focal, (depth - shift)[..., None]),
+            axis=-1,
+        )
+        raw = {
+            "points": raw_points[None],
+            "normal": np.broadcast_to(
+                np.array((0.0, 0.0, 1.0), dtype=np.float32),
+                (1, height, width, 3),
+            ),
+            "mask": np.ones((1, height, width), dtype=np.float32),
+            "metric_scale": np.array((metric_scale,), dtype=np.float32),
+        }
+
+        raw["mask"][0, 0, :4] = [0.0, 0.2, 0.7, 1.0]
+        result = self.moge.postprocess(raw)
+
+        np.testing.assert_allclose(result["depth"], depth * metric_scale, rtol=1e-3)
+        np.testing.assert_allclose(result["points"][..., 2], result["depth"], rtol=1e-5)
+        np.testing.assert_array_equal(result["mask"], raw["mask"][0])
+
+    def test_postprocess_preserves_mask_for_nonpositive_depth(self):
+        mask = np.array([[[0.2, 0.7], [0.0, 1.0]]], dtype=np.float32)
+        raw = {
+            "points": np.zeros((1, 2, 2, 3), dtype=np.float32),
+            "normal": np.zeros((1, 2, 2, 3), dtype=np.float32),
+            "mask": mask,
+            "metric_scale": np.ones(1, dtype=np.float32),
+        }
+        with patch.object(self.moge, "recover_focal_shift", return_value=(1.0, -1.0)) as recover:
+            result = self.moge.postprocess(raw)
+        np.testing.assert_array_equal(recover.call_args.args[1], mask[0] > 0.5)
+        np.testing.assert_array_equal(result["mask"], mask[0])
+        self.assertTrue((result["depth"] < 0).all())
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ai/test_model_moge3_onnx.py

@@ -0,0 +1,139 @@
+from dataclasses import replace
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+
+from server.models import moge, onnx_moge3
+from server import model_catalog, model_manager
+
+
+def test_sparse_neighbors_and_pooling_follow_voxel_coordinates():
+    coord = np.zeros((1, 16, 32, 3), dtype=np.float32)
+    inputs = onnx_moge3.build_sparse_inputs(coord)
+    neighbors = inputs["neighbor_0"]
+    assert np.array_equal(neighbors[:, 13], np.arange(512))
+    assert neighbors[0, 10] == 512  # Missing pixel at x=-1.
+    assert neighbors[0, 16] == 1
+    for level in range(4):
+        children = inputs[f"pool_{level}"]
+        parent = inputs[f"up_{level}"]
+        for row, values in enumerate(children):
+            valid = values < len(parent)
+            assert np.all(parent[values[valid]] == row)
+    shifted = coord.copy()
+    shifted[..., 2] += 1
+    for key, values in inputs.items():
+        assert np.array_equal(values, onnx_moge3.build_sparse_inputs(shifted)[key])
+
+
+def test_inference_runs_three_refinements_and_preserves_image_dimensions():
+    calls = []
+    coord = np.zeros((1, 16, 16, 3), dtype=np.float32)
+    backbone = object()
+    refiner = object()
+
+    def run(session, names, inputs, *, release_memory):
+        calls.append((session, inputs.copy(), release_memory))
+        if session is backbone:
+            return [
+                coord,
+                np.zeros((1, 1026, 1, 1), np.float32),
+                np.ones((1, 3, 16, 16), np.float32),
+                np.ones((1, 1, 16, 16), np.float32),
+                np.zeros((1, 1), np.float32),
+            ]
+        assert "neighbor_4" in inputs
+        return [np.full((1, 16, 16), 0.1, np.float32)]
+
+    with (
+        patch.object(onnx_moge3, "run_session", side_effect=run),
+        patch.object(
+            onnx_moge3,
+            "postprocess",
+            side_effect=lambda value, **kwargs: (value, kwargs),
+        ),
+    ):
+        output, options = moge.infer_image(
+            onnx_moge3.Moge3Session(backbone, refiner),
+            np.zeros((7, 11, 3), np.float32),
+            9,
+            fov_x=90,
+        )
+    assert len(calls) == 4
+    assert calls[0][1]["num_tokens"] == 3600
+    assert [call[2] for call in calls[1:]] == [False, False, True]
+    assert output["points"].shape == (1, 7, 11, 3)
+    assert np.allclose(output["points"][..., 2], np.exp(0.3))
+    assert options["fov_x"] == 90
+
+
+def test_bilinear_resize_preserves_edges_and_pixel_centers():
+    image = np.array([[0, 2], [4, 6]], np.float32)
+    assert np.array_equal(onnx_moge3._resize(image, 2, 2), image)
+    assert onnx_moge3._resize(image, 1, 1)[0, 0] == 3
+
+
+def test_moge3_requires_all_verified_model_files(tmp_path):
+    model = replace(
+        model_catalog.MOGE3_VITL,
+        r2_files=(
+            (
+                "backbone.onnx",
+                1,
+                "ca978112ca1bbdcafac231b39a23dc4da786eff8147c4e72b9807785afee48bb",
+            ),
+            (
+                "refiner.onnx",
+                1,
+                "3e23e8160039594a33894f6564e1b1348bbd7a0088d42c4acb73eeaed59c009d",
+            ),
+        ),
+    )
+    manager = model_manager.ModelManager(tmp_path)
+    directory = manager.directory("MOGE3_VITL")
+    directory.mkdir()
+    with patch.object(
+        model_manager.model_catalog, "get_downloadable_model", return_value=model
+    ):
+        (directory / "backbone.onnx").write_bytes(b"a")
+        assert not manager.ready("MOGE3_VITL")
+        (directory / "refiner.onnx").write_bytes(b"b")
+        assert manager.ready("MOGE3_VITL")
+        (directory / "refiner.onnx").write_bytes(b"c")
+        assert not manager.ready("MOGE3_VITL")
+
+
+def test_cache_reuses_moge3_and_releases_it_when_switching():
+    cache = model_manager.ModelSessionCache()
+    session = onnx_moge3.Moge3Session(object(), object())
+    with patch.object(onnx_moge3, "create_session", return_value=session) as load:
+        assert cache.get_moge("moge-3-vitl-onnx", "cpu") is session
+        assert cache.get_moge("moge-3-vitl-onnx", "cpu") is session
+        load.assert_called_once()
+    from server.models import onnx_moge2
+
+    with patch.object(onnx_moge2, "create_session", return_value=object()):
+        assert cache.get_moge("moge-2-vits-normal-onnx", "cpu") is not session
+    cache.close()
+    assert cache.moge_model is None
+
+
+def test_moge3_catalog_uses_onnx_assets_and_geometry_family():
+    model = model_catalog.get_downloadable_model("MOGE3_VITL")
+    assert model.ready_patterns == ("backbone.onnx", "refiner.onnx")
+    assert model_catalog.model_record(model)["family"] == "moge"
+    assert model.huggingface_repository == ""
+
+
+def test_cache_recovers_after_moge3_session_load_failure():
+    cache = model_manager.ModelSessionCache()
+    session = onnx_moge3.Moge3Session(object(), object())
+    with patch.object(
+        onnx_moge3, "create_session", side_effect=[RuntimeError("load failed"), session]
+    ):
+        with pytest.raises(RuntimeError, match="load failed"):
+            cache.get_moge("moge-3-vitl-onnx", "cpu")
+        assert cache.moge_model is None
+        assert cache.moge_key is None
+        assert cache.get_moge("moge-3-vitl-onnx", "cpu") is session

tests/ai/test_model_onnx_runtime.py

@@ -0,0 +1,199 @@
+import sys
+import unittest
+from types import ModuleType, SimpleNamespace
+from unittest.mock import patch
+
+
+from server.models import onnx_runtime
+
+
+class FakeSession:
+    def __init__(self, providers):
+        self.providers = providers
+        self.calls = []
+
+    def get_providers(self):
+        return self.providers
+
+    def run(self, output_names, input_feed, run_options=None):
+        self.calls.append((output_names, input_feed, run_options))
+        return ["prediction"]
+
+
+class OnnxRuntimeTest(unittest.TestCase):
+    def test_directml_disables_parallel_execution_and_memory_pattern(self):
+        options = SimpleNamespace(enable_mem_pattern=True, execution_mode="parallel")
+        runtime = ModuleType("onnxruntime")
+        runtime.get_available_providers = lambda: ["DmlExecutionProvider"]
+        runtime.SessionOptions = lambda: options
+        runtime.ExecutionMode = SimpleNamespace(ORT_SEQUENTIAL="sequential")
+        calls = []
+        runtime.InferenceSession = lambda path, **kwargs: calls.append(kwargs)
+        with patch.dict(sys.modules, {"onnxruntime": runtime}):
+            onnx_runtime.create_session("model.onnx", "directml")
+        self.assertFalse(options.enable_mem_pattern)
+        self.assertEqual(options.execution_mode, "sequential")
+        self.assertIs(calls[0]["sess_options"], options)
+
+    def create_runtime(self, available):
+        calls = []
+        runtime = ModuleType("onnxruntime")
+        runtime.get_available_providers = lambda: available
+        runtime.InferenceSession = lambda path, providers: calls.append(
+            (path, providers)
+        )
+        return runtime, calls
+
+    def test_cuda_session_uses_shrinkable_arena(self):
+        runtime, calls = self.create_runtime(
+            ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        )
+
+        with patch.dict(sys.modules, {"onnxruntime": runtime}):
+            onnx_runtime.create_session("model.onnx", "cuda")
+
+        self.assertEqual(
+            calls[0],
+            (
+                "model.onnx",
+                [
+                    (
+                        "CUDAExecutionProvider",
+                        {"arena_extend_strategy": "kSameAsRequested"},
+                    )
+                ],
+            ),
+        )
+
+    def test_coreml_and_cpu_use_plain_session_providers(self):
+        runtime, calls = self.create_runtime(
+            ["CoreMLExecutionProvider", "CPUExecutionProvider"]
+        )
+
+        with patch.dict(sys.modules, {"onnxruntime": runtime}):
+            onnx_runtime.create_session("model.onnx", "auto")
+
+        self.assertEqual(
+            calls[0][1],
+            [
+                "CoreMLExecutionProvider",
+                "CPUExecutionProvider",
+            ],
+        )
+
+    def test_model_options_only_configure_coreml_provider(self):
+        runtime, calls = self.create_runtime(
+            ["CoreMLExecutionProvider", "CPUExecutionProvider"]
+        )
+
+        with patch.dict(sys.modules, {"onnxruntime": runtime}):
+            onnx_runtime.create_session(
+                "model.onnx",
+                "coreml",
+                coreml_options={"RequireStaticInputShapes": "1"},
+            )
+
+        self.assertEqual(
+            calls[0][1],
+            [
+                (
+                    "CoreMLExecutionProvider",
+                    {"RequireStaticInputShapes": "1"},
+                ),
+                "CPUExecutionProvider",
+            ],
+        )
+
+    def test_coreml_device_selects_coreml_provider(self):
+        available = [
+            "CoreMLExecutionProvider",
+            "CPUExecutionProvider",
+        ]
+
+        self.assertEqual(
+            onnx_runtime.select_providers("coreml", available),
+            ["CoreMLExecutionProvider", "CPUExecutionProvider"],
+        )
+        self.assertEqual(
+            onnx_runtime.select_providers(
+                "coreml",
+                ["CPUExecutionProvider"],
+            ),
+            ["CPUExecutionProvider"],
+        )
+
+
+    def test_completing_cuda_run_shrinks_memory_arena(self):
+        configs = []
+
+        class FakeRunOptions:
+            def add_run_config_entry(self, key, value):
+                configs.append((key, value))
+
+        runtime = ModuleType("onnxruntime")
+        runtime.RunOptions = FakeRunOptions
+        session = FakeSession(["CUDAExecutionProvider", "CPUExecutionProvider"])
+
+        with patch.dict(sys.modules, {"onnxruntime": runtime}):
+            result = onnx_runtime.run_session(session, ["output"], {"input": 1})
+
+        self.assertEqual(result, ["prediction"])
+        self.assertIsInstance(session.calls[0][2], FakeRunOptions)
+        self.assertEqual(
+            configs,
+            [("memory.enable_memory_arena_shrinkage", "gpu:0")],
+        )
+
+    def test_intermediate_and_non_cuda_runs_use_plain_session_run(self):
+        cuda = FakeSession(["CUDAExecutionProvider"])
+        directml = FakeSession(["DmlExecutionProvider"])
+
+        onnx_runtime.run_session(
+            cuda,
+            ["output"],
+            {"input": 1},
+            release_memory=False,
+        )
+        result = onnx_runtime.run_session(
+            directml,
+            ["output"],
+            {"input": 2},
+        )
+
+        self.assertIsNone(cuda.calls[0][2])
+        self.assertIsNone(directml.calls[0][2])
+        self.assertEqual(result, ["prediction"])
+
+    def test_installed_runtime_accepts_memory_arena_shrinkage_option(self):
+        import onnxruntime
+
+        run_options = onnxruntime.RunOptions()
+        run_options.add_run_config_entry(
+            "memory.enable_memory_arena_shrinkage",
+            "gpu:0",
+        )
+
+
+    def test_auto_provider_prefers_accelerated_runtime(self):
+        providers = onnx_runtime.select_providers(
+            "auto",
+            ["CPUExecutionProvider", "DmlExecutionProvider"],
+        )
+
+        self.assertEqual(
+            providers,
+            ["DmlExecutionProvider", "CPUExecutionProvider"],
+        )
+
+
+    def test_explicit_unavailable_provider_falls_back_to_cpu(self):
+        providers = onnx_runtime.select_providers(
+            "cuda",
+            ["CPUExecutionProvider"],
+        )
+
+        self.assertEqual(providers, ["CPUExecutionProvider"])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ai/test_model_upscale_onnx.py

@@ -0,0 +1,164 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+import numpy as np
+from PIL import Image
+
+
+from server.models import onnx_upscale
+
+
+class FakeSession:
+    def get_providers(self):
+        return ["CPUExecutionProvider"]
+
+    def get_inputs(self):
+        return [SimpleNamespace(name="input")]
+
+    def get_outputs(self):
+        return [SimpleNamespace(name="output")]
+
+    def run(self, output_names, feeds):
+        self.output_names = output_names
+        self.feeds = feeds
+        source = feeds["input"]
+        return [np.repeat(np.repeat(source, 4, axis=2), 4, axis=3)]
+
+
+class UpscaleOnnxTest(unittest.TestCase):
+    def test_wider_context_preserves_pixels_across_odd_and_small_images(self):
+        for width, height in [(1, 1), (7, 5), (321, 163)]:
+            pixels = np.random.default_rng(7).integers(0, 256, (height, width, 3), dtype=np.uint8)
+            result = onnx_upscale.infer(FakeSession(), Image.fromarray(pixels), tile_border=48)
+            expected = np.repeat(np.repeat(pixels, 4, axis=0), 4, axis=1)
+            np.testing.assert_array_equal(np.asarray(result), expected)
+
+
+    def test_preprocess_matches_dynamic_rgb_layout(self):
+        image = Image.new("RGB", (7, 5), (64, 128, 192))
+
+        result = onnx_upscale.preprocess(image)
+
+        self.assertEqual(result.shape, (1, 3, 5, 7))
+        self.assertEqual(result.dtype, np.float32)
+        self.assertTrue(result.flags.c_contiguous)
+        np.testing.assert_allclose(
+            result[0, :, 0, 0],
+            np.array((64, 128, 192), dtype=np.float32) / 255.0,
+        )
+
+    def test_infer_returns_direct_four_x_rgb_prediction(self):
+        session = FakeSession()
+        image = Image.new("RGB", (7, 5), (10, 20, 30))
+
+        result = onnx_upscale.infer(session, image)
+
+        self.assertEqual(result.mode, "RGB")
+        self.assertEqual(result.size, (28, 20))
+        self.assertEqual(session.output_names, ["output"])
+        self.assertEqual(session.feeds["input"].shape, (1, 3, 256, 256))
+        self.assertEqual(result.getpixel((0, 0)), (10, 20, 30))
+
+    def test_infer_returns_raw_rgb_for_rgba_input(self):
+        session = FakeSession()
+        pixels = np.zeros((5, 7, 4), dtype=np.uint8)
+        pixels[..., :3] = (10, 20, 30)
+        pixels[..., 3] = np.arange(7, dtype=np.uint8) * 40
+        image = Image.fromarray(pixels, mode="RGBA")
+
+        result = onnx_upscale.infer(session, image)
+
+        self.assertEqual(result.mode, "RGB")
+        self.assertEqual(result.size, (28, 20))
+        self.assertEqual(result.getpixel((0, 0)), (10, 20, 30))
+
+    def test_tiles_use_reflected_context_at_outer_image_edges(self):
+        class RecordingSession(FakeSession):
+            def run(self, output_names, feeds):
+                self.input = feeds["input"].copy()
+                return super().run(output_names, feeds)
+
+        pixels = np.zeros((5, 7, 3), dtype=np.uint8)
+        pixels[..., 0] = np.arange(7, dtype=np.uint8) * 20
+        session = RecordingSession()
+
+        onnx_upscale.infer(
+            session,
+            Image.fromarray(pixels, mode="RGB"),
+        )
+
+        red = session.input[0, 0]
+        self.assertAlmostEqual(red[16, 22], 120 / 255.0)
+        self.assertAlmostEqual(red[16, 23], 100 / 255.0)
+
+
+    def test_all_models_split_large_inputs_into_fixed_tiles(self):
+        class CountingSession(FakeSession):
+            def __init__(self):
+                self.calls = 0
+
+            def run(self, output_names, feeds):
+                self.calls += 1
+                return super().run(output_names, feeds)
+
+        session = CountingSession()
+        result = onnx_upscale.infer(
+            session,
+            Image.new("RGB", (449, 225), (1, 2, 3)),
+        )
+
+        self.assertEqual(result.size, (1796, 900))
+        self.assertEqual(session.calls, 6)
+        self.assertEqual(session.feeds["input"].shape, (1, 3, 256, 256))
+
+    def test_only_final_tile_releases_memory(self):
+        releases = []
+        session = FakeSession()
+
+        def run(model, output_names, feeds, *, release_memory):
+            releases.append(release_memory)
+            return model.run(output_names, feeds)
+
+        with patch.object(onnx_upscale, "run_session", side_effect=run):
+            result = onnx_upscale.infer(
+                session,
+                Image.new("RGB", (449, 225), (1, 2, 3)),
+            )
+
+        self.assertEqual(result.size, (1796, 900))
+        self.assertEqual(releases, [False, False, False, False, False, True])
+
+    def test_non_final_image_does_not_release_its_final_tile(self):
+        releases = []
+        session = FakeSession()
+
+        def run(model, output_names, feeds, *, release_memory):
+            releases.append(release_memory)
+            return model.run(output_names, feeds)
+
+        with patch.object(onnx_upscale, "run_session", side_effect=run):
+            onnx_upscale.infer(
+                session,
+                Image.new("RGB", (449, 225), (1, 2, 3)),
+                release_memory=False,
+            )
+
+        self.assertEqual(releases, [False] * 6)
+
+    def test_provider_unicode_error_is_presented_as_resource_error(self):
+        class InvalidProviderSession(FakeSession):
+            def run(self, _output_names, _feeds):
+                raise UnicodeDecodeError("utf-8", b"\xc4", 0, 1, "invalid")
+
+        with self.assertRaisesRegex(
+            onnx_upscale.UpscaleResourceError,
+            "provider failed",
+        ):
+            onnx_upscale.infer(
+                InvalidProviderSession(),
+                Image.new("RGB", (7, 5)),
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/ai/test_panorama.py

@@ -0,0 +1,133 @@
+"""Analytic projection, production inference and artifact lifecycle checks."""
+
+import json
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+from PIL import Image
+
+from tests.anyimage.server.support import FakeJobContext, load_server_module
+
+geometry = load_server_module("geometry.panorama")
+onnx = load_server_module("models.onnx_moge2")
+job = load_server_module("jobs.panorama")
+
+
+def test_known_fov_recovers_shift_metric_and_excludes_invisible_points():
+    rays = geometry.view_rays(64)
+    depth = 2 + 0.3 * rays[..., 0]
+    points = rays * depth[..., None]
+    points[..., 2] -= 0.7
+    visible = np.ones((64, 64), bool)
+    visible[:8] = False
+    points[:8] = np.nan
+    raw = {"points": points[None], "normal": np.zeros_like(points)[None],
+           "mask": np.ones((1, 64, 64)), "metric_scale": [1.5]}
+    prediction = onnx.postprocess(raw, fov_x=90, source_valid=visible)
+    np.testing.assert_allclose(prediction["depth"][visible], depth[visible] * 1.5, rtol=1e-5)
+    assert not prediction["mask"][:8].any()
+    assert np.isfinite(prediction["points"]).all()
+
+
+def test_periodic_fusion_removes_view_scale_offsets():
+    bases = geometry.camera_bases()
+    distances = []
+    for scale, basis in zip(np.exp(np.linspace(-0.4, 0.4, 12)), bases):
+        rays = geometry.view_rays(128) @ basis
+        rays /= np.linalg.norm(rays, axis=-1, keepdims=True)
+        distances.append((3 + 0.4 * rays[..., 0] + 0.2 * rays[..., 2]) * scale)
+    merged, valid = geometry.merge_distances(distances, np.ones((12, 128, 128), bool), bases, 64, 32)
+    rays = geometry.panorama_directions(64, 32)
+    expected = 3 + 0.4 * rays[..., 0] + 0.2 * rays[..., 2]
+    merged *= np.median(expected / merged)
+    assert valid.all()
+    assert np.mean(abs(merged / expected - 1)) < 0.002
+    assert np.max(abs(merged[:, [0, -1]] / expected[:, [0, -1]] - 1)) < 0.005
+    sampled = geometry.sample_panorama(expected, np.array(((1, 1e-8, 0), (1, -1e-8, 0))))
+    np.testing.assert_allclose(sampled[0], sampled[1], atol=1e-6)
+
+
+@pytest.fixture
+def panorama_job(tmp_path, monkeypatch):
+    source = tmp_path / "source.png"
+    context = FakeJobContext(tmp_path / "result", "generate-panorama-geometry",
+                             SimpleNamespace(directory=lambda _: tmp_path, get_moge=lambda *_: object()))
+    calls, pixels = [], []
+
+    def infer(session, color, level, **options):
+        calls.append((color, level, options))
+        rays = geometry.view_rays(color.shape[0])
+        points = rays / np.linalg.norm(rays, axis=-1, keepdims=True) * 3
+        return {"points": points, "mask": np.ones(color.shape[:2])}
+
+    monkeypatch.setattr(job, "infer", infer)
+    write = job.write_float_exr
+    monkeypatch.setattr(job, "write_float_exr", lambda rgba, path: (pixels.append(rgba.copy()), write(rgba, path))[1])
+    parameters = {"input": str(source), "model": "MOGE2", "device": "cpu", "resolution_level": 3,
+                  "max_input_size": 32}
+    return source, context, parameters, calls, pixels
+
+
+def test_job_outputs_geometry_combines_alpha_once_and_skips_transparent_views(panorama_job):
+    source, context, parameters, calls, pixels = panorama_job
+    image = np.full((32, 64, 4), 255, np.uint8)
+    image[:24, :, 3] = 0
+    image[26:28, :, 3] = 128
+    Image.fromarray(image).save(source)
+    result = job.run(context, parameters)
+    assert 0 < len(calls) < 12
+    assert all(level == 3 and opts["fov_x"] == 90 for _, level, opts in calls)
+    assert all(np.all(color[~opts["source_valid"]] == 0) for color, _, opts in calls)
+    assert {path.name for path in context.directory.iterdir()} == {"depth.exr", "depth.json"}
+    alpha = pixels[0][..., 3]
+    assert pixels[0].shape == image.shape
+    assert not alpha[:24].any()
+    assert np.all(alpha <= image[..., 3] / 255 + 1e-7)
+    assert np.any(np.isclose(alpha, 128 / 255))
+    np.testing.assert_allclose(pixels[0][..., 0], pixels[0][..., 2])
+    metadata = json.loads((context.directory / result["depth_metadata"]).read_text())
+    assert metadata["projection"] == "equirectangular"
+    assert metadata["image_size"] == [64, 32]
+    assert set(metadata) == {"projection", "image_size"}
+
+
+def test_full_resolution_fusion_preserves_detail_and_periodic_alignment():
+    bases = geometry.camera_bases()
+    distances = []
+    def radius(rays):
+        angle = np.arctan2(rays[..., 1], rays[..., 0])
+        return 3 + 0.15 * np.sin(80 * angle) * (1 - rays[..., 2] ** 2) ** 4
+    for scale, basis in zip(np.exp(np.linspace(-0.2, 0.2, 12)), bases):
+        rays = geometry.view_rays(512) @ basis
+        rays /= np.linalg.norm(rays, axis=-1, keepdims=True)
+        distances.append(radius(rays) * scale)
+    merged, valid = geometry.merge_distances(distances, np.ones((12, 512, 512), bool),
+                                              bases, 1024, 512)
+    expected = radius(geometry.panorama_directions(1024, 512))
+    merged *= np.median(expected / merged)
+    assert merged.shape == (512, 1024) and valid.all()
+    assert np.sqrt(np.mean((merged - expected) ** 2)) < 0.001
+    assert np.max(abs(merged[:, [0, -1]] / expected[:, [0, -1]] - 1)) < 0.01
+
+
+@pytest.mark.parametrize("kind", ["transparent", "invalid_model", "wrong_ratio", "cancel"])
+def test_failed_jobs_leave_no_partial_artifacts(panorama_job, monkeypatch, kind):
+    source, context, parameters, calls, _ = panorama_job
+    Image.new("RGBA", (64 if kind != "wrong_ratio" else 32, 32),
+              (255, 255, 255, 0 if kind == "transparent" else 255)).save(source)
+    if kind == "invalid_model":
+        monkeypatch.setattr(job, "infer", lambda *a, **k: {"points": np.ones((32, 32, 3)), "mask": np.zeros((32, 32))})
+    if kind == "cancel":
+        count = [0]
+        def cancel():
+            count[0] += 1
+            if count[0] == 6:
+                raise RuntimeError("Cancelled")
+        context.check_cancelled = cancel
+    with pytest.raises((ValueError, RuntimeError)):
+        job.run(context, parameters)
+    assert not list(context.directory.glob("depth.*"))
+    assert not (context.directory / "foreground.png").exists()
+    if kind == "transparent":
+        assert not calls

tests/ai/test_prediction_artifacts.py

@@ -0,0 +1,104 @@
+from types import SimpleNamespace
+from unittest.mock import patch
+import numpy as np
+import pytest
+from PIL import Image
+
+
+from server.geometry import prediction_artifacts as artifacts
+from server.models.geometry import GeometryFrame
+from server.models.onnx_moge2 import depth_to_points
+
+
+def make_frame(depth):
+    depth = np.asarray(depth, dtype=np.float32)
+    height, width = depth.shape
+    normalized = np.array([[1.2, 0, 0.4], [0, 1.4, 0.6], [0, 0, 1]], np.float32)
+    intrinsics = normalized.copy()
+    intrinsics[0] *= width
+    intrinsics[1] *= height
+    return GeometryFrame(
+        depth=depth,
+        validity=np.ones_like(depth),
+        intrinsics=intrinsics,
+        points=depth_to_points(depth, normalized),
+        normal=np.zeros((*depth.shape, 3), np.float32),
+    )
+
+
+def silhouette():
+    alpha = np.zeros((25, 35), np.float32)
+    alpha[4:21, 4:22] = 1
+    alpha[4:21, 22:25] = 0.2
+    depth = np.full(alpha.shape, 20, np.float32)
+    depth[alpha == 1] = 2
+    return make_frame(depth), alpha
+
+
+def test_alpha_resize_retains_threshold_values(tmp_path):
+    rgba = np.zeros((1, 4, 4), np.uint8)
+    rgba[0, :, 3] = [25, 26, 242, 243]
+    path = tmp_path / "alpha.png"
+    Image.fromarray(rgba).save(path)
+    alpha = artifacts._image_alpha(path, (2, 8))
+    np.testing.assert_array_equal(
+        alpha >= 0.1, [[False] * 2 + [True] * 6] * 2
+    )
+    np.testing.assert_array_equal(
+        alpha >= 0.95, [[False] * 6 + [True] * 2] * 2
+    )
+
+
+@pytest.mark.parametrize(
+    "generate_depth,normal_mode",
+    [
+        (True, "OBJECT"),
+        (True, "TANGENT"),
+        (True, "NONE"),
+        (False, "OBJECT"),
+    ],
+)
+def test_artifacts_write_raw_requested_depth_and_preserve_normals(
+    tmp_path, generate_depth, normal_mode
+):
+    frame, alpha = silhouette()
+    rgba = np.full((*alpha.shape, 4), 255, np.uint8)
+    rgba[..., 3] = np.round(alpha * 255).astype(np.uint8)
+    path = tmp_path / "source.png"
+    Image.fromarray(rgba).save(path)
+    context = SimpleNamespace(
+        directory=tmp_path,
+        resource=lambda _: None,
+        progress=lambda *_: None,
+        check_cancelled=lambda: None,
+    )
+    with (
+        patch.object(artifacts.moge, "infer", return_value=frame) as infer,
+        patch.object(
+            artifacts, "write_depth_texture",
+            return_value=tmp_path / "depth.exr",
+        ) as depth,
+        patch.object(
+            artifacts, "write_depth_metadata", return_value=tmp_path / "depth.json"
+        ) as metadata,
+        patch.object(
+            artifacts, "_write_normal", return_value=tmp_path / "normal.png"
+        ) as normal,
+    ):
+        artifacts.generate_moge_artifacts(
+            context, {}, path, generate_depth=generate_depth, normal_mode=normal_mode,
+        )
+    infer.assert_called_once()
+    if normal_mode != "NONE":
+        assert normal.call_args.args[0] is frame
+    else:
+        normal.assert_not_called()
+    if generate_depth:
+        written_frame = depth.call_args.args[0]
+        np.testing.assert_array_equal(depth.call_args.kwargs["alpha"], alpha)
+        assert metadata.call_args.args[0] is written_frame
+        assert written_frame is frame
+        assert frame.depth[12, 24] == 20
+    else:
+        metadata.assert_not_called()
+        depth.assert_not_called()

tests/ai/test_server_jobs.py

@@ -0,0 +1,259 @@
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+import numpy as np
+from tests.anyimage.server.support import load_server_module, FakeJobContext
+
+
+class ServerJobsTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.resolution = load_server_module("media.resolution")
+
+        cls.background = load_server_module("jobs.remove_background")
+
+        cls.background_model = load_server_module("models.background")
+
+        cls.onnx_upscale = load_server_module("models.onnx_upscale")
+
+        cls.depth_plane_job = load_server_module("jobs.depth_plane")
+
+        cls.upscale_job = load_server_module("jobs.upscale")
+
+
+    def test_plane_job_requests_depth_and_matching_normal_space(self):
+        for plane_type, normal_space in (("DEPTH", "OBJECT"), ("RELIEF", "TANGENT")):
+            normal_key = f"{normal_space.lower()}_normal"
+            with tempfile.TemporaryDirectory() as directory:
+                root = Path(directory)
+                input_path = root / "input.png"
+                input_path.touch()
+                context = FakeJobContext(root / "result", "generate-depth-plane-geometry")
+                depth_texture = context.directory / "depth.exr"
+                metadata = context.directory / "depth.json"
+                normal_texture = context.directory / f"{normal_space.lower()}-normal.png"
+                parameters = {"input": str(input_path), "plane_type": plane_type}
+                with (
+                    patch.object(
+                        self.depth_plane_job,
+                        "validate_input_path",
+                        return_value=input_path,
+                    ),
+                    patch.object(
+                        self.depth_plane_job,
+                        "generate_moge_artifacts",
+                        return_value={
+                            "depth": depth_texture,
+                            "depth_metadata": metadata,
+                            normal_key: normal_texture,
+                        },
+                    ) as generate,
+                ):
+                    result = self.depth_plane_job.run(context, parameters)
+
+            self.assertEqual(
+                result,
+                {
+                    "depth": "depth.exr",
+                    "depth_metadata": "depth.json",
+                    normal_key: f"{normal_space.lower()}-normal.png",
+                },
+            )
+            generate.assert_called_once_with(
+                context,
+                parameters,
+                input_path,
+                generate_depth=True,
+                normal_mode=normal_space,
+            )
+
+
+    def test_ben2_background_removal_writes_fast_transparent_png(self):
+        from PIL import Image
+        onnx_ben2 = load_server_module("models.onnx_ben2")
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            input_path = root / "input.png"
+            Image.new("RGBA", (2, 1), (255, 64, 32, 128)).save(input_path)
+            context = FakeJobContext(root / "result", "remove-background")
+            with (
+                patch.object(onnx_ben2, "create_session", return_value=object()),
+                patch.object(onnx_ben2, "infer_alpha", return_value=np.full((1, 2), 128 / 255, np.float32)) as infer,
+                patch.object(Image.Image, "save", autospec=True, side_effect=Image.Image.save) as save,
+            ):
+                result = self.background.run(context, {"input": str(input_path), "model": "BEN2_BASE", "device": "cpu"})
+            self.assertEqual(result, {"foreground": "foreground.png"})
+            self.assertEqual([p.name for p in context.directory.glob("*.png")], ["foreground.png"])
+            with Image.open(root / "result" / "foreground.png") as image:
+                self.assertEqual(image.getpixel((0, 0)), (255, 64, 32, 64))
+            self.assertEqual(infer.call_args.kwargs, {"release_memory": True})
+            self.assertEqual(save.call_args.kwargs, {"format": "PNG", "compress_level": 1})
+        self.assertIn("Background removed in", context.status["message"])
+
+
+    def test_ben2_preserves_input_alpha_and_hides_transparent_rgb(self):
+        from PIL import Image
+
+        onnx_ben2 = load_server_module("models.onnx_ben2")
+        source_pixels = np.asarray(
+            (((255, 0, 0, 0), (0, 255, 0, 128)),),
+            dtype=np.uint8,
+        )
+        model_result = np.full((1, 2), 128 / 255, np.float32)
+        received = []
+
+        def infer(_model, image, **_options):
+            received.append(np.asarray(image.convert("RGB"), dtype=np.uint8))
+            return model_result
+
+        with tempfile.TemporaryDirectory() as directory:
+            input_path = Path(directory) / "input.png"
+            Image.fromarray(source_pixels, mode="RGBA").save(input_path)
+            with (
+                patch.object(onnx_ben2, "create_session", return_value=object()),
+                patch.object(onnx_ben2, "infer_alpha", side_effect=infer),
+            ):
+                result, _load_ms, _inference_ms = self.background_model.infer_foreground(
+                    input_path,
+                    "BEN2-ONNX",
+                    "cpu",
+                    preserve_input_alpha=True,
+                )
+                raw_result, _load_ms, _inference_ms = self.background_model.infer_foreground(
+                    input_path,
+                    "BEN2-ONNX",
+                    "cpu",
+                )
+
+        inference_pixels = received[0]
+        result_pixels = np.asarray(result.convert("RGBA"), dtype=np.uint8)
+        self.assertEqual(tuple(inference_pixels[0, 0]), (0, 0, 0))
+        self.assertEqual(tuple(inference_pixels[0, 1]), (0, 128, 0))
+        self.assertEqual(tuple(received[1][0, 0]), (0, 0, 0))
+        self.assertEqual(tuple(received[1][0, 1]), (0, 128, 0))
+        self.assertEqual(tuple(result_pixels[0, 0]), (255, 0, 0, 0))
+        self.assertEqual(tuple(result_pixels[0, 1]), (0, 255, 0, 64))
+        self.assertEqual(
+            tuple(np.asarray(raw_result.convert("RGBA"))[0, 0]),
+            (0, 0, 0, 128),
+        )
+
+
+    def test_background_alpha_job_preserves_float_precision(self):
+        from PIL import Image
+        onnx_ben2 = load_server_module("models.onnx_ben2")
+        alpha = np.array([[0.123456, 0.987654]], dtype=np.float32)
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            path = root / "input.png"
+            Image.new("RGB", (2, 1)).save(path)
+            context = FakeJobContext(root / "result", "remove-background")
+            with (
+                patch.object(onnx_ben2, "create_session", return_value=object()),
+                patch.object(onnx_ben2, "infer_alpha", return_value=alpha),
+            ):
+                result = self.background.run(context, {"input": str(path), "model": "BEN2_BASE", "device": "cpu", "output_kind": "alpha"})
+            self.assertEqual(result, {"alpha": "alpha.npy"})
+            np.testing.assert_array_equal(np.load(context.directory / result["alpha"], allow_pickle=False), alpha)
+            self.assertFalse(list(context.directory.glob("*.png")))
+
+
+    def test_upscale_job_returns_one_image_file(self):
+        from PIL import Image
+
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            input_path = root / "input.png"
+            input_path.touch()
+            context = FakeJobContext(root / "result", "upscale-image")
+            with patch.object(
+                self.upscale_job,
+                "infer_one",
+                return_value=Image.new("RGB", (8, 8), (10, 20, 30)),
+            ) as infer_one:
+                result = self.upscale_job.run(
+                    context,
+                    {
+                        "input": str(input_path),
+                        "model": "REALESRGAN_GENERAL_WDN_X4V3",
+
+                        "device": "cpu",
+                    },
+                )
+
+            self.assertEqual(infer_one.call_count, 1)
+            self.assertEqual(result, {"upscale": "upscale.png"})
+            colors = sorted(context.directory.glob("*.png"))
+            self.assertEqual(
+                [path.name for path in colors],
+                ["upscale.png"],
+            )
+            self.assertIn(
+                (
+                    "INFO",
+                    "Upscale model: REALESRGAN_GENERAL_WDN_X4V3; "
+                    "maximum AI input size: 2048 px",
+                ),
+                context.logs,
+            )
+
+
+    def test_upscale_job_preserves_source_alpha(self):
+        from PIL import Image
+
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            source_path = root / "source.png"
+            pixels = np.zeros((3, 5, 4), dtype=np.uint8)
+            pixels[..., :3] = (10, 20, 30)
+            pixels[..., 3] = np.arange(5, dtype=np.uint8) * 60
+            Image.fromarray(pixels, mode="RGBA").save(source_path)
+            context = FakeJobContext(root / "result", "upscale-image")
+            prediction = Image.new("RGB", (20, 12), (10, 20, 30))
+
+            with patch.object(
+                self.onnx_upscale,
+                "infer",
+                return_value=prediction,
+            ):
+                self.upscale_job.run(
+                    context,
+                    {
+                        "input": str(source_path),
+                        "model": "REALESRGAN_GENERAL_WDN_X4V3",
+
+                        "device": "cpu",
+                    },
+                )
+
+            with Image.open(context.directory / "upscale.png") as result:
+                self.assertEqual(result.mode, "RGBA")
+                self.assertEqual(result.size, (10, 6))
+                alpha = np.asarray(result.getchannel("A"))
+                expected_alpha = np.asarray(
+                    Image.fromarray(pixels[:, :, 3], mode="L").resize(
+                        result.size,
+                        Image.Resampling.LANCZOS,
+                    )
+                )
+                np.testing.assert_array_equal(alpha, expected_alpha)
+                self.assertLess(
+                    int(alpha[:, 0].max()),
+                    int(alpha[:, -1].min()),
+                )
+
+
+    def test_upscale_input_uses_the_global_long_edge_limit(self):
+        from PIL import Image
+
+        source = Image.new("RGB", (3000, 1500))
+        limited = self.resolution.limit_image(source, 2048)
+
+        self.assertEqual(limited.size, (2048, 1024))
+        self.assertEqual(
+            self.resolution.limit_image(Image.new("RGB", (800, 600)), 2048).size,
+            (800, 600),
+        )
+        with self.assertRaisesRegex(ValueError, "between 1 and 8192"):
+            self.resolution.limit_image(source, 9000)