Initial public release

tests/anyimage/server/runtime_support.py

@@ -0,0 +1,51 @@
+import importlib
+import tempfile
+import time
+import unittest
+from pathlib import Path
+
+
+class ServerTestCase(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.server_module = importlib.import_module("server")
+        importlib.import_module("server.geometry.prediction_artifacts")
+        importlib.import_module("server.jobs.cutout")
+        from blendjob import JobServer
+        from server.model_manager import ModelManager, ModelSessionCache
+        from server.models import onnx_runtime
+
+        cls.JobServer = JobServer
+
+        cls.ModelManager, cls.ModelSessionCache = ModelManager, ModelSessionCache
+        cls.onnx_runtime = onnx_runtime
+        cls.app_module = importlib.import_module("server.app")
+
+
+    def setUp(self):
+        self.jobs_temp = tempfile.TemporaryDirectory()
+        self.addCleanup(self.jobs_temp.cleanup)
+        self.storage_root = Path(self.jobs_temp.name)
+
+
+    def create_cache(self):
+        return self.ModelSessionCache()
+
+
+    def create_server(self):
+        source = self.server_module.server
+        server = self.JobServer("AnyImage Job Server", storage_root=self.storage_root)
+        self.addCleanup(server.close)
+        server.handlers.update(source.handlers)
+        server.resource_factories.update(source.resource_factories)
+        server.bind(self.storage_root)
+        return server
+
+
+    def wait_for_job(self, context, timeout=2.0):
+        deadline = time.monotonic() + timeout
+        while time.monotonic() < deadline:
+            if context._snapshot()["state"] in {"succeeded", "failed", "cancelled"}:
+                return context._snapshot()
+            time.sleep(0.01)
+        self.fail(f"Job did not finish: {context._snapshot()}")

tests/anyimage/server/support.py

@@ -0,0 +1,56 @@
+from tests.support.paths import PROJECT_ROOT
+import sys
+from pathlib import Path
+
+
+def load_server_module(name):
+    addon_directory = PROJECT_ROOT / "src" / "anyimage"
+    sys.path.insert(0, str(addon_directory))
+    try:
+        return __import__(f"server.{name}", fromlist=[name])
+    finally:
+        sys.path.remove(str(addon_directory))
+
+
+class FakeModels:
+    def __init__(self, directory="BEN2-ONNX"):
+        self.model_directory = Path(directory)
+
+    def directory(self, _model):
+        return self.model_directory
+
+    def get_background(self, _model_directory, _device):
+        from server.models.onnx_ben2 import create_session
+
+        return create_session(self.model_directory, "cpu"), 0.0
+
+
+    def get_upscale(self, _model_directory, _device):
+        return object()
+
+
+class FakeJobContext:
+    def __init__(self, directory, job_type, models=None):
+        self.directory = Path(directory)
+        self.job_type = job_type
+        self.models = models or FakeModels()
+        self.status = {}
+        self.logs = []
+
+    def resource(self, name):
+        if name != "model_manager":
+            raise KeyError(name)
+        return self.models
+
+    def progress(self, progress, message, **details):
+        self.status = {
+            "progress": progress,
+            "message": message,
+            **details,
+        }
+
+    def log(self, message, *, level="INFO"):
+        self.logs.append((level, message))
+
+    def check_cancelled(self):
+        return None

tests/anyimage/server/test_input_validation.py

@@ -0,0 +1,74 @@
+
+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", ".JPG", ".webp", ".jpeg"])
+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,kind", [(".png", "image"), (".mp4", "image")])
+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_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
+
+
+
+
+@pytest.mark.parametrize("suffix", [".mp4", ".mts", ".gif", ".tiff", ""])
+def test_unsupported_formats_are_rejected(media, tmp_path, suffix):
+    path = tmp_path / f"input{suffix}"
+    path.touch()
+    with pytest.raises(RuntimeError, match="Unsupported input format"):
+        media.validate_input_path(path)
+
+
+def test_frame_directory_is_rejected(media, tmp_path):
+    Image.new("RGB", (2, 2)).save(tmp_path / "frame.png")
+    with pytest.raises(RuntimeError, match="Source image file does not exist"):
+        media.validate_input_path(tmp_path)
+
+
+def test_animated_png_is_rejected(media, tmp_path):
+    path = tmp_path / "animated.png"
+    Image.new("RGB", (2, 2), "red").save(
+        path, save_all=True, append_images=[Image.new("RGB", (2, 2), "blue")], duration=100,
+    )
+    with pytest.raises(RuntimeError, match="Only still images"):
+        with media.open_image(path):
+            pytest.fail("Animated images must not reach processing")
+
+
+def test_ai_tasks_support_cooperative_cancellation(media):
+    with pytest.raises(media.TaskCancelled):
+        media.check_cancelled(lambda: True)
+    media.check_cancelled(lambda: False)

tests/anyimage/server/test_job_server.py

@@ -0,0 +1,173 @@
+import threading
+import time
+from types import SimpleNamespace
+from unittest.mock import patch
+from tests.anyimage.server.runtime_support import ServerTestCase
+
+class JobServerTest(ServerTestCase):
+    def test_dedicated_server_runs_only_one_job_at_a_time(self):
+        server = self.JobServer("Test", storage_root=self.storage_root)
+        self.addCleanup(server.close)
+        first_started = threading.Event()
+        release_first = threading.Event()
+        order = []
+
+        @server.job("example")
+        def example(_context, parameters):
+            order.append(parameters["name"])
+            if parameters["name"] == "first":
+                first_started.set()
+                release_first.wait(1.0)
+
+        first = server.submit("example", {"name": "first"}, job_id="job")
+        self.assertTrue(first_started.wait(1.0))
+        second = server.submit("example", {"name": "second"}, job_id="other")
+
+        self.assertIsNotNone(first)
+        self.assertIsNotNone(second)
+        self.assertTrue(server.snapshot("instance")["busy"])
+        self.assertEqual(server.snapshot("instance")["queued_jobs"], 1)
+        release_first.set()
+        self.wait_for_job(first)
+        self.wait_for_job(second)
+        self.assertEqual(order, ["first", "second"])
+
+
+    def test_queued_job_can_be_cancelled_before_it_runs(self):
+        server = self.JobServer("Test", storage_root=self.storage_root)
+        self.addCleanup(server.close)
+        first_started = threading.Event()
+        release_first = threading.Event()
+        executed = []
+
+        @server.job("example")
+        def example(_context, parameters):
+            executed.append(parameters["name"])
+            if parameters["name"] == "first":
+                first_started.set()
+                release_first.wait(1.0)
+
+        first = server.submit("example", {"name": "first"})
+        self.assertTrue(first_started.wait(1.0))
+        second = server.submit("example", {"name": "second"})
+
+        self.assertIs(server.cancel(second.job_id), second)
+        release_first.set()
+        self.wait_for_job(first)
+        second_status = self.wait_for_job(second)
+
+        self.assertEqual(second_status["state"], "cancelled")
+        self.assertEqual(executed, ["first"])
+
+
+    def test_cancelling_job_keeps_server_and_model_cache_alive(self):
+        fake_models = SimpleNamespace(
+            snapshot=lambda: {"loaded": {"background": "BEN2_BASE", "geometry": ""}},
+            clear=lambda: None,
+            close=lambda: None,
+        )
+        with patch.object(self.app_module, "ModelManager", return_value=fake_models):
+            server = self.create_server()
+        started = threading.Event()
+
+        @server.job("cancellable")
+        def cancellable(context, _parameters):
+            started.set()
+            while not context._is_cancelled():
+                time.sleep(0.001)
+            context.check_cancelled()
+
+        context = server.submit("cancellable", {}, job_id="job")
+        self.assertTrue(started.wait(1.0))
+        self.assertIsNotNone(server.cancel("job"))
+        self.assertTrue(context._is_cancelled())
+        status = self.wait_for_job(context)
+
+        self.assertEqual(status["state"], "cancelled")
+        self.assertFalse(server.shutdown_event.is_set())
+        self.assertEqual(
+            server.snapshot("instance")["resources"]["model_manager"]["loaded"]["background"],
+            "BEN2_BASE",
+        )
+
+
+    def test_job_progress_is_kept_in_server_memory(self):
+        server = self.JobServer("Test", storage_root=self.storage_root)
+        self.addCleanup(server.close)
+
+        @server.job("example")
+        def example(context, _parameters):
+            context.progress(0.5, "Half way")
+
+        context = server.submit("example", {})
+        self.wait_for_job(context)
+
+        self.assertRegex(context.job_id, r"^[0-9a-f]{32}$")
+        self.assertEqual(context.directory.name, context.job_id)
+        self.assertEqual(context._snapshot()["state"], "succeeded")
+        self.assertEqual(context._snapshot()["progress"], 1.0)
+
+
+    def test_runtime_storage_root_owns_the_jobs_directory(self):
+        storage_root = self.storage_root
+        server = self.JobServer("Test")
+        self.addCleanup(server.close)
+        server.bind(storage_root)
+
+        @server.job("example")
+        def example(_context, _parameters):
+            return {"value": 42}
+
+        context = server.submit("example", {})
+        self.wait_for_job(context)
+
+        self.assertEqual(server.jobs_directory, storage_root / "jobs")
+        self.assertEqual(context.storage_root, storage_root.resolve())
+        self.assertEqual(context._snapshot()["result"], {"value": 42})
+
+
+    def test_job_server_owns_and_clears_shared_resources(self):
+        events = []
+        resource = SimpleNamespace(
+            snapshot=lambda: {"loaded": "model"},
+            clear=lambda: events.append("clear"),
+            close=lambda: events.append("close"),
+        )
+        server = self.JobServer("Test", storage_root=self.storage_root)
+        self.addCleanup(server.close)
+        server.add_resource("model_manager", resource)
+
+        @server.job("example")
+        def example(context, _parameters):
+            self.assertIs(context.resource("model_manager"), resource)
+            release.wait(1.0)
+
+        release = threading.Event()
+        context = server.submit("example", {}, job_id="job")
+        self.assertFalse(server.clear_resource("model_manager"))
+        release.set()
+        self.wait_for_job(context)
+        self.assertTrue(server.clear_resource("model_manager"))
+        server.close()
+
+        self.assertEqual(
+            server.snapshot("instance")["resources"]["model_manager"]["loaded"],
+            "model",
+        )
+        self.assertEqual(events, ["clear", "close"])
+
+
+    def test_server_registers_tasks_instead_of_model_names(self):
+        server = self.create_server()
+        self.assertEqual(
+            set(server.handlers),
+            {
+                "download-model",
+                "download-required-models",
+                "remove-background",
+                "upscale-image",
+                "generate-depth-plane-geometry",
+                "generate-panorama-geometry",
+                "generate-cutout-artifacts",
+            },
+        )

tests/anyimage/server/test_model_download.py

@@ -0,0 +1,170 @@
+from tests.support.paths import PROJECT_ROOT
+import os
+import sys
+import unittest
+from pathlib import Path
+import tempfile
+from types import SimpleNamespace
+from unittest.mock import patch
+
+
+def load_download_module():
+    addon_directory = PROJECT_ROOT / "src" / "anyimage"
+    sys.path.insert(0, str(addon_directory))
+    try:
+        return __import__(
+            "server.model_download",
+            fromlist=["model_download"],
+        )
+    finally:
+        sys.path.remove(str(addon_directory))
+
+
+class FakeTqdm:
+    def __init__(self, *args, **kwargs):
+        self.n = kwargs.get("initial", 0)
+        self.total = kwargs.get("total", 0)
+
+    def update(self, amount=1):
+        self.n += amount
+        return True
+
+    def refresh(self, *args, **kwargs):
+        return True
+
+
+class ModelDownloadTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.download = load_download_module()
+
+    def create_progress(self, description="Downloading bytes", **options):
+        reports = []
+        progress_class = self.download.create_download_progress_class(
+            lambda value, message: reports.append((value, message)),
+            "owner/model",
+            tqdm_type=options.pop("tqdm_type", FakeTqdm),
+            clock=options.pop("clock", lambda: 1.0),
+            **options,
+        )
+        progress = progress_class(
+            desc=description,
+            total=100,
+            initial=0,
+            unit="B",
+        )
+        return progress, reports
+
+    def test_https_connection_bounds_connect_and_idle_read_time(self):
+        connect_timeouts = []
+        read_timeouts = []
+
+        def connect(connection):
+            connect_timeouts.append(connection.timeout)
+            connection.sock = SimpleNamespace(
+                settimeout=read_timeouts.append,
+            )
+
+        connection = self.download._ShortConnectHTTPSConnection("example.com")
+        with patch.object(
+            self.download.http.client.HTTPSConnection,
+            "connect",
+            new=connect,
+        ):
+            connection.connect()
+
+        self.assertEqual(connect_timeouts, [10.0])
+        self.assertEqual(read_timeouts, [60.0])
+        self.assertIsNone(connection.timeout)
+
+    def test_reports_aggregated_download_bytes(self):
+        progress, reports = self.create_progress()
+
+        progress.update(25)
+
+        self.assertAlmostEqual(reports[-1][0], 0.275)
+        self.assertEqual(reports[-1][1], "Downloading owner/model")
+
+    def test_reports_hugging_face_aggregated_bar(self):
+        progress, reports = self.create_progress(
+            "Downloading (incomplete total...)"
+        )
+
+        progress.update(25)
+
+        self.assertAlmostEqual(reports[-1][0], 0.275)
+
+    def test_reports_download_speed(self):
+        class FakeTqdmWithRate(FakeTqdm):
+            @property
+            def format_dict(self):
+                return {"rate": 14_000_000}
+
+        progress, reports = self.create_progress(
+            tqdm_type=FakeTqdmWithRate,
+            label="MoGe-2",
+        )
+
+        progress.update(50)
+
+        self.assertEqual(
+            reports[-1][1],
+            "Downloading MoGe-2 13.4 MB/s",
+        )
+
+    def test_progress_callback_has_no_eta_or_state_protocol(self):
+        times = iter((0.0, 10.0))
+        progress, reports = self.create_progress(clock=lambda: next(times))
+
+        progress.update(50)
+
+        self.assertEqual(len(reports[-1]), 2)
+
+    def test_ignores_unrelated_progress_bars(self):
+        progress, reports = self.create_progress("Reconstructing")
+
+        progress.update(25)
+
+        self.assertEqual(reports, [])
+
+    def test_download_endpoints_uses_direct_endpoint(self):
+        with patch.dict(os.environ, {"HF_ENDPOINT": ""}):
+            endpoints = self.download.download_endpoints()
+
+        self.assertEqual(endpoints, ["https://huggingface.co"])
+
+    def test_download_endpoints_respects_environment_override(self):
+        with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example"}):
+            endpoints = self.download.download_endpoints()
+
+        self.assertEqual(endpoints, ["https://hf.example"])
+
+    def test_download_endpoints_strips_trailing_slash(self):
+        with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example/"}):
+            endpoints = self.download.download_endpoints()
+
+        self.assertEqual(endpoints, ["https://hf.example"])
+
+    def test_declared_download_skips_hf_when_no_public_onnx_exists(self):
+        args = SimpleNamespace(
+            repository="",
+            progress=lambda *_args: None,
+            cancel_check=lambda: None,
+        )
+        with tempfile.TemporaryDirectory() as directory, patch.object(
+            self.download, "download_from_hf"
+        ) as hugging_face, patch.object(
+            self.download, "download_from_r2"
+        ) as mirror, patch.object(
+            self.download, "_verify_model_complete"
+        ):
+            self.download._download_declared_or_fail(
+                args, Path(directory), "Model", 0.05, 0.90
+            )
+
+        hugging_face.assert_not_called()
+        mirror.assert_called_once()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/anyimage/server/test_model_download_jobs.py

@@ -0,0 +1,154 @@
+import os
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+from types import ModuleType, SimpleNamespace
+from unittest.mock import patch
+from tests.anyimage.server.support import load_server_module
+
+
+class ModelDownloadJobsTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.download = load_server_module("model_download")
+
+
+    def test_download_model_uses_actual_request_without_preflight(self):
+        calls = []
+        fake_hugging_face = ModuleType("huggingface_hub")
+        fake_hugging_face.snapshot_download = lambda **options: calls.append(
+            options
+        )
+        args = SimpleNamespace(
+            repository="Ruicheng/moge-2-vits-normal-onnx",
+            model_dir="model",
+            progress=lambda *_args, **_kwargs: None,
+            allow_pattern=[],
+        )
+
+        with tempfile.TemporaryDirectory() as directory, patch.dict(
+            sys.modules,
+            {"huggingface_hub": fake_hugging_face},
+        ), patch.object(
+            self.download,
+            "create_download_progress_class",
+            return_value=object,
+        ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
+            self.download,
+            "_verify_model_complete",
+            return_value=None,
+        ):
+            args.model_dir = str(Path(directory) / "models" / "MoGe2")
+            self.download.download_model(args)
+
+        self.assertEqual(len(calls), 1)
+
+
+    def test_download_model_refuses_success_when_files_missing(self):
+        fake_hugging_face = ModuleType("huggingface_hub")
+        fake_hugging_face.snapshot_download = lambda **options: None
+        args = SimpleNamespace(
+            repository="PramaLLC/BEN2",
+            model_dir="model",
+            progress=lambda *_args, **_kwargs: None,
+            allow_pattern=[],
+        )
+
+        with tempfile.TemporaryDirectory() as directory, patch.dict(
+            sys.modules,
+            {"huggingface_hub": fake_hugging_face},
+        ), patch.object(
+            self.download,
+            "create_download_progress_class",
+            return_value=object,
+        ), patch.dict(os.environ, {"HF_ENDPOINT": ""}):
+            args.model_dir = str(Path(directory) / "model")
+            with self.assertRaisesRegex(
+                RuntimeError,
+                "All download sources failed",
+            ):
+                self.download.download_model(args)
+
+
+    def test_download_model_reports_all_failed_sources(self):
+        args = SimpleNamespace(
+            repository="Ruicheng/moge-2-vits-normal-onnx",
+            model_dir="model",
+            progress=lambda *_args, **_kwargs: None,
+            allow_pattern=[],
+            r2_directory="moge-2-vits-normal-onnx",
+            r2_file=["config.json|100|abc"],
+        )
+
+        with tempfile.TemporaryDirectory() as directory, patch.object(
+            self.download,
+            "download_from_hf",
+            side_effect=OSError("direct down"),
+        ), patch.object(
+            self.download,
+            "download_from_r2",
+            side_effect=OSError("mirror down"),
+        ):
+            args.model_dir = str(Path(directory) / "model")
+            with self.assertRaisesRegex(RuntimeError, "models.omoolab.xyz"):
+                self.download.download_model(args)
+
+
+    def test_download_model_prefers_direct_hf_then_r2(self):
+        args = SimpleNamespace(
+            repository="Ruicheng/moge-2-vits-normal-onnx",
+            model_dir="model",
+            progress=lambda *_args, **_kwargs: None,
+            allow_pattern=["config.json"],
+            r2_directory="moge-2-vits-normal-onnx",
+            r2_file=["config.json|100|abc"],
+        )
+        order = []
+
+        def fake_download_from_hf(*_args, **_kwargs):
+            order.append("hf")
+            raise OSError("hf down")
+
+        def fake_download_from_r2(*_args, **_kwargs):
+            order.append("r2")
+
+        with tempfile.TemporaryDirectory() as directory, patch.object(
+            self.download,
+            "download_from_hf",
+            side_effect=fake_download_from_hf,
+        ), patch.object(
+            self.download,
+            "download_from_r2",
+            side_effect=fake_download_from_r2,
+        ), patch.object(
+            self.download,
+            "_verify_model_complete",
+            return_value=None,
+        ):
+            args.model_dir = str(Path(directory) / "models" / "MoGe2")
+            self.download.download_model(args)
+
+        self.assertEqual(order, ["hf", "r2"])
+
+
+    def test_download_from_r2_downloads_declared_files(self):
+        args = SimpleNamespace(
+            repository="Ruicheng/moge-2-vits-normal-onnx",
+            model_dir="model",
+            progress=lambda *_args, **_kwargs: None,
+            allow_pattern=["config.json"],
+            r2_directory="moge-2-vits-normal-onnx",
+            r2_file=["config.json|100|abc"],
+        )
+        downloaded = []
+
+        with tempfile.TemporaryDirectory() as directory, patch.object(
+            self.download,
+            "_download_file_from_url",
+            side_effect=lambda *args, **kwargs: downloaded.append(args[1].name),
+        ):
+            args.model_dir = str(Path(directory) / "model")
+            self.download.download_from_r2(args, args.model_dir)
+
+        self.assertEqual(downloaded, ["config.json"])