nan
2ff7b04f56
- Introduced a new memory management policy for ONNX Runtime to release temporary GPU memory after inference completion. - Established a shared inference entry point for all ONNX models, differentiating between long-lived sessions and temporary memory arenas. - Updated CUDA sessions to utilize a shrinkable memory arena, reclaiming GPU memory post-inference. - Ensured consistent memory release semantics across BEN2, DA3, MoGe-2, and Upscale models while reusing existing sessions. - Added lifecycle tests to validate temporary memory release and session reuse. - Modified model adapters to integrate the new memory management features, ensuring backward compatibility with existing inference behavior. - Enhanced unit tests to cover new memory release scenarios and verify correct behavior across different execution providers.
1109 lines
38 KiB
Python
1109 lines
38 KiB
Python
import os
|
|
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
|
|
|
|
|
|
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
|
|
|
|
|
|
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.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(
|
|
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_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 = []
|
|
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_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
|
|
|
|
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, (20, 12))
|
|
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)
|
|
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()
|