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

feat(node_assets): add image plane geometry nodes and validation scripts

Browse Source 
- 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.
This commit is contained in:
马南
2026-09-08 23:24:27 +08:00
parent 196b964efc
commit 4c902ed7db
155 changed files with 9700 additions and 10393 deletions
Download Patch File Download Diff File Expand all files Collapse all files

753
tests/test_server_jobs.py
View File

@@ -1,4 +1,3 @@
import os
import sys
import tempfile
import unittest
@@ -6,129 +5,25 @@ from contextlib import nullcontext
from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import patch
import numpy as np
PROJECT_ROOT = Path(__file__).resolve().parents[1]
def load_server_module(name):
addon_directory = PROJECT_ROOT / "src" / "anyimage"
sys.path.insert(0, str(addon_directory))
try:
return __import__(f"server.{name}", fromlist=[name])
finally:
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="model"):
self.model_directory = Path(directory)
def directory(self, _model):
return self.model_directory
def get_ben2(self, _model_directory, _device):
from server.models.onnx_ben2 import create_session
return create_session(self.model_directory, "cpu"), 0.0
def get_da3(self, _model_directory, _device):
return object()
def get_upscale(self, _model_directory, _device):
return object()
class FakeJobContext:
def __init__(self, directory, job_type, models=None):
self.directory = Path(directory)
self.job_type = job_type
self.models = models or FakeModels()
self.status = {}
self.logs = []
def resource(self, name):
if name != "model_manager":
raise KeyError(name)
return self.models
def progress(self, progress, message, **details):
self.status = {
"progress": progress,
"message": message,
**details,
}
def log(self, message, *, level="INFO"):
self.logs.append((level, message))
def check_cancelled(self):
return None
from server_test_support import load_server_module, FakeJobContext
class ServerJobsTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.media = load_server_module("media.input")
cls.resolution = load_server_module("media.resolution")
cls.geometry = load_server_module("geometry")
cls.background = load_server_module("jobs.remove_background")
cls.ben2 = load_server_module("models.ben2")
cls.da3_model = load_server_module("models.da3")
cls.moge2_model = load_server_module("models.moge2")
cls.upscale_model = load_server_module("models.upscale")
cls.onnx_upscale = load_server_module("models.onnx_upscale")
cls.debug_da3_outputs = load_server_module("outputs.debug_da3")
cls.debug_moge2_outputs = load_server_module("outputs.debug_moge2")
cls.debug_jobs = load_server_module("jobs.debug_predictions")
cls.depth_plane_job = load_server_module("jobs.depth_plane")
cls.upscale_job = load_server_module("jobs.upscale")
cls.download = load_server_module("model_download")
def test_upscale_progress_omits_single_frame_count(self):
self.assertEqual(
@@ -140,16 +35,17 @@ class ServerJobsTest(unittest.TestCase):
"Upscaling image (2/3)",
)
def test_depth_plane_job_requests_z_depth_and_tangent_normal_together(self):
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")
z_depth = context.directory / "z-depth.exr"
depth_texture = context.directory / "depth.exr"
metadata = context.directory / "depth.json"
tangent_normal = context.directory / "tangent-normal.png"
parameters = {"input": str(input_path)}
parameters = {"input": str(input_path), "generate_depth": True}
with (
patch.object(
self.depth_plane_job,
@@ -161,7 +57,7 @@ class ServerJobsTest(unittest.TestCase):
self.depth_plane_job,
"generate_moge2_artifacts",
return_value={
"z_depth": z_depth,
"depth": depth_texture,
"depth_metadata": metadata,
"tangent_normal": tangent_normal,
},
@@ -172,7 +68,7 @@ class ServerJobsTest(unittest.TestCase):
self.assertEqual(
result,
{
"z_depth": "z-depth.exr",
"depth": "depth.exr",
"depth_metadata": "depth.json",
"tangent_normal": "tangent-normal.png",
},
@@ -181,318 +77,10 @@ class ServerJobsTest(unittest.TestCase):
context,
parameters,
input_path,
depth_mode="Z",
generate_depth=True,
normal_mode="TANGENT",
)
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_copy_color_images_numbers_a_still_image(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
source = root / "source.png"
source.write_bytes(b"frame")
output = root / "result"
output.mkdir()
self.media.copy_color_images([source], output)
self.assertEqual(
(output / "image.0001.png").read_bytes(),
b"frame",
)
def test_copy_frames_skips_a_still_image_by_default(self):
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
source = root / "source.png"
source.touch()
output = root / "result"
output.mkdir()
self.media.copy_frames([source], output)
self.assertEqual(list(output.iterdir()), [])
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 type"):
self.media.prepare_input(input_path, 16, directory)
def test_download_model_uses_actual_request_without_preflight(self):
calls = []
fake_hugging_face = ModuleType("huggingface_hub")
fake_hugging_face.snapshot_download = lambda **options: calls.append(
options
)
args = SimpleNamespace(
repository="depth-anything/DA3-SMALL",
model_dir="model",
progress=lambda *_args, **_kwargs: None,
allow_pattern=[],
)
with tempfile.TemporaryDirectory() as directory, patch.dict(
sys.modules,
{"huggingface_hub": fake_hugging_face},
), patch.object(
self.download,
"create_download_progress_class",
return_value=object,
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.download,
"_verify_model_complete",
return_value=None,
):
args.model_dir = str(Path(directory) / "models" / "DA3")
self.download.download_model(args)
self.assertEqual(len(calls), 1)
def test_download_model_refuses_success_when_files_missing(self):
fake_hugging_face = ModuleType("huggingface_hub")
fake_hugging_face.snapshot_download = lambda **options: None
args = SimpleNamespace(
repository="PramaLLC/BEN2",
model_dir="model",
progress=lambda *_args, **_kwargs: None,
allow_pattern=[],
)
with tempfile.TemporaryDirectory() as directory, patch.dict(
sys.modules,
{"huggingface_hub": fake_hugging_face},
), patch.object(
self.download,
"create_download_progress_class",
return_value=object,
), patch.dict(os.environ, {"HF_ENDPOINT": ""}):
args.model_dir = str(Path(directory) / "model")
with self.assertRaisesRegex(
RuntimeError,
"All download sources failed",
):
self.download.download_model(args)
def test_download_model_reports_all_failed_sources(self):
args = SimpleNamespace(
repository="depth-anything/DA3-SMALL",
model_dir="model",
progress=lambda *_args, **_kwargs: None,
allow_pattern=[],
r2_directory="DA3Small",
r2_file=["config.json|100|abc"],
)
with tempfile.TemporaryDirectory() as directory, patch.object(
self.download,
"download_from_hf",
side_effect=OSError("direct down"),
), patch.object(
self.download,
"download_from_r2",
side_effect=OSError("mirror down"),
):
args.model_dir = str(Path(directory) / "model")
with self.assertRaisesRegex(RuntimeError, "models.omoolab.xyz"):
self.download.download_model(args)
def test_download_model_prefers_direct_hf_then_r2(self):
args = SimpleNamespace(
repository="depth-anything/DA3-SMALL",
model_dir="model",
progress=lambda *_args, **_kwargs: None,
allow_pattern=["config.json"],
r2_directory="DA3Small",
r2_file=["config.json|100|abc"],
)
order = []
def fake_download_from_hf(*_args, **_kwargs):
order.append("hf")
raise OSError("hf down")
def fake_download_from_r2(*_args, **_kwargs):
order.append("r2")
with tempfile.TemporaryDirectory() as directory, patch.object(
self.download,
"download_from_hf",
side_effect=fake_download_from_hf,
), patch.object(
self.download,
"download_from_r2",
side_effect=fake_download_from_r2,
), patch.object(
self.download,
"_verify_model_complete",
return_value=None,
):
args.model_dir = str(Path(directory) / "models" / "DA3")
self.download.download_model(args)
self.assertEqual(order, ["hf", "r2"])
def test_download_from_r2_downloads_declared_files(self):
args = SimpleNamespace(
repository="depth-anything/DA3-SMALL",
model_dir="model",
progress=lambda *_args, **_kwargs: None,
allow_pattern=["config.json"],
r2_directory="DA3Small",
r2_file=["config.json|100|abc"],
)
downloaded = []
with tempfile.TemporaryDirectory() as directory, patch.object(
self.download,
"_download_file_from_url",
side_effect=lambda *args, **kwargs: downloaded.append(args[1].name),
):
args.model_dir = str(Path(directory) / "model")
self.download.download_from_r2(args, args.model_dir)
self.assertEqual(downloaded, ["config.json"])
def test_ben2_background_removal_writes_fast_transparent_png(self):
inference_options = []
@@ -591,6 +179,7 @@ class ServerJobsTest(unittest.TestCase):
)
self.assertIn("Background removed in", context.status["message"])
def test_ben2_preserves_input_alpha_and_hides_transparent_rgb(self):
from PIL import Image
@@ -638,6 +227,7 @@ class ServerJobsTest(unittest.TestCase):
(10, 20, 30, 128),
)
def test_ben2_background_removal_processes_each_frame(self):
inference_options = []
@@ -740,250 +330,6 @@ class ServerJobsTest(unittest.TestCase):
[False, True],
)
def test_da3_debug_saves_only_direct_api_prediction(self):
prediction = SimpleNamespace(
depth=np.full((1, 2, 2), 2.0, dtype=np.float32),
conf=np.ones((1, 2, 2), dtype=np.float32),
extrinsics=np.eye(4, dtype=np.float32)[None, :3],
intrinsics=np.eye(3, dtype=np.float32)[None],
processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8),
is_metric=0,
sky=None,
scale_factor=None,
)
fake_da3 = ModuleType("server.models.onnx_da3")
calls = []
def fake_infer(_session, images, **options):
calls.append((images, options))
return prediction
fake_da3.infer = fake_infer
fake_runtime = ModuleType("server.models.onnx_runtime")
fake_runtime.provider_label = lambda _session: "CPU"
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
input_path = root / "input.png"
input_path.touch()
with patch.dict(
sys.modules,
{
"server.models.onnx_da3": fake_da3,
"server.models.onnx_runtime": fake_runtime,
},
):
context = FakeJobContext(root / "result", "debug-da3")
self.debug_jobs.run_da3(
context,
{
"input": str(input_path),
"model": "DA3",
"video_frames": 1,
"device": "cpu",
"process_res": 518,
"process_res_method": "lower_bound_resize",
"ref_view_strategy": "middle",
},
)
with np.load(context.directory / "prediction.npz") as result:
self.assertEqual(
set(result.files),
{
"depth",
"conf",
"extrinsics",
"intrinsics",
"processed_images",
"is_metric",
},
)
np.testing.assert_array_equal(result["conf"], prediction.conf)
self.assertFalse(any(context.directory.glob("*.glb")))
self.assertFalse(any(context.directory.glob("*.png")))
self.assertEqual(calls[0][0], [str(input_path)])
self.assertEqual(calls[0][1]["process_res"], 518)
self.assertTrue(calls[0][1]["release_memory"])
def test_da3_per_frame_releases_only_after_final_frame(self):
releases = []
def prediction():
return SimpleNamespace(
depth=np.ones((1, 2, 2), dtype=np.float32),
conf=np.ones((1, 2, 2), dtype=np.float32),
extrinsics=np.eye(4, dtype=np.float32)[None, :3],
intrinsics=np.eye(3, dtype=np.float32)[None],
processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8),
)
fake_da3 = ModuleType("server.models.onnx_da3")
def infer(_session, _images, **options):
releases.append(options["release_memory"])
return prediction()
fake_da3.infer = infer
fake_runtime = ModuleType("server.models.onnx_runtime")
fake_runtime.provider_label = lambda _session: "CPU"
parameters = {
"model": "DA3",
"device": "cpu",
"process_res": 56,
"process_res_method": "upper_bound_resize",
"ref_view_strategy": "middle",
"inference_mode": "per_frame",
}
with patch.dict(
sys.modules,
{
"server.models.onnx_da3": fake_da3,
"server.models.onnx_runtime": fake_runtime,
},
):
self.da3_model.infer(
FakeModels(),
parameters,
[Path("a.png"), Path("b.png")],
)
self.assertEqual(releases, [False, True])
def test_moge2_debug_saves_each_direct_api_dictionary(self):
prediction = {
"points": np.ones((2, 2, 3), dtype=np.float32),
"depth": np.ones((2, 2), dtype=np.float32),
"normal": np.ones((2, 2, 3), dtype=np.float32),
"mask": np.ones((2, 2), dtype=bool),
"intrinsics": np.eye(3, dtype=np.float32),
}
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", "debug-moge2")
with patch.object(
self.debug_jobs,
"infer_moge2",
return_value=prediction,
) as infer_one:
self.debug_jobs.run_moge2(
context,
{
"input": str(frames),
"model": "MOGE2_VITS_NORMAL",
"video_frames": 2,
"device": "cpu",
"resolution_level": 5,
},
)
self.assertEqual(infer_one.call_count, 2)
self.assertEqual(
[call.kwargs["release_memory"] for call in infer_one.call_args_list],
[False, True],
)
archives = sorted(context.directory.glob("prediction.*.npz"))
self.assertEqual(len(archives), 2)
with np.load(archives[0]) as result:
self.assertEqual(set(result.files), set(prediction))
self.assertFalse(any(context.directory.glob("*.glb")))
def test_moge2_batch_releases_only_after_final_image(self):
prediction = {
"points": np.ones((2, 2, 3), dtype=np.float32),
"depth": np.ones((2, 2), dtype=np.float32),
"normal": np.ones((2, 2, 3), dtype=np.float32),
"mask": np.ones((2, 2), dtype=bool),
"intrinsics": np.eye(3, dtype=np.float32),
}
with patch.object(
self.moge2_model,
"infer_one",
return_value=prediction,
) as infer_one:
frames = self.moge2_model.infer(
FakeModels(),
{"model": "MOGE2_VITS_NORMAL", "device": "cpu"},
[Path("a.png"), Path("b.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],
)
def test_ben2_debug_saves_direct_api_foreground(self):
class Foreground:
def save(self, path):
Path(path).write_bytes(b"foreground")
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
input_path = root / "input.png"
input_path.touch()
context = FakeJobContext(root / "result", "debug-ben2")
with patch.object(
self.debug_jobs,
"infer_foreground",
return_value=(Foreground(), 1.0, 2.0),
) as infer_foreground:
self.debug_jobs.run_ben2(
context,
{
"input": str(input_path),
"model": "BEN2_BASE",
"video_frames": 1,
"device": "cpu",
},
)
infer_foreground.assert_called_once()
self.assertTrue(
infer_foreground.call_args.kwargs["release_memory"]
)
self.assertEqual(
(context.directory / "foreground.png").read_bytes(),
b"foreground",
)
def test_upscale_debug_saves_each_direct_api_prediction(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", "debug-upscale")
with patch.object(
self.debug_jobs,
"infer_upscale",
return_value=Image.new("RGB", (8, 8), (10, 20, 30)),
) as infer_one:
self.debug_jobs.run_upscale(
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],
)
predictions = sorted(context.directory.glob("prediction.*.png"))
self.assertEqual(len(predictions), 2)
with Image.open(predictions[0]) as prediction:
self.assertEqual(prediction.size, (8, 8))
self.assertEqual(prediction.mode, "RGB")
def test_upscale_job_saves_importable_upscale_sequence(self):
from PIL import Image
@@ -1030,6 +376,7 @@ class ServerJobsTest(unittest.TestCase):
context.logs,
)
def test_upscale_job_preserves_source_alpha(self):
from PIL import Image
@@ -1074,6 +421,7 @@ class ServerJobsTest(unittest.TestCase):
int(alpha[:, -1].min()),
)
def test_upscale_input_uses_the_global_long_edge_limit(self):
from PIL import Image
@@ -1087,68 +435,3 @@ class ServerJobsTest(unittest.TestCase):
)
with self.assertRaisesRegex(ValueError, "between 1 and 8192"):
self.resolution.limit_image(source, 9000)
def test_upscale_resource_error_releases_session_and_is_readable(self):
from PIL import Image
class Models:
def __init__(self, directory):
self.model_directory = directory
self.released = False
def directory(self, _model):
return self.model_directory
def get_upscale(self, _directory, _device):
return object()
def release_upscale(self):
self.released = True
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "input.png"
Image.new("RGB", (8, 8)).save(path)
models = Models(Path(directory))
with patch.object(
self.onnx_upscale,
"infer",
side_effect=self.onnx_upscale.UpscaleResourceError(
"provider memory exhausted"
),
):
with self.assertRaisesRegex(
RuntimeError,
"Reduce Maximum AI Input Size or select CPU",
):
self.upscale_model.infer_one(
models,
{
"model": "REALESRGAN_GENERAL_WDN_X4V3",
"device": "directml",
"max_input_size": 2048,
},
path,
)
self.assertTrue(models.released)
def test_da3_debug_writer_keeps_api_values_without_derived_attributes(self):
prediction = SimpleNamespace(
depth=np.ones((1, 2, 3), dtype=np.float32),
conf=np.arange(6, dtype=np.float32).reshape(1, 2, 3),
intrinsics=np.eye(3, dtype=np.float32)[None],
extrinsics=np.eye(4, dtype=np.float32)[None],
)
with tempfile.TemporaryDirectory() as directory:
output_path = self.debug_da3_outputs.write_prediction(
prediction,
Path(directory) / "prediction.npz",
)
with np.load(output_path) as attributes:
self.assertEqual(attributes["depth"].shape, (1, 2, 3))
self.assertEqual(attributes["conf"].shape, (1, 2, 3))
self.assertNotIn("confidence_percentile", attributes.files)
if __name__ == "__main__":
unittest.main()