icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add new functionality for image rectification and depth metadata handling

Browse Source 
- Implemented a preview tool for rectifying images with interactive aspect controls in `preview.py`.
- Added tests for depth metadata loading to ensure consistent schema usage in `test_depth_metadata.py`.
- Developed comprehensive tests for cutout calibration processes in `test_cutout_calibration.py`, including stability checks and depth calibration.
- Created tests for frame compositing operations in `test_frame_compositing.py`, ensuring correct handling of depth and transparency.
- Introduced material slot inheritance functionality in `material_slots.py` to manage geometry material assignments.
- Built shared shader nodes for material processing in `shader.py`, including roughness and bump mapping setups.
This commit is contained in:
马南
2026-09-12 19:20:06 +08:00
parent e8268504c2
commit 938d95db0f
175 changed files with 3812 additions and 6864 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
tests/anyimage/server/runtime_support.py
View File

@@ -1,28 +1,14 @@
from tests.support.paths import PROJECT_ROOT
import importlib.util
import sys
import importlib
import tempfile
import time
import unittest
from pathlib import Path
def load_server_module():
server_directory = PROJECT_ROOT / "src" / "anyimage" / "server"
specification = importlib.util.spec_from_file_location(
"server",
server_directory / "__init__.py",
submodule_search_locations=[str(server_directory)],
)
module = importlib.util.module_from_spec(specification)
sys.modules["server"] = module
specification.loader.exec_module(module)
return module
class ServerTestCase(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.server_module = load_server_module()
cls.server_module = importlib.import_module("server")
importlib.import_module("server.geometry.prediction_artifacts")
importlib.import_module("server.jobs.cutout")
from blendjob import JobServer
@@ -33,8 +19,7 @@ class ServerTestCase(unittest.TestCase):
cls.ModelManager, cls.ModelSessionCache = ModelManager, ModelSessionCache
cls.onnx_runtime = onnx_runtime
cls.server = cls.server_module.server
cls.app_module = sys.modules["server.app"]
cls.app_module = importlib.import_module("server.app")
def setUp(self):
@@ -50,6 +35,7 @@ class ServerTestCase(unittest.TestCase):
def create_server(self):
source = self.server_module.server
server = self.JobServer("AnyImage Job Server", storage_root=self.storage_root)
self.addCleanup(server.close)
server.handlers.update(source.handlers)
server.resource_factories.update(source.resource_factories)
server.bind(self.storage_root)

44
tests/anyimage/server/support.py
View File

@@ -1,7 +1,6 @@
from tests.support.paths import PROJECT_ROOT
import sys
from pathlib import Path
from types import SimpleNamespace
def load_server_module(name):
@@ -13,49 +12,6 @@ def load_server_module(name):
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="BEN2-ONNX"):
self.model_directory = Path(directory)

23
tests/anyimage/server/test_background_models.py
View File

