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

Enhance device management and model handling in SnapForm

Browse Source 
- Updated frame extraction logic to allow for single frame extraction.
- Integrated model readiness checks for BEN2 and DA3 models in relevant operators.
- Added directory name attribute to model specifications for better organization.
- Improved device selection logic for production and debug tasks.
- Refactored preferences to expose device settings and ensure they reflect available hardware.
- Enhanced runtime device detection and management, including support for CUDA and MPS.
- Updated tests to cover new device management features and model directory handling.
- Removed deprecated download model tests to streamline codebase.
This commit is contained in:
马南
2026-08-08 20:16:11 +08:00
parent aee460fe7f
commit 9c896213a5
18 changed files with 699 additions and 196 deletions
Download Patch File Download Diff File Expand all files Collapse all files

211
tests/test_addon.py
View File

@@ -1356,6 +1356,104 @@ class AddonTest(unittest.TestCase):
["GIANT_1_1", "NESTED_GIANT_LARGE_1_1"],
)
def test_device_enums_list_only_supported_devices(self):
runtime = importlib.import_module("snapform.runtime")
with patch.object(
runtime,
"supported_devices",
return_value=["cuda", "cpu"],
):
items = runtime.device_enum_items(None, None)
self.assertEqual([item[0] for item in items], ["CUDA", "CPU"])
self.assertEqual(items[0][2], "NVIDIA CUDA")
def test_device_enums_default_to_the_fastest_supported_device(self):
for class_type in (
self.addon.SnapFormSettings,
self.addon.SnapFormPreferences,
):
with self.subTest(class_type=class_type.__name__):
self.assertEqual(
class_type.__annotations__["device"]["default"],
0,
)
def test_preferences_expose_production_device(self):
preferences = self.addon.SnapFormPreferences
self.assertIn("device", preferences.__annotations__)
items_callback = preferences.__annotations__["device"]["items"]
with patch.object(
self.addon.runtime,
"supported_devices",
return_value=["cpu"],
):
items = items_callback(None, None)
self.assertEqual([item[0] for item in items], ["CPU"])
def test_remove_background_poll_requires_ben2_model(self):
context = SimpleNamespace(
object=SimpleNamespace(
type="EMPTY",
empty_display_type="IMAGE",
data=object(),
)
)
with patch.object(
self.addon.runtime,
"host_busy",
return_value=False,
), patch.object(
self.addon.runtime,
"model_ready",
return_value=False,
):
self.assertFalse(self.addon.RemoveImageBackground.poll(context))
with patch.object(
self.addon.runtime,
"host_busy",
return_value=False,
), patch.object(
self.addon.runtime,
"model_ready",
return_value=True,
):
self.assertTrue(self.addon.RemoveImageBackground.poll(context))
def test_convert_to_depth_plane_poll_requires_da3_model(self):
context = SimpleNamespace(
object=SimpleNamespace(
type="EMPTY",
empty_display_type="IMAGE",
data=object(),
)
)
with patch.object(
self.addon.runtime,
"host_busy",
return_value=False,
), patch.object(
self.addon.runtime,
"model_ready",
side_effect=lambda key: key == "BASE",
):
self.assertTrue(
self.addon.ConvertImageToDepthPlane.poll(context)
)
with patch.object(
self.addon.runtime,
"host_busy",
return_value=False,
), patch.object(
self.addon.runtime,
"model_ready",
return_value=False,
):
self.assertFalse(
self.addon.ConvertImageToDepthPlane.poll(context)
)
def test_debug_panel_is_closed_by_default(self):
self.assertIn(
"DEFAULT_CLOSED",
@@ -1485,6 +1583,95 @@ class AddonTest(unittest.TestCase):
self.assertEqual(remove_background.call_args.args[1], input_path)
self.assertEqual(remove_background.call_args.args[4], "BEN2_BASE")
def test_debug_jobs_use_the_scene_device(self):
settings = SimpleNamespace(
background_input_path="input.png",
device="CPU",
)
operator = self.addon.RunJob()
operator.job_type = self.addon.operators.BACKGROUND_REMOVAL_JOB
with TemporaryDirectory() as directory:
root = Path(directory)
input_path = root / "input.png"
input_path.touch()
settings.background_input_path = str(input_path)
paths = {
"log": root / "job.log",
"status": root / "job.json",
"result": root / "result",
}
with patch.object(
self.addon.runtime,
"new_job_paths",
return_value=paths,
), patch.object(
self.addon.runtime,
"runtime_ready",
return_value=True,
), patch.object(
self.addon.runtime,
"model_ready",
return_value=True,
), patch.object(
self.addon.runtime,
"remove_background_request",
return_value=self.addon.runtime.HostJobRequest(
{"job_id": "test"}
),
) as remove_background:
operator._prepare_command(None, settings)
self.assertEqual(
remove_background.call_args.kwargs["device"],
"CPU",
)
def test_production_jobs_use_the_preferences_device(self):
settings = SimpleNamespace(device="CPU", video_frames=16)
operator = self.addon.RunJob()
operator.job_type = self.addon.operators.IMAGE_BACKGROUND_JOB
with TemporaryDirectory() as directory:
root = Path(directory)
input_path = root / "input.png"
input_path.touch()
operator.snapform_input_path = str(input_path)
paths = {
"log": root / "job.log",
"status": root / "job.json",
"result": root / "result",
}
with patch.object(
self.addon.runtime,
"new_job_paths",
return_value=paths,
), patch.object(
self.addon.runtime,
"runtime_ready",
return_value=True,
), patch.object(
self.addon.runtime,
"model_ready",
return_value=True,
), patch.object(
self.addon.runtime,
"production_device",
return_value="MPS",
), patch.object(
self.addon.runtime,
"remove_background_request",
return_value=self.addon.runtime.HostJobRequest(
{"job_id": "test"}
),
) as remove_background:
operator._prepare_command(None, settings)
self.assertEqual(
remove_background.call_args.kwargs["device"],
"MPS",
)
def test_image_background_job_uses_operator_input(self):
settings = SimpleNamespace(device="AUTO")
operator = self.addon.RunJob()
@@ -1804,8 +1991,8 @@ class AddonTest(unittest.TestCase):
return_value="C:/uv/uv.exe",
), patch.object(
runtime,
"preferences",
return_value=SimpleNamespace(torch_backend="CU128"),
"machine_devices",
return_value=["cuda", "cpu"],
), patch.object(
runtime,
"runtime_directory",
@@ -1821,11 +2008,15 @@ class AddonTest(unittest.TestCase):
self.assertEqual(command[0], sys.executable)
self.assertIn("--uv", command)
self.assertIn("C:/uv/uv.exe", command)
self.assertEqual(
command[command.index("--backend") + 1],
"cu128",
)
self.assertTrue(
any("bootstrap_runtime.py" in str(value) for value in command)
)
def test_install_runtime_command_passes_configured_pypi_index(self):
def test_install_runtime_command_auto_selects_backend(self):
runtime = importlib.import_module("snapform.runtime")
with patch.object(
runtime,
@@ -1833,11 +2024,8 @@ class AddonTest(unittest.TestCase):
return_value="C:/uv/uv.exe",
), patch.object(
runtime,
"preferences",
return_value=SimpleNamespace(
torch_backend="CU128",
pypi_index_url="https://pypi.tuna.tsinghua.edu.cn/simple",
),
"machine_devices",
return_value=["cpu"],
), patch.object(
runtime,
"runtime_directory",
@@ -1849,11 +2037,12 @@ class AddonTest(unittest.TestCase):
):
command = runtime.install_runtime_command(Path("status.json"))
index_position = command.index("--pypi-index")
self.assertNotIn("--pypi-index", command)
self.assertEqual(
command[index_position + 1],
"https://pypi.tuna.tsinghua.edu.cn/simple",
command[command.index("--backend") + 1],
"cpu",
)
self.assertIn("C:/uv/uv.exe", command)
def test_completed_export_opens_job_directory(self):
calls = []

2
tests/test_ai_host.py
View File

@@ -109,7 +109,7 @@ class HostTest(unittest.TestCase):
self.assertIs(first, second)
self.assertEqual(loads, ["model"])
self.assertEqual(cache.status()["da3"], "da3_model")
self.assertEqual(cache.status()["da3"], "model")
def test_host_tracks_clients_and_rejects_shared_shutdown(self):
fake_models = SimpleNamespace(

192
tests/test_bootstrap_runtime.py
View File

@@ -62,7 +62,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
) as mocked_run, patch.object(
bootstrap,
"run_install",
) as mocked_run_install:
) as mocked_run_install, patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
commands = [call.args[0] for call in mocked_run.call_args_list]
@@ -113,10 +117,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", host_command)
self.assertIn("setuptools", host_command)
self.assertNotIn("ninja", host_command)
check_command = next(
command
for command in commands
if any("import torch;" in str(value) for value in command)
check_command = bootstrap.verify_command(
bootstrap.environment_python(runtime_dir)
)
self.assertTrue(
any("import torch;" in str(value) for value in check_command)
)
self.assertTrue(
any(
@@ -130,6 +135,12 @@ class BootstrapRuntimeTest(unittest.TestCase):
self.assertFalse(
any("import gsplat" in str(value) for value in check_command)
)
self.assertTrue(
any(
"SNAPFORM_DEVICES=" in str(value)
for value in check_command
)
)
default_downloads = [
command
for command in commands
@@ -190,7 +201,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
), patch.object(
bootstrap,
"run_install",
) as mocked_run_install:
) as mocked_run_install, patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
commands = [
@@ -372,7 +387,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
bootstrap,
"run_stage",
side_effect=fake_run_stage,
), patch.object(bootstrap, "run"):
), patch.object(bootstrap, "run"), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
starts = [item[1] for item in stage_ranges]
@@ -604,34 +623,6 @@ class BootstrapRuntimeTest(unittest.TestCase):
"https://pypi.tuna.tsinghua.edu.cn/simple",
)
def test_child_environment_injects_registry_proxy(self):
bootstrap = load_bootstrap_runtime()
with patch.object(
bootstrap.urllib.request,
"getproxies",
return_value={
"http": "http://proxy:8080",
"https": "http://proxy:8080",
},
):
env = bootstrap.child_environment()
self.assertEqual(env["HTTP_PROXY"], "http://proxy:8080")
self.assertEqual(env["HTTPS_PROXY"], "http://proxy:8080")
def test_child_environment_keeps_explicit_proxy(self):
bootstrap = load_bootstrap_runtime()
with patch.object(
bootstrap.urllib.request,
"getproxies",
return_value={"https": "http://proxy:8080"},
):
env = bootstrap.child_environment(
{"HTTPS_PROXY": "http://explicit:3128"}
)
self.assertEqual(env["HTTPS_PROXY"], "http://explicit:3128")
def test_ready_json_records_installed_packages(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
@@ -657,6 +648,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
bootstrap,
"run_stage",
return_value=["torch==2.9.1+cu128", "numpy==1.26.4"],
), patch.object(
bootstrap,
"run_verify",
return_value={"cuda": True, "mps": False},
):
bootstrap.install(args)
@@ -669,6 +664,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
ready["packages"],
["numpy==1.26.4", "torch==2.9.1+cu128"],
)
self.assertEqual(ready["devices"], ["cuda", "cpu"])
def test_broken_runtime_is_cleaned_before_reinstall(self):
bootstrap = load_bootstrap_runtime()
@@ -693,6 +689,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
), patch.object(
bootstrap,
"run_install",
), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
@@ -728,6 +728,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
), patch.object(
bootstrap,
"run_install",
), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
@@ -758,7 +762,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
bootstrap,
"run",
side_effect=fake_run,
), patch.object(bootstrap, "run_install"):
), patch.object(bootstrap, "run_install"), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
self.assertTrue((runtime_dir / "snapform_ready.json").is_file())
@@ -872,6 +880,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
bootstrap,
"run_install",
side_effect=fake_run_install,
), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
@@ -918,6 +930,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
), patch.object(
bootstrap,
"run_install",
), patch.object(
bootstrap,
"run_verify",
return_value={},
):
bootstrap.install(args)
@@ -995,6 +1011,110 @@ class BootstrapRuntimeTest(unittest.TestCase):
self.assertIn(f"depth-anything-3 @ {da3_local}", ai_command)
self.assertIn(f"ben2 @ {ben2_local}", ai_command)
def test_install_stages_download_models_into_prefixed_directories(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
runtime_dir = root / "runtime"
runtime_dir.mkdir()
sources = root / "sources"
sources.mkdir()
(sources / "depth-anything-3.tar.gz").write_bytes(b"da3")
(sources / "ben2.tar.gz").write_bytes(b"ben2")
uv_path = root / "uv"
uv_path.write_bytes(b"uv")
args = SimpleNamespace(
uv=str(uv_path),
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
status=str(root / "status.json"),
)
stages = bootstrap.install_stages(
args,
runtime_dir,
root / "python",
uv_path,
)
downloads = {
label: command
for label, command, _environment, kind in stages
if kind == "download"
}
self.assertIn("da3_small", " ".join(map(str, downloads["da3 small"])))
self.assertIn("da3_base", " ".join(map(str, downloads["da3 base"])))
self.assertIn("ben2_base", " ".join(map(str, downloads["ben2 base"])))
def test_run_verify_returns_detected_devices(self):
bootstrap = load_bootstrap_runtime()
class FakeProcess:
def __init__(self, lines):
self.stdout = lines
def wait(self):
return 0
fake_process = FakeProcess(
['SNAPFORM_DEVICES={"cuda": false, "mps": true}']
)
with patch.object(
bootstrap.subprocess,
"Popen",
return_value=fake_process,
):
report = bootstrap.run_verify(
"python",
"status.json",
clock=lambda: 0.0,
)
self.assertFalse(report["cuda"])
self.assertTrue(report["mps"])
def test_run_verify_writes_heartbeat_progress(self):
bootstrap = load_bootstrap_runtime()
class FakeProcess:
def __init__(self, lines):
self.stdout = lines
def wait(self):
return 0
clock_state = {"now": 0.0}
def fake_clock():
clock_state["now"] += 6.0
return clock_state["now"]
def delayed_empty_lines():
time.sleep(0.6)
if False:
yield
fake_process = FakeProcess(delayed_empty_lines())
with patch.object(
bootstrap.subprocess,
"Popen",
return_value=fake_process,
), patch.object(bootstrap, "write_status") as mocked_write:
report = bootstrap.run_verify(
"python",
"status.json",
clock=fake_clock,
)
messages = [call.args[2] for call in mocked_write.call_args_list]
self.assertTrue(
any("verifying runtime" in message for message in messages)
)
self.assertEqual(report, {})
def test_install_uv_extracts_bundled_archive(self):
bootstrap = load_bootstrap_runtime()
with tempfile.TemporaryDirectory() as directory:

15
tests/test_image_depth.py
View File

@@ -209,6 +209,21 @@ class ImageDepthTest(unittest.TestCase):
(1, 1),
)
def test_extract_frame_range_honors_single_frame_setting(self):
source_object = SimpleNamespace(
data=SimpleNamespace(source="MOVIE", frame_duration=97),
image_user=SimpleNamespace(
frame_start=1,
frame_offset=70,
frame_duration=1,
),
)
self.assertEqual(
self.image_depth.extract_frame_range(source_object),
(71, 1),
)
def test_configure_image_user_preserves_cyclic_setting(self):
image_user = SimpleNamespace(
frame_start=0,

18
tests/test_models.py
View File

@@ -54,3 +54,21 @@ class ModelTest(unittest.TestCase):
("config.json", "model.safetensors"),
)
self.assertNotIn("BEN2_BASE", self.models.MODELS)
def test_da3_model_directories_share_the_da3_prefix(self):
for model_key, model in self.models.MODELS.items():
with self.subTest(model_key=model_key):
self.assertTrue(
model.directory_name.startswith("da3_"),
model.directory_name,
)
self.assertEqual(
self.models.model_directory_name(model_key),
model.directory_name,
)
def test_ben2_model_directory_uses_ben2_prefix(self):
self.assertEqual(
self.models.model_directory_name("BEN2_BASE"),
"ben2_base",
)

