nan
4c902ed7db
- Implemented `image_plane.py` to build the Image Plane Geometry Nodes group with customizable parameters for geometry, subdivide, and thickness. - Created `image_plane_validation.py` to validate the functionality of the image plane, ensuring correct geometry generation and material assignment. - Developed `material_layers.py` to construct shader node groups for image materials, including depth and normal handling. - Authored `plane_layouts.py` to define node arrangements for image geometry assets. - Introduced `shadeless.py` to create a shadeless shader node group that displays color without scene lighting while preserving shadows and reflections.
438 lines
15 KiB
Python
438 lines
15 KiB
Python
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_depth_plane_job_requests_depth_texture_and_tangent_normal_together(self):
|
|
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"
|
|
tangent_normal = context.directory / "tangent-normal.png"
|
|
parameters = {"input": str(input_path), "generate_depth": True}
|
|
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,
|
|
"tangent_normal": tangent_normal,
|
|
},
|
|
) as generate,
|
|
):
|
|
result = self.depth_plane_job.run(context, parameters)
|
|
|
|
self.assertEqual(
|
|
result,
|
|
{
|
|
"depth": "depth.exr",
|
|
"depth_metadata": "depth.json",
|
|
"tangent_normal": "tangent-normal.png",
|
|
},
|
|
)
|
|
generate.assert_called_once_with(
|
|
context,
|
|
parameters,
|
|
input_path,
|
|
generate_depth=True,
|
|
normal_mode="TANGENT",
|
|
)
|
|
|
|
|
|
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)
|