Add new tests and functionality for image processing and geometry handling

- Introduced color_image.py for preparing standalone color images for tests.
- Added hdr_image.py to create HDR textures for testing.
- Implemented test_image_depth_panorama.py to evaluate panorama topology, UVs, and depth discontinuities.
- Created test_shadeless.py to validate the opacity and HDR color endpoints of the Shadeless group.
- Developed moge3_export.py for exporting dynamic MoGe-3 backbone and ONNX graphs.
- Added material.py to inherit material slots onto generated geometry.
- Built image_depth_panorama.py to construct a quad panorama with depth separation and sampling.

tests/anyimage/server/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/anyimage/server/test_edge_depth.py

@@ -204,7 +204,7 @@ def test_artifacts_use_corrected_depth_and_original_normals(
         check_cancelled=lambda: None,
     )
     with (
-        patch.object(artifacts.moge2, "infer", return_value=[frame]),
+        patch.object(artifacts.moge, "infer", return_value=[frame]),
         patch.object(
             artifacts, "write_depth_texture",
             return_value=tmp_path / "depth.exr",
@@ -217,7 +217,7 @@ def test_artifacts_use_corrected_depth_and_original_normals(
         ) as normal,
         patch.object(artifacts, "extend_edge_depth", wraps=extend_edge_depth) as repair,
     ):
-        artifacts.generate_moge2_artifacts(
+        artifacts.generate_moge_artifacts(
             context, {}, path, generate_depth=generate_depth, normal_mode=normal_mode,
         )
     if normal_mode != "NONE":

tests/anyimage/server/test_geometry_artifacts.py

@@ -19,7 +19,7 @@ class GeometryArtifactsTest(ServerTestCase):
             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:
+        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]])
@@ -246,10 +246,10 @@ class GeometryArtifactsTest(ServerTestCase):
 
         texture = importlib.import_module("server.geometry.depth_texture")
         with (
-            patch.object(artifacts.moge2, "infer", return_value=[frame]) as infer,
-            patch.object(texture, "_write_float_exr", wraps=texture._write_float_exr) as write,
+            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_moge2_artifacts(
+            result = artifacts.generate_moge_artifacts(
                 context,
                 {"max_input_size": 2048},
                 input_path,
@@ -284,7 +284,7 @@ class GeometryArtifactsTest(ServerTestCase):
         normal = context.directory / "object-normal.png"
         with patch.object(
             job,
-            "generate_moge2_artifacts",
+            "generate_moge_artifacts",
             return_value={
                 "depth": depth_texture,
                 "depth_metadata": metadata,

tests/anyimage/server/test_input_validation.py

@@ -0,0 +1,93 @@
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+from PIL import Image
+
+from tests.anyimage.server.support import load_server_module
+
+
+@pytest.fixture
+def media():
+    return load_server_module("media.input")
+
+
+@pytest.mark.parametrize("suffix", [".PNG", ".MP4", ".webp", ".mts"])
+def test_supported_format_is_case_insensitive(media, tmp_path, suffix):
+    path = tmp_path / f"input{suffix}"
+    path.touch()
+    assert media.validate_input_path(path) == path
+
+
+@pytest.mark.parametrize("suffix", [".tiff", ""])
+def test_frontend_and_backend_share_format_reason(media, tmp_path, suffix):
+    path = tmp_path / f"input{suffix}"
+    path.touch()
+    with pytest.raises(RuntimeError) as preflight:
+        media.validate_input_path(path)
+    with pytest.raises(RuntimeError) as backend:
+        media.prepare_input(path, 16, tmp_path)
+    assert str(preflight.value) == str(backend.value)
+    assert f"Unsupported input format: {suffix or '(no extension)'}" in str(backend.value)
+    assert "Supported images: .jpeg, .jpg, .png, .webp." in str(backend.value)
+    assert "Supported videos:" in str(backend.value)
+
+
+@pytest.mark.parametrize("suffix,kind", [(".png", "image"), (".mp4", "video")])
+def test_missing_path_precedes_format_validation(media, tmp_path, suffix, kind):
+    with pytest.raises(RuntimeError, match=f"Source {kind} file does not exist"):
+        media.validate_input_path(tmp_path / f"missing{suffix}")
+
+
+def test_directory_filtering_and_empty_reason(media, tmp_path):
+    (tmp_path / "note.txt").touch()
+    (tmp_path / "frame.exr").touch()
+    with pytest.raises(RuntimeError, match="No supported image frames found"):
+        media.validate_input_path(tmp_path)
+    for name in ["b.PNG", "a.jpg"]:
+        (tmp_path / name).touch()
+    frames, temporary = media.prepare_input(tmp_path, 1, tmp_path, start_frame=2)
+    assert frames == [tmp_path / "b.PNG"]
+    assert temporary is None
+
+
+def test_image_subset_rejects_video(media, tmp_path):
+    path = tmp_path / "clip.mp4"
+    path.touch()
+    with pytest.raises(RuntimeError) as error:
+        media.validate_input_path(path, images_only=True)
+    assert "Supported images:" in str(error.value)
+    assert "Supported videos:" not in str(error.value)
+
+
+def test_image_decode_error_keeps_cause(media, tmp_path):
+    path = tmp_path / "broken.png"
+    path.write_bytes(b"broken")
+    with pytest.raises(RuntimeError, match="Unable to read image") as error:
+        with media.open_image(path):
+            pytest.fail("Broken image must not reach processing")
+    assert isinstance(error.value.__cause__, OSError)
+
+
+def test_processing_error_is_not_relabelled(media, tmp_path):
+    path = tmp_path / "valid.png"
+    Image.new("RGB", (2, 2)).save(path)
+    failure = OSError("inference failed")
+    with pytest.raises(OSError) as error:
+        with media.open_image(path):
+            raise failure
+    assert error.value is failure
+
+
+def test_video_decode_error_closes_container(media, tmp_path):
+    failure = ValueError("broken packet")
+    def frames():
+        raise failure
+        yield
+    container = SimpleNamespace(streams=SimpleNamespace(video=[object()]),
+                                decode=lambda _stream: frames(), close=Mock())
+    with pytest.raises(RuntimeError, match="Unable to read video") as error:
+        media.extract_video_frames("broken.mp4", tmp_path, 2,
+                                   SimpleNamespace(open=lambda _path: container))
+    assert error.value.__cause__ is failure
+    container.close.assert_called_once()

tests/anyimage/server/test_job_server.py

@@ -165,6 +165,7 @@ class JobServerTest(ServerTestCase):
                 "remove-background",
                 "upscale-image",
                 "generate-depth-plane-geometry",
+                "generate-panorama-geometry",
                 "generate-cutout-artifacts",
             },
         )

tests/anyimage/server/test_media_input.py

@@ -103,22 +103,6 @@ class MediaInputTest(unittest.TestCase):
         self.assertIsNone(temporary_directory)
 
 
-    def test_copy_color_images_numbers_a_still_image(self):
-        with tempfile.TemporaryDirectory() as directory:
-            root = Path(directory)
-            source = root / "source.png"
-            source.write_bytes(b"frame")
-            output = root / "result"
-            output.mkdir()
-
-            self.media.copy_color_images([source], output)
-
-            self.assertEqual(
-                (output / "image.0001.png").read_bytes(),
-                b"frame",
-            )
-
-
     def test_extract_video_frames_reads_frames_sequentially(self):
         with tempfile.TemporaryDirectory() as directory:
             root = Path(directory)
@@ -179,5 +163,5 @@ class MediaInputTest(unittest.TestCase):
             input_path = Path(directory) / "input.txt"
             input_path.touch()
 
-            with self.assertRaisesRegex(RuntimeError, "Unsupported input type"):
+            with self.assertRaisesRegex(RuntimeError, "Unsupported input format: .txt"):
                 self.media.prepare_input(input_path, 16, directory)

tests/anyimage/server/test_model_ben2_onnx.py

@@ -36,13 +36,13 @@ class Ben2OnnxTest(unittest.TestCase):
                 names, feeds
             ),
         ) as run:
-            result = onnx_ben2.infer(
+            result = onnx_ben2.infer_alpha(
                 session,
                 Image.new("RGB", (40, 20)),
                 release_memory=False,
             )
 
-        self.assertEqual(result.size, (40, 20))
+        self.assertEqual(result.shape, (20, 40))
         self.assertFalse(run.call_args.kwargs["release_memory"])
 
     def test_preprocess_matches_model_layout(self):
@@ -57,13 +57,18 @@ class Ben2OnnxTest(unittest.TestCase):
         session = FakeSession()
         image = Image.new("RGB", (40, 20), (10, 20, 30))
 
-        result = onnx_ben2.infer(session, image)
+        result = onnx_ben2.infer_alpha(session, image)
 
-        self.assertEqual(result.mode, "RGBA")
-        self.assertEqual(result.size, image.size)
+        self.assertEqual(result.dtype, np.float32)
+        self.assertEqual(result.shape, (20, 40))
         self.assertEqual(session.feeds["input"].shape, (1, 3, 1024, 1024))
-        alpha = np.asarray(result.getchannel("A"))
-        self.assertLess(int(alpha[:, 0].max()), int(alpha[:, -1].min()))
+        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__":

tests/anyimage/server/test_model_geometry.py

@@ -2,7 +2,7 @@ from unittest.mock import patch
 import numpy as np
 
 
-from anyimage.server.models import moge2 as moge2_model
+from anyimage.server.models import moge as moge2_model
 from anyimage.server.models.geometry import GeometryFrame
 
 

tests/anyimage/server/test_model_inference.py

@@ -12,7 +12,7 @@ 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.moge2")
+        cls.moge2_model = load_server_module("models.moge")
 
         cls.upscale_model = load_server_module("models.upscale")
 

tests/anyimage/server/test_model_manager.py

@@ -83,7 +83,7 @@ class ModelManagerTest(ServerTestCase):
         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.get_moge("new-moge", "cuda"), new_moge)
 
         self.assertIs(cache.ben2_model, ben2)
         self.assertIs(cache.upscale_model, upscale)
@@ -167,10 +167,10 @@ class ModelManagerTest(ServerTestCase):
                 "release_ben2",
             )
 