@@ -13,19 +13,14 @@ from tests.anyimage.server.support import FakeJobContext
@pytest.mark.parametrize("key", model_catalog.BACKGROUND_MODELS)
@pytest.mark.parametrize("kind", ["foreground", "frames", "alpha"])
def test_models_dispatch_keep_alpha_and_reuse_sessions(tmp_path, key, kind):
@pytest.mark.parametrize("kind", ["foreground", "alpha"])
def test_models_dispatch_keep_alpha(tmp_path, key, kind):
source = tmp_path / "source.png"
Image.new("RGBA", (7, 5), (60, 120, 180, 128)).save(source)
if kind == "frames":
source = tmp_path / "frames"
source.mkdir()
for index in (1, 2):
Image.new("RGBA", (7, 5), (60, 120, 180, 128)).save(source / f"{index}.png")
manager = ModelManager(tmp_path / "models")
context = FakeJobContext(tmp_path / "job", "remove-background", manager)
adapter = background.model_adapter(manager.directory(key))
parameters = {"input": str(source), "model": key, "device": "directml", "video_frames": 2}
parameters = {"input": str(source), "model": key, "device": "directml"}
if kind == "alpha":
parameters["output_kind"] = kind
with patch.object(adapter, "create_session", return_value=object()) as create, patch.object(
@@ -33,7 +28,7 @@ def test_models_dispatch_keep_alpha_and_reuse_sessions(tmp_path, key, kind):
) as infer:
remove_background.run(context, parameters)
create.assert_called_once_with(manager.directory(key), "cpu" if key == "BIREFNET_HR_MATTING" else "directml")
assert infer.call_count == (2 if kind == "frames" else 1)
infer.assert_called_once()
if kind == "alpha":
np.testing.assert_allclose(np.load(context.directory / "alpha.npy"), .4)
else:
@@ -63,13 +58,3 @@ def test_readiness_checks_model_and_license_corruption(tmp_path, key):
assert manager.ready(key)
(directory / "LICENSE.txt").write_bytes(b"damaged!")
assert not manager.ready(key)
@pytest.mark.parametrize("key", ["PISA_SR", "SAFMN_REAL_X4", "MOGE2_VITL_NORMAL"])
def test_retired_job_keys_fail_before_loading(key):
from server.models import upscale
with pytest.raises(ValueError, match="Unknown upscale model"):
upscale.infer_one(None, {"model": key}, "unused.png")
with pytest.raises(ValueError, match="Unknown background removal model"):
remove_background.run(None, {"model": key})

69
tests/anyimage/server/test_geometry_artifacts.py
View File

