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 server_test_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)