nan
d08d0e9c1a
- Modified status messages and stage labels in `test_addon.py` for clarity and consistency. - Added tests to ensure proper host termination and command execution in `test_addon.py`. - Enhanced `test_bootstrap_runtime.py` with additional tests for installation behavior, including handling of CUDA and MPS devices, and improved error handling for disk space and network checks. - Updated download progress reporting in `test_download_progress.py` to include download speed and improved message formatting. - Added new tests in `test_worker.py` for device selection, NPZ file validation, and model downloading behavior with Hugging Face endpoints. - Bumped version in `uv.lock` to 0.9.30.
506 lines
18 KiB
Python
506 lines
18 KiB
Python
import importlib.util
|
|
import inspect
|
|
import json
|
|
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_worker():
|
|
specification = importlib.util.spec_from_file_location(
|
|
"snapform_worker",
|
|
PROJECT_ROOT / "src" / "snapform" / "worker.py",
|
|
)
|
|
module = importlib.util.module_from_spec(specification)
|
|
specification.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class WorkerTest(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.worker = load_worker()
|
|
|
|
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.worker.collect_images(root)
|
|
|
|
self.assertEqual([path.name for path in images], ["a.jpg", "b.PNG"])
|
|
|
|
def test_reconstruct_reports_status_before_loading_torch(self):
|
|
source = inspect.getsource(self.worker.reconstruct)
|
|
|
|
self.assertLess(
|
|
source.index('"Starting AI runtime"'),
|
|
source.index("import torch"),
|
|
)
|
|
|
|
def test_ai_tasks_support_cooperative_cancellation(self):
|
|
with self.assertRaises(self.worker.TaskCancelled):
|
|
self.worker.check_cancelled(lambda: True)
|
|
|
|
self.worker.check_cancelled(lambda: False)
|
|
self.assertIn(
|
|
"cancel_check=cancel_check",
|
|
inspect.getsource(self.worker.reconstruct),
|
|
)
|
|
self.assertIn(
|
|
"cancel_check=cancel_check",
|
|
inspect.getsource(self.worker.remove_background),
|
|
)
|
|
|
|
def test_select_device_supports_mps_and_auto_fallback(self):
|
|
mps_torch = SimpleNamespace(
|
|
cuda=SimpleNamespace(is_available=lambda: False),
|
|
backends=SimpleNamespace(
|
|
mps=SimpleNamespace(is_available=lambda: True),
|
|
),
|
|
)
|
|
mps_torch.device = lambda device_type: SimpleNamespace(
|
|
type=device_type
|
|
)
|
|
|
|
self.assertEqual(
|
|
self.worker.select_device("mps", mps_torch).type,
|
|
"mps",
|
|
)
|
|
self.assertEqual(
|
|
self.worker.select_device("auto", mps_torch).type,
|
|
"mps",
|
|
)
|
|
self.assertEqual(
|
|
self.worker.select_device("cpu", mps_torch).type,
|
|
"cpu",
|
|
)
|
|
|
|
cuda_torch = SimpleNamespace(
|
|
cuda=SimpleNamespace(is_available=lambda: True),
|
|
)
|
|
cuda_torch.device = lambda device_type: SimpleNamespace(
|
|
type=device_type
|
|
)
|
|
self.assertEqual(
|
|
self.worker.select_device("auto", cuda_torch).type,
|
|
"cuda",
|
|
)
|
|
|
|
cpu_only_torch = SimpleNamespace(
|
|
cuda=SimpleNamespace(is_available=lambda: False),
|
|
backends=SimpleNamespace(
|
|
mps=SimpleNamespace(is_available=lambda: False),
|
|
),
|
|
)
|
|
cpu_only_torch.device = lambda device_type: SimpleNamespace(
|
|
type=device_type
|
|
)
|
|
self.assertEqual(
|
|
self.worker.select_device("auto", cpu_only_torch).type,
|
|
"cpu",
|
|
)
|
|
with self.assertRaisesRegex(RuntimeError, "MPS was requested"):
|
|
self.worker.select_device("mps", cpu_only_torch)
|
|
|
|
def test_optional_background_removal_runs_before_da3_model_load(self):
|
|
source = inspect.getsource(self.worker.reconstruct)
|
|
|
|
self.assertLess(
|
|
source.index("infer_foreground("),
|
|
source.index("DepthAnything3.from_pretrained"),
|
|
)
|
|
args = self.worker.parser().parse_args(
|
|
[
|
|
"reconstruct",
|
|
"--input",
|
|
"input.png",
|
|
"--model-dir",
|
|
"da3",
|
|
"--output-dir",
|
|
"result",
|
|
"--export-format",
|
|
"mini_npz",
|
|
"--status",
|
|
"status.json",
|
|
"--remove-background-model-dir",
|
|
"ben2",
|
|
"--process-res",
|
|
"1024",
|
|
"--process-res-method",
|
|
"upper_bound_resize",
|
|
"--conf-thresh-percentile",
|
|
"40",
|
|
"--num-max-points",
|
|
"1000",
|
|
"--show-cameras",
|
|
"0",
|
|
"--gs-views-interval",
|
|
"1",
|
|
"--video-frames",
|
|
"16",
|
|
"--device",
|
|
"auto",
|
|
"--use-ray-pose",
|
|
"1",
|
|
"--ref-view-strategy",
|
|
"first",
|
|
]
|
|
)
|
|
self.assertEqual(args.remove_background_model_dir, "ben2")
|
|
|
|
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.worker.prepare_input(
|
|
image_path,
|
|
16,
|
|
directory,
|
|
)
|
|
|
|
self.assertEqual(images, [image_path])
|
|
self.assertIsNone(temporary_directory)
|
|
|
|
def test_npz_is_complete_requires_readable_archive(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
valid_path = root / "valid.npz"
|
|
np.savez_compressed(
|
|
valid_path,
|
|
depth=np.zeros((4, 4)),
|
|
)
|
|
self.assertTrue(self.worker.npz_is_complete(valid_path))
|
|
|
|
empty_path = root / "empty.npz"
|
|
empty_path.write_bytes(b"PK\x05\x06" + b"\x00" * 18)
|
|
self.assertFalse(self.worker.npz_is_complete(empty_path))
|
|
|
|
def test_wait_for_exported_paths_waits_for_complete_npz(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
npz_path = root / "exports" / "mini_npz" / "results.npz"
|
|
npz_path.parent.mkdir(parents=True)
|
|
npz_path.touch()
|
|
with patch.object(
|
|
self.worker,
|
|
"npz_is_complete",
|
|
side_effect=[False, True],
|
|
):
|
|
result = self.worker.wait_for_exported_paths(
|
|
root,
|
|
"mini_npz",
|
|
timeout=2.0,
|
|
)
|
|
self.assertEqual(result, [npz_path])
|
|
|
|
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.worker.prepare_input(input_path, 16, directory)
|
|
|
|
def test_download_passes_optional_file_patterns(self):
|
|
calls = []
|
|
fake_hugging_face = ModuleType("huggingface_hub")
|
|
fake_hugging_face.snapshot_download = lambda **options: calls.append(
|
|
options
|
|
)
|
|
args = SimpleNamespace(
|
|
repository="PramaLLC/BEN2",
|
|
model_dir="model",
|
|
status="status.json",
|
|
allow_pattern=["config.json", "model.safetensors"],
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
|
sys.modules,
|
|
{"huggingface_hub": fake_hugging_face},
|
|
), patch.object(
|
|
self.worker,
|
|
"create_download_progress_class",
|
|
return_value=object,
|
|
), patch.object(
|
|
self.worker,
|
|
"resolve_hf_endpoint",
|
|
return_value=(self.worker.HF_ENDPOINT, False),
|
|
):
|
|
args.model_dir = str(Path(directory) / "model")
|
|
args.status = str(Path(directory) / "status.json")
|
|
self.worker.download_model(args)
|
|
|
|
self.assertEqual(
|
|
calls[0]["allow_patterns"],
|
|
["config.json", "model.safetensors"],
|
|
)
|
|
|
|
def test_download_model_switches_to_mirror_before_hf_import(self):
|
|
calls = []
|
|
fake_hugging_face = ModuleType("huggingface_hub")
|
|
fake_hugging_face.snapshot_download = lambda **options: calls.append(
|
|
options
|
|
)
|
|
args = SimpleNamespace(
|
|
repository="PramaLLC/BEN2",
|
|
model_dir="model",
|
|
status="status.json",
|
|
allow_pattern=[],
|
|
)
|
|
|
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
|
sys.modules,
|
|
{"huggingface_hub": fake_hugging_face},
|
|
), patch.object(
|
|
self.worker,
|
|
"create_download_progress_class",
|
|
return_value=object,
|
|
), patch.object(
|
|
self.worker,
|
|
"resolve_hf_endpoint",
|
|
return_value=("https://hf-mirror.com", True),
|
|
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
|
self.worker,
|
|
"write_status",
|
|
) as mocked_write:
|
|
args.model_dir = str(Path(directory) / "model")
|
|
args.status = str(Path(directory) / "status.json")
|
|
self.worker.download_model(args)
|
|
self.assertEqual(
|
|
os.environ["HF_ENDPOINT"],
|
|
"https://hf-mirror.com",
|
|
)
|
|
messages = [call.args[2] for call in mocked_write.call_args_list]
|
|
self.assertTrue(any("(mirror)" in message for message in messages))
|
|
|
|
def test_ben2_background_removal_writes_fast_transparent_png(self):
|
|
inference_options = []
|
|
save_options = []
|
|
|
|
class FakeImage:
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
return None
|
|
|
|
def convert(self, mode):
|
|
self.mode = mode
|
|
return self
|
|
|
|
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("ben2")
|
|
fake_ben2.BEN_Base = FakeModel
|
|
fake_pil = ModuleType("PIL")
|
|
fake_pil.Image = SimpleNamespace(open=lambda _path: FakeImage())
|
|
|
|
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
|
sys.modules,
|
|
{
|
|
"torch": fake_torch,
|
|
"ben2": fake_ben2,
|
|
"PIL": fake_pil,
|
|
},
|
|
):
|
|
root = Path(directory)
|
|
input_path = root / "input.png"
|
|
input_path.touch()
|
|
args = SimpleNamespace(
|
|
input=str(input_path),
|
|
model_dir="model",
|
|
output_dir=str(root / "result"),
|
|
status=str(root / "status.json"),
|
|
device="cpu",
|
|
)
|
|
self.worker.remove_background(args)
|
|
status = json.loads(Path(args.status).read_text(encoding="utf-8"))
|
|
|
|
self.assertTrue((root / "result" / "foreground.png").is_file())
|
|
|
|
self.assertEqual(
|
|
inference_options[0][1],
|
|
{"refine_foreground": False},
|
|
)
|
|
self.assertEqual(
|
|
save_options[0],
|
|
{"format": "PNG", "compress_level": 1},
|
|
)
|
|
self.assertIn("Background removed in", status["message"])
|
|
|
|
def test_uses_official_da3_export_parameters_for_all_four_formats(self):
|
|
inference_options = []
|
|
export_calls = []
|
|
|
|
class FakeModel:
|
|
@classmethod
|
|
def from_pretrained(cls, model_dir):
|
|
if model_dir != "model":
|
|
raise AssertionError(f"Unexpected model directory: {model_dir}")
|
|
return cls()
|
|
|
|
def to(self, device):
|
|
if device.type != "cpu":
|
|
raise AssertionError(f"Unexpected device: {device.type}")
|
|
return self
|
|
|
|
def eval(self):
|
|
return None
|
|
|
|
def inference(self, images, **options):
|
|
inference_options.append((images, options))
|
|
return SimpleNamespace()
|
|
|
|
def fake_export(prediction, export_format, export_dir, **options):
|
|
export_calls.append((prediction, export_format, export_dir, options))
|
|
output_directory = Path(export_dir)
|
|
if export_format == "glb":
|
|
(output_directory / "scene.glb").touch()
|
|
elif export_format == "gs_ply":
|
|
ply_directory = output_directory / "gs_ply"
|
|
ply_directory.mkdir()
|
|
(ply_directory / "0000.ply").touch()
|
|
elif export_format == "depth_vis":
|
|
depth_directory = output_directory / "depth_vis"
|
|
depth_directory.mkdir()
|
|
(depth_directory / "0000.jpg").touch()
|
|
else:
|
|
mini_directory = output_directory / "exports" / "mini_npz"
|
|
mini_directory.mkdir(parents=True)
|
|
np.savez_compressed(
|
|
mini_directory / "results.npz",
|
|
depth=np.zeros((14, 14)),
|
|
)
|
|
|
|
fake_torch = ModuleType("torch")
|
|
fake_torch.device = lambda device_type: SimpleNamespace(
|
|
type=device_type,
|
|
)
|
|
fake_torch.inference_mode = nullcontext
|
|
fake_api = ModuleType("depth_anything_3.api")
|
|
fake_api.DepthAnything3 = FakeModel
|
|
fake_export_module = ModuleType("depth_anything_3.utils.export")
|
|
fake_export_module.export = fake_export
|
|
fake_utils = ModuleType("depth_anything_3.utils")
|
|
fake_utils.__path__ = []
|
|
fake_package = ModuleType("depth_anything_3")
|
|
fake_package.__path__ = []
|
|
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
input_path = root / "input.png"
|
|
input_path.touch()
|
|
with patch.dict(
|
|
sys.modules,
|
|
{
|
|
"torch": fake_torch,
|
|
"depth_anything_3": fake_package,
|
|
"depth_anything_3.api": fake_api,
|
|
"depth_anything_3.utils": fake_utils,
|
|
"depth_anything_3.utils.export": fake_export_module,
|
|
},
|
|
):
|
|
for export_format in self.worker.EXPORT_FORMATS:
|
|
args = SimpleNamespace(
|
|
input=str(input_path),
|
|
video_frames=16,
|
|
output_dir=str(root / export_format),
|
|
export_format=export_format,
|
|
status=str(root / f"{export_format}.json"),
|
|
device="cpu",
|
|
model_dir="model",
|
|
process_res=518,
|
|
process_res_method="lower_bound_resize",
|
|
use_ray_pose=1,
|
|
ref_view_strategy="middle",
|
|
conf_thresh_percentile=35.0,
|
|
num_max_points=12345,
|
|
show_cameras=0,
|
|
gs_views_interval=2,
|
|
)
|
|
self.worker.reconstruct(args)
|
|
status = json.loads(
|
|
Path(args.status).read_text(encoding="utf-8")
|
|
)
|
|
self.assertEqual(status["progress"], 1.0)
|
|
self.assertIn(export_format, status["message"])
|
|
|
|
self.assertEqual(len(inference_options), 4)
|
|
self.assertEqual(len(export_calls), 4)
|
|
for images, options in inference_options:
|
|
self.assertEqual(images, [str(input_path)])
|
|
self.assertEqual(options["process_res"], 518)
|
|
self.assertEqual(
|
|
options["process_res_method"],
|
|
"lower_bound_resize",
|
|
)
|
|
self.assertTrue(options["use_ray_pose"])
|
|
self.assertEqual(options["ref_view_strategy"], "middle")
|
|
for _prediction, export_format, _export_dir, options in export_calls:
|
|
if export_format == "glb":
|
|
self.assertEqual(
|
|
options,
|
|
{
|
|
"glb": {
|
|
"conf_thresh_percentile": 35.0,
|
|
"num_max_points": 12345,
|
|
"show_cameras": False,
|
|
"export_depth_vis": False,
|
|
}
|
|
},
|
|
)
|
|
elif export_format == "gs_ply":
|
|
self.assertEqual(
|
|
options,
|
|
{"gs_ply": {"gs_views_interval": 2}},
|
|
)
|
|
else:
|
|
self.assertEqual(options, {})
|
|
|
|
infer_gs_values = [options["infer_gs"] for _, options in inference_options]
|
|
self.assertEqual(infer_gs_values, [False, True, False, False])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|