@@ -149,73 +149,6 @@ class GeometryArtifactsTest(ServerTestCase):
np.testing.assert_array_equal(outputs[0], first[..., :3])
def test_geometry_writers_reject_missing_or_non_finite_fields(self):
import numpy as np
depth = importlib.import_module("server.geometry.depth_texture")
normal = importlib.import_module("server.geometry.normal_texture")
GeometryFrame = importlib.import_module(
"server.models.geometry"
).GeometryFrame
base = {
"depth": np.ones((1, 2), dtype=np.float32),
"validity": np.ones((1, 2), dtype=np.float32),
"intrinsics": np.eye(3, dtype=np.float32),
}
with self.assertRaisesRegex(ValueError, "point field"):
depth.write_depth_texture(
GeometryFrame(**base),
self.storage_root / "missing-points.exr",
alpha=base["validity"],
)
point_frame = GeometryFrame(
**base,
points=np.asarray([[(0.0, 0.0, 1.0), (np.nan, 0.0, 1.0)]], dtype=np.float32),
)
with self.assertRaisesRegex(ValueError, "Point field"):
depth.write_depth_texture(
point_frame,
self.storage_root / "nan-points.exr",
alpha=base["validity"],
)
depth_frame = GeometryFrame(
**base, points=np.asarray([[(0, 0, 1), (0, 0, np.inf)]], dtype=np.float32),
)
with self.assertRaisesRegex(ValueError, "Point field"):
depth.write_depth_texture(
depth_frame,
self.storage_root / "inf-depth.exr",
alpha=base["validity"],
)
with self.assertRaisesRegex(ValueError, "normal field"):
normal.write_camera_normal_texture(
GeometryFrame(**base),
np.ones((1, 2), dtype=bool),
(0, 0),
(2, 1),
(2, 1),
self.storage_root / "missing-normal.png",
)
normal_frame = GeometryFrame(
**base,
normal=np.asarray(
[[(0.0, 0.0, 1.0), (0.0, np.nan, 1.0)]],
dtype=np.float32,
),
)
with self.assertRaisesRegex(ValueError, "Normal field"):
normal.write_camera_normal_texture(
normal_frame,
np.ones((1, 2), dtype=bool),
(0, 0),
(2, 1),
(2, 1),
self.storage_root / "nan-normal.png",
)
self.assertFalse(any(self.storage_root.glob("*.exr")))
self.assertFalse(any(self.storage_root.glob("*-normal.png")))
def test_moge_artifact_orchestration_uses_one_geometry_frame(self):
import numpy as np
from PIL import Image
@@ -246,7 +179,7 @@ class GeometryArtifactsTest(ServerTestCase):
texture = importlib.import_module("server.geometry.depth_texture")
with (
patch.object(artifacts.moge, "infer", return_value=[frame]) as infer,
patch.object(artifacts.moge, "infer", return_value=frame) as infer,
patch.object(texture, "write_float_exr", wraps=texture.write_float_exr) as write,
):
result = artifacts.generate_moge_artifacts(

83
tests/anyimage/server/test_input_validation.py
View File

@@ -1,5 +1,3 @@
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
from PIL import Image
@@ -12,54 +10,19 @@ def media():
return load_server_module("media.input")
@pytest.mark.parametrize("suffix", [".PNG", ".MP4", ".webp", ".mts"])
@pytest.mark.parametrize("suffix", [".PNG", ".JPG", ".webp", ".jpeg"])
def test_supported_format_is_case_insensitive(media, tmp_path, suffix):
path = tmp_path / f"input{suffix}"
path.touch()
assert media.validate_input_path(path) == path
@pytest.mark.parametrize("suffix", [".tiff", ""])
def test_frontend_and_backend_share_format_reason(media, tmp_path, suffix):
path = tmp_path / f"input{suffix}"
path.touch()
with pytest.raises(RuntimeError) as preflight:
media.validate_input_path(path)
with pytest.raises(RuntimeError) as backend:
media.prepare_input(path, 16, tmp_path)
assert str(preflight.value) == str(backend.value)
assert f"Unsupported input format: {suffix or '(no extension)'}" in str(backend.value)
assert "Supported images: .jpeg, .jpg, .png, .webp." in str(backend.value)
assert "Supported videos:" in str(backend.value)
@pytest.mark.parametrize("suffix,kind", [(".png", "image"), (".mp4", "video")])
@pytest.mark.parametrize("suffix,kind", [(".png", "image"), (".mp4", "image")])
def test_missing_path_precedes_format_validation(media, tmp_path, suffix, kind):
with pytest.raises(RuntimeError, match=f"Source {kind} file does not exist"):
media.validate_input_path(tmp_path / f"missing{suffix}")
def test_directory_filtering_and_empty_reason(media, tmp_path):
(tmp_path / "note.txt").touch()
(tmp_path / "frame.exr").touch()
with pytest.raises(RuntimeError, match="No supported image frames found"):
media.validate_input_path(tmp_path)
for name in ["b.PNG", "a.jpg"]:
(tmp_path / name).touch()
frames, temporary = media.prepare_input(tmp_path, 1, tmp_path, start_frame=2)
assert frames == [tmp_path / "b.PNG"]
assert temporary is None
def test_image_subset_rejects_video(media, tmp_path):
path = tmp_path / "clip.mp4"
path.touch()
with pytest.raises(RuntimeError) as error:
media.validate_input_path(path, images_only=True)
assert "Supported images:" in str(error.value)
assert "Supported videos:" not in str(error.value)
def test_image_decode_error_keeps_cause(media, tmp_path):
path = tmp_path / "broken.png"
path.write_bytes(b"broken")
@@ -79,15 +42,33 @@ def test_processing_error_is_not_relabelled(media, tmp_path):
assert error.value is failure
def test_video_decode_error_closes_container(media, tmp_path):
failure = ValueError("broken packet")
def frames():
raise failure
yield
container = SimpleNamespace(streams=SimpleNamespace(video=[object()]),
decode=lambda _stream: frames(), close=Mock())
with pytest.raises(RuntimeError, match="Unable to read video") as error:
media.extract_video_frames("broken.mp4", tmp_path, 2,
SimpleNamespace(open=lambda _path: container))
assert error.value.__cause__ is failure
container.close.assert_called_once()
@pytest.mark.parametrize("suffix", [".mp4", ".mts", ".gif", ".tiff", ""])
def test_unsupported_formats_are_rejected(media, tmp_path, suffix):
path = tmp_path / f"input{suffix}"
path.touch()
with pytest.raises(RuntimeError, match="Unsupported input format"):
media.validate_input_path(path)
def test_frame_directory_is_rejected(media, tmp_path):
Image.new("RGB", (2, 2)).save(tmp_path / "frame.png")
with pytest.raises(RuntimeError, match="Source image file does not exist"):
media.validate_input_path(tmp_path)
def test_animated_png_is_rejected(media, tmp_path):
path = tmp_path / "animated.png"
Image.new("RGB", (2, 2), "red").save(
path, save_all=True, append_images=[Image.new("RGB", (2, 2), "blue")], duration=100,
)
with pytest.raises(RuntimeError, match="Only still images"):
with media.open_image(path):
pytest.fail("Animated images must not reach processing")
def test_ai_tasks_support_cooperative_cancellation(media):
with pytest.raises(media.TaskCancelled):
media.check_cancelled(lambda: True)
media.check_cancelled(lambda: False)

8
tests/anyimage/server/test_job_server.py
View File

@@ -7,6 +7,7 @@ from tests.anyimage.server.runtime_support import ServerTestCase
class JobServerTest(ServerTestCase):
def test_dedicated_server_runs_only_one_job_at_a_time(self):
server = self.JobServer("Test", storage_root=self.storage_root)
self.addCleanup(server.close)
first_started = threading.Event()
release_first = threading.Event()
order = []
@@ -26,16 +27,15 @@ class JobServerTest(ServerTestCase):
self.assertIsNotNone(second)
self.assertTrue(server.snapshot("instance")["busy"])
self.assertEqual(server.snapshot("instance")["queued_jobs"], 1)
self.assertNotIn("clients", server.snapshot("instance"))
release_first.set()
self.wait_for_job(first)
self.wait_for_job(second)
self.assertEqual(order, ["first", "second"])
server.close()
def test_queued_job_can_be_cancelled_before_it_runs(self):
server = self.JobServer("Test", storage_root=self.storage_root)
self.addCleanup(server.close)
first_started = threading.Event()
release_first = threading.Event()
executed = []
@@ -58,7 +58,6 @@ class JobServerTest(ServerTestCase):
self.assertEqual(second_status["state"], "cancelled")
self.assertEqual(executed, ["first"])
server.close()
def test_cancelling_job_keeps_server_and_model_cache_alive(self):
@@ -94,6 +93,7 @@ class JobServerTest(ServerTestCase):
def test_job_progress_is_kept_in_server_memory(self):
server = self.JobServer("Test", storage_root=self.storage_root)
self.addCleanup(server.close)
@server.job("example")
def example(context, _parameters):
@@ -111,6 +111,7 @@ class JobServerTest(ServerTestCase):
def test_runtime_storage_root_owns_the_jobs_directory(self):
storage_root = self.storage_root
server = self.JobServer("Test")
self.addCleanup(server.close)
server.bind(storage_root)
@server.job("example")
@@ -133,6 +134,7 @@ class JobServerTest(ServerTestCase):
close=lambda: events.append("close"),
)
server = self.JobServer("Test", storage_root=self.storage_root)
self.addCleanup(server.close)
server.add_resource("model_manager", resource)
@server.job("example")