-            first_moge2 = cache.get_moge2("moge2", "cuda")
+            first_moge2 = cache.get_moge("moge2", "cuda")
             assert_reused(
                 first_moge2,
-                lambda: cache.get_moge2("moge2", "cuda"),
+                lambda: cache.get_moge("moge2", "cuda"),
                 "release_geometry",
             )
 

tests/anyimage/server/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/anyimage/server/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, "video_frames": 1, "start_frame": 1}
+    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 metadata["reference_distance"] == pytest.approx(3, rel=1e-5)
+
+
+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 / "image.0001.png").exists()
+    if kind == "transparent":
+        assert not calls

tests/anyimage/server/test_server_jobs.py

@@ -54,10 +54,9 @@ class ServerJobsTest(unittest.TestCase):
                         "input_frames",
                         return_value=nullcontext([input_path]),
                     ),
-                    patch.object(self.depth_plane_job, "copy_color_images"),
                     patch.object(
                         self.depth_plane_job,
-                        "generate_moge2_artifacts",
+                        "generate_moge_artifacts",
                         return_value={
                             "depth": depth_texture,
                             "depth_metadata": metadata,
@@ -85,100 +84,23 @@ class ServerJobsTest(unittest.TestCase):
 
 
     def test_ben2_background_removal_writes_fast_transparent_png(self):
-        inference_options = []
-        save_options = []
-
-        class FakeImage:
-            size = (1, 1)
-
-            def __enter__(self):
-                return self
-
-            def __exit__(self, *_args):
-                return None
-
-            def convert(self, mode):
-                self.mode = mode
-                return self
-
-            def getchannel(self, _name):
-                return self
-
-            def putalpha(self, _alpha):
-                return None
-
-            def paste(self, *_args):
-                return None
-
-            def save(self, path, **options):
-                save_options.append(options)
-                Path(path).touch()
-
-        class FakeModel:
-            @classmethod
-            def from_pretrained(cls, model_dir):
-                self.assertEqual(model_dir, "model")
-                return cls()
-
-            def to(model_self, device):
-                self.assertEqual(device.type, "cpu")
-                return model_self
-
-            def eval(self):
-                return None
-
-            def inference(self, image, **options):
-                inference_options.append((image, options))
-                return FakeImage()
-
-        fake_torch = ModuleType("torch")
-        fake_torch.cuda = SimpleNamespace(
-            is_available=lambda: False,
-            synchronize=lambda: None,
-        )
-        fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
-        fake_torch.inference_mode = nullcontext
-        fake_ben2 = ModuleType("server.models.onnx_ben2")
-        fake_ben2.create_session = lambda _model_dir, _device: object()
-        fake_ben2.infer = lambda _model, image, **options: (
-            inference_options.append((image, options)) or FakeImage()
-        )
-        fake_pil = ModuleType("PIL")
-        fake_pil.Image = SimpleNamespace(
-            open=lambda _path: FakeImage(),
-            new=lambda *_args, **_kwargs: FakeImage(),
-        )
-        fake_pil.ImageChops = SimpleNamespace(
-            multiply=lambda _source, _foreground: FakeImage()
-        )
-
-        with tempfile.TemporaryDirectory() as directory, patch.dict(
-            sys.modules,
-            {
-                "torch": fake_torch,
-                "server.models.onnx_ben2": fake_ben2,
-                "PIL": fake_pil,
-            },
-        ):
+        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"
-            input_path.touch()
+            Image.new("RGBA", (2, 1), (255, 64, 32, 128)).save(input_path)
             context = FakeJobContext(root / "result", "remove-background")
-            self.background.run(
-                context,
-                {"input": str(input_path), "model": "BEN2", "device": "cpu"},
-            )
-
-            self.assertTrue((root / "result" / "image.0001.png").is_file())
-
-        self.assertEqual(
-            inference_options[0][1],
-            {"release_memory": True},
-        )
-        self.assertEqual(
-            save_options[0],
-            {"format": "PNG", "compress_level": 1},
-        )
+            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,
+            ):
+                self.background.run(context, {"input": str(input_path), "model": "BEN2", "device": "cpu"})
+            with Image.open(root / "result" / "image.0001.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"])
 
 
@@ -190,7 +112,7 @@ class ServerJobsTest(unittest.TestCase):
             (((255, 0, 0, 0), (0, 255, 0, 128)),),
             dtype=np.uint8,
         )
