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, "copy_color_images"), patch.object( self.depth_plane_job, "generate_moge2_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): 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_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)