167
tests/anyimage/server/test_media_input.py
View File

@@ -1,167 +0,0 @@
import tempfile
import unittest
from pathlib import Path
from tests.anyimage.server.support import load_server_module, _make_fake_av
class MediaInputTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.media = load_server_module("media.input")
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_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 format: .txt"):
self.media.prepare_input(input_path, 16, directory)

11
tests/anyimage/server/test_model_birefnet_onnx.py
View File

@@ -23,14 +23,3 @@ def test_probability_alpha_keeps_transparency_and_source_dimensions(size):
assert run.call_args.kwargs["release_memory"] is False
feed = run.call_args.args[2]["image"]
np.testing.assert_allclose(feed[0, :, 0, 0], (np.array([64, 128, 192]) / 255 - [.485, .456, .406]) / [.229, .224, .225], atol=1e-6)
@pytest.mark.parametrize("value", [np.nan, np.inf, -0.1, 1.1])
def test_invalid_alpha_is_rejected(value):
session = SimpleNamespace(
get_inputs=lambda: [SimpleNamespace(name="image", shape=[1, 3, 1024, 1024])],
get_outputs=lambda: [SimpleNamespace(name="alpha")],
)
with patch.object(onnx_birefnet, "run_session", return_value=[np.full((1, 1, 1024, 1024), value)]):
with pytest.raises(RuntimeError, match="invalid alpha"):
onnx_birefnet.infer_alpha(session, Image.new("RGB", (2, 3)))