-        model_result = Image.new("RGBA", (2, 1), (10, 20, 30, 128))
+        model_result = np.full((1, 2), 128 / 255, np.float32)
         received = []
 
         def infer(_model, image, **_options):
@@ -202,7 +124,7 @@ class ServerJobsTest(unittest.TestCase):
             Image.fromarray(source_pixels, mode="RGBA").save(input_path)
             with (
                 patch.object(onnx_ben2, "create_session", return_value=object()),
-                patch.object(onnx_ben2, "infer", side_effect=infer),
+                patch.object(onnx_ben2, "infer_alpha", side_effect=infer),
             ):
                 result, _load_ms, _inference_ms = self.ben2.infer_foreground(
                     input_path,
@@ -226,111 +148,48 @@ class ServerJobsTest(unittest.TestCase):
         self.assertEqual(tuple(result_pixels[0, 1]), (0, 255, 0, 64))
         self.assertEqual(
             tuple(np.asarray(raw_result.convert("RGBA"))[0, 0]),
-            (10, 20, 30, 128),
+            (0, 0, 0, 128),
         )
 
 
     def test_ben2_background_removal_processes_each_frame(self):
-        inference_options = []
-
-        class FakeImage:
-            size = (1, 1)
-
-            def __enter__(self):
-                return self
-
-            def __exit__(self, *_args):
-                return None
-
-            def convert(self, mode):
-                self.mode = mode
-                return self
-
-            def getchannel(self, _name):
-                return self
-
-            def putalpha(self, _alpha):
-                return None
-
-            def paste(self, *_args):
-                return None
-
-            def save(self, path, **options):
-                Path(path).touch()
-
-        class FakeModel:
-            @classmethod
-            def from_pretrained(cls, model_dir):
-                return cls()
-
-            def to(model_self, device):
-                return model_self
-
-            def eval(self):
-                return None
-
-            def inference(self, image, **options):
-                inference_options.append((image, options))
-                return FakeImage()
-
-        fake_torch = ModuleType("torch")
-        fake_torch.cuda = SimpleNamespace(
-            is_available=lambda: False,
-            synchronize=lambda: None,
-        )
-        fake_torch.device = lambda device_type: SimpleNamespace(type=device_type)
-        fake_torch.inference_mode = nullcontext
-        fake_ben2 = ModuleType("server.models.onnx_ben2")
-        fake_ben2.create_session = lambda _model_dir, _device: object()
-        fake_ben2.infer = lambda _model, image, **options: (
-            inference_options.append((image, options)) or FakeImage()
-        )
-        fake_pil = ModuleType("PIL")
-        fake_pil.Image = SimpleNamespace(
-            open=lambda _path: FakeImage(),
-            new=lambda *_args, **_kwargs: FakeImage(),
-        )
-        fake_pil.ImageChops = SimpleNamespace(
-            multiply=lambda _source, _foreground: FakeImage()
-        )
-
-        with tempfile.TemporaryDirectory() as directory, patch.dict(
-            sys.modules,
-            {
-                "torch": fake_torch,
-                "server.models.onnx_ben2": fake_ben2,
-                "PIL": fake_pil,
-            },
-        ):
+        from PIL import Image
+        onnx_ben2 = load_server_module("models.onnx_ben2")
+        with tempfile.TemporaryDirectory() as directory:
             root = Path(directory)
-            input_directory = root / "frames"
-            input_directory.mkdir()
-            (input_directory / "frame-0001.jpg").touch()
-            (input_directory / "frame-0002.jpg").touch()
+            frames = root / "frames"
+            frames.mkdir()
+            for index in (1, 2):
+                Image.new("RGB", (2, 1), (index * 50, 0, 0)).save(frames / f"frame.{index:04d}.png")
             context = FakeJobContext(root / "result", "remove-background")
-            self.background.run(
-                context,
-                {
-                    "input": str(input_directory),
-                    "model": "BEN2",
-                    "device": "cpu",
-                    "video_frames": 8,
-                },
-            )
+            with (
+                patch.object(onnx_ben2, "create_session", return_value=object()),
+                patch.object(onnx_ben2, "infer_alpha", return_value=np.ones((1, 2), np.float32)) as infer,
+            ):
+                self.background.run(context, {"input": str(frames), "model": "BEN2", "device": "cpu", "video_frames": 8})
+            self.assertEqual([call.kwargs["release_memory"] for call in infer.call_args_list], [False, True])
+            for index in (1, 2):
+                with Image.open(context.directory / f"image.{index:04d}.png") as image:
+                    self.assertEqual(image.getpixel((0, 0)), (index * 50, 0, 0, 255))
 
-            result_directory = context.directory
-            self.assertTrue(
-                (result_directory / "image.0001.png").is_file()
-            )
-            self.assertTrue(
-                (result_directory / "image.0002.png").is_file()
-            )
 
-        self.assertEqual(len(inference_options), 2)
-        self.assertEqual(
-            [options["release_memory"] for _image, options in inference_options],
-            [False, True],
-        )
+    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", "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_saves_importable_upscale_sequence(self):