73
tests/test_worker.py
View File

@@ -385,79 +385,6 @@ class WorkerTest(unittest.TestCase):
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,
):
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"],
)
self.assertEqual(calls[0]["endpoint"], self.worker.HF_ENDPOINT)
def test_download_model_falls_back_to_modelscope_when_direct_fails(self):
calls = []
fake_hugging_face = ModuleType("huggingface_hub")
def fake_snapshot_download(**options):
calls.append(options)
raise OSError("Connection reset by peer")
fake_hugging_face.snapshot_download = fake_snapshot_download
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.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
self.worker,
"download_from_modelscope",
) as mocked_modelscope:
def fake_download_from_modelscope(*args, **kwargs):
model_dir = Path(args[1])
model_dir.mkdir(parents=True, exist_ok=True)
(model_dir / "config.json").write_text("{}")
(model_dir / "model.safetensors").write_bytes(b"x")
mocked_modelscope.side_effect = fake_download_from_modelscope
args.model_dir = str(Path(directory) / "model")
args.status = str(Path(directory) / "status.json")
self.worker.download_model(args)
self.assertEqual(len(calls), 1)
self.assertEqual(calls[0]["endpoint"], self.worker.HF_ENDPOINT)
self.assertEqual(mocked_modelscope.call_count, 1)
def test_download_model_skips_unreachable_endpoint(self):
calls = []
fake_hugging_face = ModuleType("huggingface_hub")