104
tests/anyimage/server/test_model_catalog.py
View File

@@ -1,104 +0,0 @@
from tests.support.paths import PROJECT_ROOT
import importlib.util
import unittest
def load_module(name, relative_path):
specification = importlib.util.spec_from_file_location(
name,
PROJECT_ROOT / relative_path,
)
module = importlib.util.module_from_spec(specification)
specification.loader.exec_module(module)
return module
class ModelCatalogTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.models = load_module(
"anyimage_server_models",
"src/anyimage/server/model_catalog.py",
)
def test_unknown_model_is_rejected(self):
with self.assertRaisesRegex(ValueError, "Unknown model"):
self.models.get_downloadable_model("UNKNOWN")
def test_three_tiers_and_retired_keys(self):
self.assertEqual(list(self.models.MOGE_MODELS), ["MOGE2_VITS_NORMAL", "MOGE2_VITB_NORMAL", "MOGE3_VITL"])
self.assertEqual(list(self.models.BACKGROUND_MODELS), ["BIREFNET_LITE", "BEN2_BASE", "BIREFNET_HR_MATTING"])
self.assertEqual(list(self.models.UPSCALE_MODELS), ["REALESRGAN_GENERAL_WDN_X4V3", "REALESRGAN_X4PLUS", "HAT_GAN_X4_SHARPER"])
self.assertEqual(len(self.models.DOWNLOADABLE_MODELS), 9)
self.assertEqual(self.models.DEFAULT_MODEL_KEYS, ("MOGE2_VITS_NORMAL", "BIREFNET_LITE", "REALESRGAN_GENERAL_WDN_X4V3"))
for key in ("DA3_BASE", "DA3_SMALL", "MOGE2_VITL_NORMAL", "SAFMN_REAL_X4", "PISA_SR"):
with self.assertRaisesRegex(ValueError, "Unknown model"):
self.models.get_downloadable_model(key)
def test_ben2_download_only_includes_required_hugging_face_files(self):
model = self.models.get_downloadable_model("BEN2_BASE")
self.assertEqual(model.repository, "onnx-community/BEN2-ONNX")
self.assertEqual(model.license_name, "Apache 2.0")
self.assertEqual(model.download_patterns, ("onnx/model_fp16.onnx",))
self.assertEqual(model.ready_patterns, ("onnx/model_fp16.onnx",))
self.assertEqual(model.r2_files[0][0], "onnx/model_fp16.onnx")
def test_ben2_model_directory_uses_repository_name(self):
self.assertEqual(
self.models.get_downloadable_model("BEN2_BASE").directory_name,
"BEN2-ONNX",
)
def test_upscale_models_use_verified_onnx_artifacts(self):
expected = {
"REALESRGAN_GENERAL_WDN_X4V3": ("realesr-general-wdn-x4v3.onnx", 4867419, "b021a837"),
"HAT_GAN_X4_SHARPER": ("model.onnx", 45874438, "032d8b03"),
"REALESRGAN_X4PLUS": ("realesrgan-x4plus.onnx", 67051618, "4ed6a45a"),
}
self.assertEqual(set(self.models.UPSCALE_MODELS), set(expected))
for key, (filename, size, checksum_prefix) in expected.items():
with self.subTest(model=key):
model = self.models.get_downloadable_model(key)
self.assertIn(filename, model.download_patterns)
self.assertIn(filename, model.ready_patterns)
self.assertEqual(model.r2_files[0][:2], (filename, size))
self.assertTrue(model.r2_files[0][2].startswith(checksum_prefix))
self.assertEqual(self.models.model_record(model)["family"], "upscale")
wdn = self.models.UPSCALE_MODELS["REALESRGAN_GENERAL_WDN_X4V3"]
self.assertEqual(wdn.repository, "xinntao/Real-ESRGAN")
self.assertEqual(wdn.huggingface_repository, "")
self.assertEqual(
self.models.UPSCALE_MODELS["HAT_GAN_X4_SHARPER"].huggingface_repository, ""
)
def test_moge2_models_use_official_onnx_models(self):
for key, variant in (("MOGE2_VITS_NORMAL", "vits"),
("MOGE2_VITB_NORMAL", "vitb")):
with self.subTest(key=key):
model = self.models.get_downloadable_model(key)
self.assertEqual(
model.repository,
f"Ruicheng/moge-2-{variant}-normal-onnx",
)
self.assertEqual(model.license_name, "MIT")
self.assertEqual(model.download_patterns, ("model.onnx",))
self.assertEqual(model.ready_patterns, ("model.onnx",))
self.assertEqual(model.r2_directory, model.repository.split("/")[-1])
self.assertEqual(model.r2_files[0][0], "model.onnx")
def test_moge2_models_match_uploaded_r2_variants(self):
self.assertEqual(
list(self.models.MOGE2_MODELS),
["MOGE2_VITS_NORMAL", "MOGE2_VITB_NORMAL"],
)
self.assertEqual(
[model.r2_directory for model in self.models.MOGE2_MODELS.values()],
[
"moge-2-vits-normal-onnx",
"moge-2-vitb-normal-onnx",
],
)

