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 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.ben2 = load_server_module("models.ben2") 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_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_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, "input_frames", return_value=nullcontext([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, ): 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"]) 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.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]), (0, 0, 0, 128), ) def test_ben2_background_removal_processes_each_frame(self): from PIL import Image onnx_ben2 = load_server_module("models.onnx_ben2") with tempfile.TemporaryDirectory() as directory: root = Path(directory) 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") 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)) 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): 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, (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)