Enhance device management and model handling in SnapForm
- 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:
@@ -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 = []
|
||||
|
||||
Reference in New Issue
Block a user