25
tests/anyimage/server/test_model_geometry.py
View File

@@ -3,26 +3,9 @@ import numpy as np
from anyimage.server.models import moge as moge2_model
from anyimage.server.models.geometry import GeometryFrame
def test_geometry_frame_requires_matching_float32_fields():
with np.testing.assert_raises_regex(ValueError, "validity"):
GeometryFrame(
depth=np.ones((2, 2), dtype=np.float32),
validity=np.ones((1, 2), dtype=np.float32),
intrinsics=np.eye(3, dtype=np.float32),
)
with np.testing.assert_raises_regex(ValueError, "normal"):
GeometryFrame(
depth=np.ones((2, 2), dtype=np.float32),
validity=np.ones((2, 2), dtype=np.float32),
intrinsics=np.eye(3, dtype=np.float32),
normal=np.ones((2, 2, 2), dtype=np.float32),
)
def test_moge_infer_returns_geometry_frames_with_model_validity():
def test_moge_infer_returns_geometry_frame_with_model_validity():
mask = np.asarray(((1.0, 0.2), (0.7, 0.0)), dtype=np.float32)
prediction = {
"depth": np.ones((2, 2), dtype=np.float32),
@@ -36,10 +19,8 @@ def test_moge_infer_returns_geometry_frames_with_model_validity():
}
with patch.object(moge2_model, "infer_one", return_value=prediction):
frames = moge2_model.infer(None, {}, ["image"], include_points=True)
frame = moge2_model.infer(None, {}, "image", include_points=True)
assert len(frames) == 1
frame = frames[0]
assert np.array_equal(frame.validity, mask.astype(np.float32))
assert np.array_equal(frame.normal, prediction["normal"])
assert np.array_equal(frame.points, prediction["points"])
@@ -59,7 +40,7 @@ def test_moge_z_depth_frame_does_not_require_points():
}
with patch.object(moge2_model, "infer_one", return_value=prediction):
frame = moge2_model.infer(None, {}, ["image"], include_points=False)[0]
frame = moge2_model.infer(None, {}, "image", include_points=False)
assert frame.points is None
assert frame.validity.all()

