import os import sys import tempfile import unittest from contextlib import nullcontext from pathlib import Path from types import ModuleType, SimpleNamespace from unittest.mock import patch import numpy as np PROJECT_ROOT = Path(__file__).resolve().parents[1] 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)) def _make_fake_av(frame_count): """Build a fake av module yielding frame_count integer-valued frames.""" class FakeImage: def __init__(self, value): self.value = value def save(self, path, format=None): Path(path).write_text(str(self.value), encoding="utf-8") class FakeFrame: def __init__(self, value): self.value = value def to_image(self): return FakeImage(self.value) class FakeContainer: def __init__(self): self.position = 0 self.closed = False self.streams = SimpleNamespace(video=[object()]) def decode(self, _stream): while self.position < frame_count: frame = FakeFrame(self.position) self.position += 1 yield frame def close(self): self.closed = True class FakeAv: def __init__(self): self.containers = [] def open(self, _path): container = FakeContainer() self.containers.append(container) return container return FakeAv() class FakeModels: def __init__(self, directory="model"): self.model_directory = Path(directory) def directory(self, _model): return self.model_directory def get_ben2(self, _model_directory, _device): from server.models.onnx_ben2 import create_session return create_session(self.model_directory, "cpu"), 0.0 def get_da3(self, _model_directory, _device): return object() 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 class ServerJobsTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.media = load_server_module("media.input") cls.resolution = load_server_module("media.resolution") cls.geometry = load_server_module("geometry") cls.background = load_server_module("jobs.remove_background") cls.ben2 = load_server_module("models.ben2") cls.da3_model = load_server_module("models.da3") cls.moge2_model = load_server_module("models.moge2") cls.upscale_model = load_server_module("models.upscale") cls.onnx_upscale = load_server_module("models.onnx_upscale") cls.debug_da3_outputs = load_server_module("outputs.debug_da3") cls.debug_moge2_outputs = load_server_module("outputs.debug_moge2") cls.debug_jobs = load_server_module("jobs.debug_predictions") cls.upscale_job = load_server_module("jobs.upscale") cls.download = load_server_module("model_download") def test_upscale_progress_omits_single_frame_count(self): self.assertEqual( self.upscale_job.frame_progress_message("Upscaling image", 0, 1), "Upscaling image", ) self.assertEqual( self.upscale_job.frame_progress_message("Upscaling image", 1, 3), "Upscaling image (2/3)", ) def test_collect_images_is_sorted_and_filters_suffixes(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "b.PNG").touch() (root / "a.jpg").touch() (root / "notes.txt").touch() images = self.media.collect_images(root) self.assertEqual([path.name for path in images], ["a.jpg", "b.PNG"]) def test_ai_tasks_support_cooperative_cancellation(self): with self.assertRaises(self.media.TaskCancelled): self.media.check_cancelled(lambda: True) self.media.check_cancelled(lambda: False) def test_prepares_a_single_image_file(self): with tempfile.TemporaryDirectory() as directory: image_path = Path(directory) / "input.png" image_path.touch() images, temporary_directory = self.media.prepare_input( image_path, 16, directory, ) self.assertEqual(images, [image_path]) self.assertIsNone(temporary_directory) def test_prepares_a_frame_directory(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) (root / "frame-0001.jpg").touch() (root / "frame-0002.png").touch() (root / "notes.txt").touch() images, temporary_directory = self.media.prepare_input( root, 16, directory, ) self.assertEqual( [path.name for path in images], ["frame-0001.jpg", "frame-0002.png"], ) self.assertIsNone(temporary_directory) def test_prepares_first_frames_of_a_frame_directory(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) for index in range(1, 6): (root / f"frame-{index:04d}.jpg").touch() images, temporary_directory = self.media.prepare_input( root, 3, directory, ) self.assertEqual( [path.name for path in images], ["frame-0001.jpg", "frame-0002.jpg", "frame-0003.jpg"], ) self.assertIsNone(temporary_directory) def test_prepares_first_frames_after_offset_in_a_frame_directory(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) for index in range(1, 6): (root / f"frame-{index:04d}.jpg").touch() images, temporary_directory = self.media.prepare_input( root, 3, directory, start_frame=2, ) self.assertEqual( [path.name for path in images], ["frame-0002.jpg", "frame-0003.jpg", "frame-0004.jpg"], ) 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_copy_frames_skips_a_still_image_by_default(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) source = root / "source.png" source.touch() output = root / "result" output.mkdir() self.media.copy_frames([source], output) self.assertEqual(list(output.iterdir()), []) def test_extract_video_frames_reads_frames_sequentially(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) fake_av = _make_fake_av(100) image_paths = self.media.extract_video_frames( root / "clip.mkv", root / "out", 10, fake_av, ) self.assertEqual( [path.name for path in image_paths], [f"frame.{index:04d}.jpg" for index in range(1, 11)], ) self.assertEqual(len(fake_av.containers), 1) self.assertEqual(fake_av.containers[0].position, 10) def test_extract_video_frames_starts_after_offset(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) fake_av = _make_fake_av(100) image_paths = self.media.extract_video_frames( root / "clip.mkv", root / "out", 10, fake_av, start_frame=71, ) extracted = [ Path(path).read_text(encoding="utf-8") for path in image_paths ] self.assertEqual(extracted, [str(frame) for frame in range(70, 80)]) self.assertEqual(fake_av.containers[0].position, 80) def test_extract_video_frames_tolerates_short_video_tail(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) image_paths = self.media.extract_video_frames( root / "clip.mkv", root / "out", 10, _make_fake_av(99), ) self.assertEqual( [path.name for path in image_paths], [f"frame.{index:04d}.jpg" for index in range(1, 11)], ) def test_rejects_unsupported_input_file(self): with tempfile.TemporaryDirectory() as directory: input_path = Path(directory) / "input.txt" input_path.touch() with self.assertRaisesRegex(RuntimeError, "Unsupported input type"): self.media.prepare_input(input_path, 16, directory) 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="depth-anything/DA3-SMALL", 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" / "DA3") 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="depth-anything/DA3-SMALL", model_dir="model", progress=lambda *_args, **_kwargs: None, allow_pattern=[], r2_directory="DA3Small", 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="depth-anything/DA3-SMALL", model_dir="model", progress=lambda *_args, **_kwargs: None, allow_pattern=["config.json"], r2_directory="DA3Small", 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" / "DA3") self.download.download_model(args) self.assertEqual(order, ["hf", "r2"]) def test_download_from_r2_downloads_declared_files(self): args = SimpleNamespace( repository="depth-anything/DA3-SMALL", model_dir="model", progress=lambda *_args, **_kwargs: None, allow_pattern=["config.json"], r2_directory="DA3Small", 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"]) 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, }, ): root = Path(directory) input_path = root / "input.png" input_path.touch() 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}, ) 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 = Image.new("RGBA", (2, 1), (10, 20, 30, 128)) 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", side_effect=infer), ): result, _load_ms, _inference_ms = self.ben2.infer_foreground( input_path, "model", "cpu", preserve_input_alpha=True, ) raw_result, _load_ms, _inference_ms = self.ben2.infer_foreground( input_path, "model", "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]), (10, 20, 30, 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, }, ): root = Path(directory) input_directory = root / "frames" input_directory.mkdir() (input_directory / "frame-0001.jpg").touch() (input_directory / "frame-0002.jpg").touch() context = FakeJobContext(root / "result", "remove-background") self.background.run( context, { "input": str(input_directory), "model": "BEN2", "device": "cpu", "video_frames": 8, }, ) 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_da3_debug_saves_only_direct_api_prediction(self): prediction = SimpleNamespace( depth=np.full((1, 2, 2), 2.0, dtype=np.float32), conf=np.ones((1, 2, 2), dtype=np.float32), extrinsics=np.eye(4, dtype=np.float32)[None, :3], intrinsics=np.eye(3, dtype=np.float32)[None], processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8), is_metric=0, sky=None, scale_factor=None, ) fake_da3 = ModuleType("server.models.onnx_da3") calls = [] def fake_infer(_session, images, **options): calls.append((images, options)) return prediction fake_da3.infer = fake_infer fake_runtime = ModuleType("server.models.onnx_runtime") fake_runtime.provider_label = lambda _session: "CPU" with tempfile.TemporaryDirectory() as directory: root = Path(directory) input_path = root / "input.png" input_path.touch() with patch.dict( sys.modules, { "server.models.onnx_da3": fake_da3, "server.models.onnx_runtime": fake_runtime, }, ): context = FakeJobContext(root / "result", "debug-da3") self.debug_jobs.run_da3( context, { "input": str(input_path), "model": "DA3", "video_frames": 1, "device": "cpu", "process_res": 518, "process_res_method": "lower_bound_resize", "ref_view_strategy": "middle", }, ) with np.load(context.directory / "prediction.npz") as result: self.assertEqual( set(result.files), { "depth", "conf", "extrinsics", "intrinsics", "processed_images", "is_metric", }, ) np.testing.assert_array_equal(result["conf"], prediction.conf) self.assertFalse(any(context.directory.glob("*.glb"))) self.assertFalse(any(context.directory.glob("*.png"))) self.assertEqual(calls[0][0], [str(input_path)]) self.assertEqual(calls[0][1]["process_res"], 518) self.assertTrue(calls[0][1]["release_memory"]) def test_da3_per_frame_releases_only_after_final_frame(self): releases = [] def prediction(): return SimpleNamespace( depth=np.ones((1, 2, 2), dtype=np.float32), conf=np.ones((1, 2, 2), dtype=np.float32), extrinsics=np.eye(4, dtype=np.float32)[None, :3], intrinsics=np.eye(3, dtype=np.float32)[None], processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8), ) fake_da3 = ModuleType("server.models.onnx_da3") def infer(_session, _images, **options): releases.append(options["release_memory"]) return prediction() fake_da3.infer = infer fake_runtime = ModuleType("server.models.onnx_runtime") fake_runtime.provider_label = lambda _session: "CPU" parameters = { "model": "DA3", "device": "cpu", "process_res": 56, "process_res_method": "upper_bound_resize", "ref_view_strategy": "middle", "inference_mode": "per_frame", } with patch.dict( sys.modules, { "server.models.onnx_da3": fake_da3, "server.models.onnx_runtime": fake_runtime, }, ): self.da3_model.infer( FakeModels(), parameters, [Path("a.png"), Path("b.png")], ) self.assertEqual(releases, [False, True]) def test_moge2_debug_saves_each_direct_api_dictionary(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 tempfile.TemporaryDirectory() as directory: root = Path(directory) frames = root / "frames" frames.mkdir() (frames / "a.png").touch() (frames / "b.png").touch() context = FakeJobContext(root / "result", "debug-moge2") with patch.object( self.debug_jobs, "infer_moge2", return_value=prediction, ) as infer_one: self.debug_jobs.run_moge2( context, { "input": str(frames), "model": "MOGE2_VITS_NORMAL", "video_frames": 2, "device": "cpu", "resolution_level": 5, }, ) self.assertEqual(infer_one.call_count, 2) self.assertEqual( [call.kwargs["release_memory"] for call in infer_one.call_args_list], [False, True], ) archives = sorted(context.directory.glob("prediction.*.npz")) self.assertEqual(len(archives), 2) with np.load(archives[0]) as result: self.assertEqual(set(result.files), set(prediction)) self.assertFalse(any(context.directory.glob("*.glb"))) def test_moge2_batch_releases_only_after_final_image(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: frames = self.moge2_model.infer( FakeModels(), {"model": "MOGE2_VITS_NORMAL", "device": "cpu"}, [Path("a.png"), Path("b.png")], include_points=True, ) self.assertEqual(len(frames), 2) self.assertEqual( [call.kwargs["release_memory"] for call in infer_one.call_args_list], [False, True], ) def test_ben2_debug_saves_direct_api_foreground(self): class Foreground: def save(self, path): Path(path).write_bytes(b"foreground") with tempfile.TemporaryDirectory() as directory: root = Path(directory) input_path = root / "input.png" input_path.touch() context = FakeJobContext(root / "result", "debug-ben2") with patch.object( self.debug_jobs, "infer_foreground", return_value=(Foreground(), 1.0, 2.0), ) as infer_foreground: self.debug_jobs.run_ben2( context, { "input": str(input_path), "model": "BEN2_BASE", "video_frames": 1, "device": "cpu", }, ) infer_foreground.assert_called_once() self.assertTrue( infer_foreground.call_args.kwargs["release_memory"] ) self.assertEqual( (context.directory / "foreground.png").read_bytes(), b"foreground", ) def test_upscale_debug_saves_each_direct_api_prediction(self): from PIL import Image with tempfile.TemporaryDirectory() as directory: root = Path(directory) frames = root / "frames" frames.mkdir() (frames / "a.png").touch() (frames / "b.png").touch() context = FakeJobContext(root / "result", "debug-upscale") with patch.object( self.debug_jobs, "infer_upscale", return_value=Image.new("RGB", (8, 8), (10, 20, 30)), ) as infer_one: self.debug_jobs.run_upscale( context, { "input": str(frames), "model": "REALESRGAN_GENERAL_WDN_X4V3", "video_frames": 2, "device": "cpu", }, ) self.assertEqual(infer_one.call_count, 2) self.assertEqual( [call.kwargs["release_memory"] for call in infer_one.call_args_list], [False, True], ) predictions = sorted(context.directory.glob("prediction.*.png")) self.assertEqual(len(predictions), 2) with Image.open(predictions[0]) as prediction: self.assertEqual(prediction.size, (8, 8)) self.assertEqual(prediction.mode, "RGB") def test_upscale_job_saves_importable_upscale_sequence(self): from PIL import Image with tempfile.TemporaryDirectory() as directory: root = Path(directory) frames = root / "frames" frames.mkdir() (frames / "a.png").touch() (frames / "b.png").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(frames), "model": "REALESRGAN_GENERAL_WDN_X4V3", "video_frames": 2, "device": "cpu", }, ) self.assertEqual(infer_one.call_count, 2) self.assertEqual( [call.kwargs["release_memory"] for call in infer_one.call_args_list], [False, True], ) self.assertEqual(result, {"upscale": "."}) colors = sorted(context.directory.glob("upscale.*.png")) self.assertEqual( [path.name for path in colors], ["upscale.0001.png", "upscale.0002.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", "video_frames": 1, "device": "cpu", }, ) with Image.open(context.directory / "upscale.0001.png") as result: self.assertEqual(result.mode, "RGBA") self.assertEqual(result.size, (20, 12)) 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) 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) def test_da3_debug_writer_keeps_api_values_without_derived_attributes(self): prediction = SimpleNamespace( depth=np.ones((1, 2, 3), dtype=np.float32), conf=np.arange(6, dtype=np.float32).reshape(1, 2, 3), intrinsics=np.eye(3, dtype=np.float32)[None], extrinsics=np.eye(4, dtype=np.float32)[None], ) with tempfile.TemporaryDirectory() as directory: output_path = self.debug_da3_outputs.write_prediction( prediction, Path(directory) / "prediction.npz", ) with np.load(output_path) as attributes: self.assertEqual(attributes["depth"].shape, (1, 2, 3)) self.assertEqual(attributes["conf"].shape, (1, 2, 3)) self.assertNotIn("confidence_percentile", attributes.files) if __name__ == "__main__": unittest.main()