13
tests/anyimage/server/test_model_inference.py
View File

@@ -19,7 +19,7 @@ class ModelInferenceTest(unittest.TestCase):
cls.onnx_upscale = load_server_module("models.onnx_upscale")
def test_moge2_batch_releases_only_after_final_image(self):
def test_moge2_infers_one_geometry_frame(self):
prediction = {
"points": np.ones((2, 2, 3), dtype=np.float32),
"depth": np.ones((2, 2), dtype=np.float32),
@@ -32,18 +32,15 @@ class ModelInferenceTest(unittest.TestCase):
"infer_one",
return_value=prediction,
) as infer_one:
frames = self.moge2_model.infer(
frame = self.moge2_model.infer(
FakeModels(),
{"model": "MOGE2_VITS_NORMAL", "device": "cpu"},
[Path("a.png"), Path("b.png")],
Path("input.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],
)
self.assertEqual(frame.depth.shape, (2, 2))
infer_one.assert_called_once()
def test_upscale_resource_error_releases_session_and_is_readable(self):

1
tests/anyimage/server/test_model_manager.py
View File

@@ -73,7 +73,6 @@ class ModelManagerTest(ServerTestCase):
},
)
self.assertIn("server.model_catalog", sys.modules)
self.assertNotIn("model_catalog", sys.modules)
def test_background_model_is_reused_for_matching_runtime_key(self):

6
tests/anyimage/server/test_model_onnx_runtime.py
View File

@@ -122,12 +122,6 @@ class OnnxRuntimeTest(unittest.TestCase):
["CPUExecutionProvider"],
)
def test_old_mps_device_is_rejected(self):
with self.assertRaises(KeyError):
onnx_runtime.select_providers(
"mps",
["CoreMLExecutionProvider", "CPUExecutionProvider"],
)
def test_completing_cuda_run_shrinks_memory_arena(self):
configs = []

14
tests/anyimage/server/test_model_upscale_onnx.py
View File

@@ -33,10 +33,6 @@ class UpscaleOnnxTest(unittest.TestCase):
expected = np.repeat(np.repeat(pixels, 4, axis=0), 4, axis=1)
np.testing.assert_array_equal(np.asarray(result), expected)
def test_nonfinite_pixels_fail_before_conversion(self):
with patch.object(onnx_upscale, "run_session", return_value=[np.full((1, 3, 1024, 1024), np.nan)]):
with self.assertRaisesRegex(RuntimeError, "non-finite"):
onnx_upscale.infer(FakeSession(), Image.new("RGB", (7, 5)))
def test_preprocess_matches_dynamic_rgb_layout(self):
image = Image.new("RGB", (7, 5), (64, 128, 192))
@@ -95,16 +91,6 @@ class UpscaleOnnxTest(unittest.TestCase):
self.assertAlmostEqual(red[16, 22], 120 / 255.0)
self.assertAlmostEqual(red[16, 23], 100 / 255.0)
def test_infer_rejects_non_four_x_output(self):
class InvalidSession(FakeSession):
def run(self, _output_names, feeds):
return [feeds["input"]]
with self.assertRaisesRegex(RuntimeError, "unexpected output shape"):
onnx_upscale.infer(
InvalidSession(),
Image.new("RGB", (7, 5)),
)
def test_all_models_split_large_inputs_into_fixed_tiles(self):
class CountingSession(FakeSession):

4
tests/anyimage/server/test_panorama.py
View File

@@ -65,7 +65,7 @@ def panorama_job(tmp_path, monkeypatch):
write = job.write_float_exr
monkeypatch.setattr(job, "write_float_exr", lambda rgba, path: (pixels.append(rgba.copy()), write(rgba, path))[1])
parameters = {"input": str(source), "model": "MOGE2", "device": "cpu", "resolution_level": 3,
"max_input_size": 32, "video_frames": 1, "start_frame": 1}
"max_input_size": 32}
return source, context, parameters, calls, pixels
@@ -128,6 +128,6 @@ def test_failed_jobs_leave_no_partial_artifacts(panorama_job, monkeypatch, kind)
with pytest.raises((ValueError, RuntimeError)):
job.run(context, parameters)
assert not list(context.directory.glob("depth.*"))
assert not (context.directory / "image.0001.png").exists()
assert not (context.directory / "foreground.png").exists()
if kind == "transparent":
assert not calls

2
tests/anyimage/server/test_prediction_artifacts.py
View File

@@ -73,7 +73,7 @@ def test_artifacts_preserve_original_geometry_with_alpha(
check_cancelled=lambda: None,
)
with (
patch.object(artifacts.moge, "infer", return_value=[frame]),
patch.object(artifacts.moge, "infer", return_value=frame),
patch.object(
artifacts, "write_depth_texture",
return_value=tmp_path / "depth.exr",

73
tests/anyimage/server/test_server_jobs.py
View File

@@ -1,9 +1,6 @@
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
@@ -25,17 +22,6 @@ class ServerJobsTest(unittest.TestCase):
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"
@@ -51,8 +37,8 @@ class ServerJobsTest(unittest.TestCase):
with (
patch.object(
self.depth_plane_job,
"input_frames",
return_value=nullcontext([input_path]),
"validate_input_path",
return_value=input_path,
),
patch.object(
self.depth_plane_job,
@@ -96,8 +82,10 @@ class ServerJobsTest(unittest.TestCase):
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_BASE", "device": "cpu"})
with Image.open(root / "result" / "image.0001.png") as image:
result = self.background.run(context, {"input": str(input_path), "model": "BEN2_BASE", "device": "cpu"})
self.assertEqual(result, {"foreground": "foreground.png"})
self.assertEqual([p.name for p in context.directory.glob("*.png")], ["foreground.png"])
with Image.open(root / "result" / "foreground.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})
@@ -152,27 +140,6 @@ class ServerJobsTest(unittest.TestCase):
)
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_BASE", "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")
@@ -192,15 +159,13 @@ class ServerJobsTest(unittest.TestCase):
self.assertFalse(list(context.directory.glob("*.png")))
def test_upscale_job_saves_importable_upscale_sequence(self):
def test_upscale_job_returns_one_image_file(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()
input_path = root / "input.png"
input_path.touch()
context = FakeJobContext(root / "result", "upscale-image")
with patch.object(
self.upscale_job,
@@ -210,23 +175,19 @@ class ServerJobsTest(unittest.TestCase):
result = self.upscale_job.run(
context,
{
"input": str(frames),
"input": str(input_path),
"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(infer_one.call_count, 1)
self.assertEqual(result, {"upscale": "upscale.png"})
colors = sorted(context.directory.glob("*.png"))
self.assertEqual(
[path.name for path in colors],
["upscale.0001.png", "upscale.0002.png"],
["upscale.png"],
)
self.assertIn(
(
@@ -261,12 +222,12 @@ class ServerJobsTest(unittest.TestCase):
{
"input": str(source_path),
"model": "REALESRGAN_GENERAL_WDN_X4V3",
"video_frames": 1,
"device": "cpu",
},
)
with Image.open(context.directory / "upscale.0001.png") as result:
with Image.open(context.directory / "upscale.png") as result:
self.assertEqual(result.mode, "RGBA")
self.assertEqual(result.size, (10, 6))
alpha = np.asarray(result.getchannel("A"))