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

feat: Update CoreML device support and server controls

Browse Source 
- Replace 'mps' with 'coreml' as the identifier for Apple CoreML Execution Provider across the codebase.
- Remove Apple chip generation checks and ensure CoreML is available on all macOS arm64 devices.
- Adjust ONNX Runtime session creation to use default CoreML configurations without forcing 'MLProgram'.
- Implement support for MoGe-2 to utilize CoreML with static input shape requirements.
- Enhance server panel to display restart and stop controls consistently across running states, including 'BUSY'.
- Update UI tests to verify new behavior for server controls and cached model display.
- Bump blendjob dependency version to 0.1.16 in packaging tests and configurations.
This commit is contained in:
马南
2026-09-02 10:50:02 +08:00
parent 7f6e47432e
commit 5a011b3b57
24 changed files with 619 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files

192
tests/test_blender_addon.py
View File

@@ -17,10 +17,20 @@ class BlenderAddonTest(unittest.TestCase):
def setUpClass(cls):
cls.remove_modules()
def draw_status_bar(_owner, _context):
pass
draw_status_bar._draw_funcs = [draw_status_bar]
def append_status_bar(draw):
draw._owner = "anyimage"
draw_status_bar._draw_funcs.append(draw)
cls.status_bar_draws.append(draw)
def remove_status_bar(draw):
draw_status_bar._draw_funcs.remove(draw)
cls.status_bar_draws.remove(draw)
fake_bpy = ModuleType("bpy")
fake_bpy.props = SimpleNamespace(
BoolProperty=lambda **options: options,
@@ -41,8 +51,9 @@ class BlenderAddonTest(unittest.TestCase):
Image=type("Image", (), {}),
Scene=type("Scene", (), {}),
STATUSBAR_HT_header=SimpleNamespace(
draw=draw_status_bar,
append=append_status_bar,
remove=lambda draw: cls.status_bar_draws.remove(draw),
remove=remove_status_bar,
),
VIEW3D_MT_object_context_menu=SimpleNamespace(
append=lambda draw: cls.context_menu_draws.append(draw),
@@ -158,26 +169,55 @@ class BlenderAddonTest(unittest.TestCase):
if hasattr(self.fake_bpy.types.Image, "anyimage_upscale_source"):
del self.fake_bpy.types.Image.anyimage_upscale_source
def test_mps_requires_m4_or_newer(self):
def test_all_apple_silicon_supports_coreml(self):
properties = self.anyimage.properties
parse = properties.parse_apple_silicon_generation
self.assertEqual(parse("Apple M2"), 2)
self.assertEqual(parse("Apple M4 Pro"), 4)
self.assertEqual(parse("Apple M10 Max"), 10)
self.assertIsNone(parse("Intel(R) Core(TM) i9"))
self.assertFalse(properties.supports_mps(3))
self.assertTrue(properties.supports_mps(4))
for generation in range(1, 5):
with (
self.subTest(generation=generation),
patch.object(
self.anyimage.runtime,
"environment_manifest",
return_value={},
),
patch.object(properties.sys, "platform", "darwin"),
patch.object(
properties.platform,
"machine",
return_value="arm64",
),
):
self.assertEqual(
properties.supported_devices(),
["cpu", "coreml"],
)
def test_unsupported_mac_hides_manifest_mps_device(self):
def test_non_apple_silicon_mac_uses_cpu(self):
properties = self.anyimage.properties
manifest = {"devices": ["mps", "cpu"]}
with (
patch.object(
self.anyimage.runtime,
"environment_manifest",
return_value={},
),
patch.object(properties.sys, "platform", "darwin"),
patch.object(
properties.platform,
"machine",
return_value="x86_64",
),
):
self.assertEqual(properties.supported_devices(), ["cpu"])
def test_manifest_coreml_requires_apple_silicon(self):
properties = self.anyimage.properties
manifest = {"devices": ["coreml", "cpu"]}
with (
patch.object(
self.anyimage.runtime,
"environment_manifest",
return_value=manifest,
),
patch.object(properties, "supports_mps", return_value=False),
patch.object(properties.sys, "platform", "linux"),
):
self.assertEqual(properties.supported_devices(), ["cpu"])
with (
@@ -186,9 +226,34 @@ class BlenderAddonTest(unittest.TestCase):
"environment_manifest",
return_value=manifest,
),
patch.object(properties, "supports_mps", return_value=True),
patch.object(properties.sys, "platform", "darwin"),
patch.object(
properties.platform,
"machine",
return_value="aarch64",
),
):
self.assertEqual(properties.supported_devices(), ["mps", "cpu"])
self.assertEqual(
properties.supported_devices(),
["cpu", "coreml"],
)
def test_manifest_rejects_old_mps_device(self):
properties = self.anyimage.properties
with (
patch.object(
self.anyimage.runtime,
"environment_manifest",
return_value={"devices": ["mps", "cpu"]},
),
patch.object(properties.sys, "platform", "darwin"),
patch.object(
properties.platform,
"machine",
return_value="arm64",
),
):
self.assertEqual(properties.supported_devices(), ["cpu"])
def test_windows_keeps_default_directml_device(self):
properties = self.anyimage.properties
@@ -2957,6 +3022,105 @@ class BlenderAddonTest(unittest.TestCase):
offenders.append(path.relative_to(source_root))
self.assertEqual(offenders, [])
def _draw_server_panel(self, state, cached_names=()):
events = SimpleNamespace(labels=[], operators=[])
class Layout:
def __init__(self, kind="layout"):
self.kind = kind
self.enabled = True
def row(self, **_options):
return Layout("row")
def column(self, **_options):
return Layout("column")
def label(self, **options):
events.labels.append(
(self.kind, options.get("text"), options.get("icon"))
)
def operator(self, identifier, **_options):
events.operators.append((identifier, self.enabled))
return SimpleNamespace()
panel_module = importlib.import_module("anyimage.panel")
panel = panel_module.ServerPanel()
panel.layout = Layout()
loaded = {
str(index): name
for index, name in enumerate(cached_names)
}
with (
patch.object(
panel_module.runtime,
"environment_ready",
return_value=True,
),
patch.object(
panel_module.runtime,
"server_busy",
return_value=state == "BUSY",
),
patch.object(
panel_module.runtime,
"server_status",
return_value={
"state": state,
"message": state.title(),
"resources": {
"model_manager": {"loaded": loaded},
},
},
),
):
panel.draw(None)
return events
def test_server_panel_shows_controls_for_every_running_state(self):
for state, label, icon in (
("READY", "Running", "RADIOBUT_ON"),
("BUSY", "Running (Busy)", "TIME"),
):
with self.subTest(state=state):
events = self._draw_server_panel(state)
operators = [identifier for identifier, _enabled in events.operators]
self.assertIn(("row", label, icon), events.labels)
self.assertIn("anyimage.restart_server", operators)
self.assertIn("anyimage.stop_server", operators)
for state in ("STOPPED", "STARTING", "ERROR"):
with self.subTest(state=state):
events = self._draw_server_panel(state)
operators = [identifier for identifier, _enabled in events.operators]
self.assertNotIn("anyimage.restart_server", operators)
self.assertNotIn("anyimage.stop_server", operators)
def test_server_panel_wraps_multiple_cached_models(self):
for cached_names, expected_labels in (
((), [("row", "Cached: None")]),
(("Model A",), [("row", "Cached: Model A")]),
(
("Model A", "Model B"),
[
("column", "Cached:"),
("column", "Model A"),
("column", "Model B"),
],
),
):
with self.subTest(cached_names=cached_names):
events = self._draw_server_panel("READY", cached_names)
cached_labels = [
(kind, text)
for kind, text, _icon in events.labels
if text != "Running"
]
self.assertEqual(cached_labels, expected_labels)
operators = [identifier for identifier, _enabled in events.operators]
self.assertIn("anyimage.clear_models", operators)
def test_status_bar_reads_runtime_state(self):
calls = []

12
tests/test_extension_packaging.py
View File

@@ -24,7 +24,7 @@ class ExtensionPackagingTest(unittest.TestCase):
pack = load_pack()
with tempfile.TemporaryDirectory() as directory:
directory_path = Path(directory)
wheel_path = directory_path / "blendjob-0.1.14-py3-none-any.whl"
wheel_path = directory_path / "blendjob-0.1.16-py3-none-any.whl"
wheel_path.write_bytes(b"wheel")
original_dist = pack.DIST_DIR
pack.DIST_DIR = directory_path / "dist"
@@ -132,7 +132,7 @@ class ExtensionPackagingTest(unittest.TestCase):
self.assertIn("server/models/upscale.py", names)
self.assertIn("server/__init__.py", names)
self.assertIn(
"wheels/blendjob-0.1.14-py3-none-any.whl",
"wheels/blendjob-0.1.16-py3-none-any.whl",
names,
)
self.assertFalse(any(name.startswith("blendjob/") for name in names))
@@ -147,7 +147,7 @@ class ExtensionPackagingTest(unittest.TestCase):
def test_reads_extension_dependencies_from_pyproject(self):
pack = load_pack()
self.assertEqual(pack.project_dependencies(), ["blendjob==0.1.14"])
self.assertEqual(pack.project_dependencies(), ["blendjob==0.1.16"])
def test_downloads_declared_wheels_for_target_platform(self):
pack = load_pack()
@@ -156,20 +156,20 @@ class ExtensionPackagingTest(unittest.TestCase):
def fake_run(command, check):
self.assertTrue(check)
(destination / "blendjob-0.1.14-py3-none-any.whl").write_bytes(b"wheel")
(destination / "blendjob-0.1.16-py3-none-any.whl").write_bytes(b"wheel")
with patch.object(pack.subprocess, "run", side_effect=fake_run) as run:
wheels = pack.download_wheels("linux-x64", destination)
command = run.call_args.args[0]
self.assertIn("blendjob==0.1.14", command)
self.assertIn("blendjob==0.1.16", command)
self.assertIn("--find-links", command)
self.assertIn(str(pack.SOURCE_DIR / "wheels"), command)
self.assertIn("--python-version=3.11", command)
self.assertIn("--platform=manylinux2014_x86_64", command)
self.assertEqual(
[wheel.name for wheel in wheels],
["blendjob-0.1.14-py3-none-any.whl"],
["blendjob-0.1.16-py3-none-any.whl"],
)
def test_rejects_unknown_platform(self):

27
tests/test_model_moge2_onnx.py
View File

@@ -79,7 +79,7 @@ class Moge2OnnxTest(unittest.TestCase):
["DmlExecutionProvider", "CPUExecutionProvider"],
)
def test_coreml_session_uses_mlprogram(self):
def test_coreml_session_requires_static_input_partitions(self):
calls = []
fake_runtime = ModuleType("onnxruntime")
fake_runtime.get_available_providers = lambda: [
@@ -91,20 +91,21 @@ class Moge2OnnxTest(unittest.TestCase):
)
with patch.dict(sys.modules, {"onnxruntime": fake_runtime}):
self.moge.create_runtime_session("model.onnx")
self.moge.create_session("moge-model", "coreml")
self.assertEqual(calls[0][0], "model.onnx")
self.assertEqual(
calls[0][1][0],
(
"CoreMLExecutionProvider",
{
"ModelFormat": "MLProgram",
"MLComputeUnits": "ALL",
"RequireStaticInputShapes": "0",
"EnableOnSubgraphs": "0",
},
),
calls[0][0],
str(Path("moge-model") / "model.onnx"),
)
self.assertEqual(
calls[0][1],
[
(
"CoreMLExecutionProvider",
{"RequireStaticInputShapes": "1"},
),
"CPUExecutionProvider",
],
)
def test_explicit_unavailable_provider_falls_back_to_cpu(self):

49
tests/test_model_onnx_runtime.py
View File

@@ -61,7 +61,7 @@ class OnnxRuntimeTest(unittest.TestCase):
),
)
def test_coreml_and_cpu_session_options_are_preserved(self):
def test_coreml_and_cpu_use_plain_session_providers(self):
runtime, calls = self.create_runtime(
["CoreMLExecutionProvider", "CPUExecutionProvider"]
)
@@ -69,17 +69,62 @@ class OnnxRuntimeTest(unittest.TestCase):
with patch.dict(sys.modules, {"onnxruntime": runtime}):
onnx_runtime.create_session("model.onnx", "auto")
self.assertEqual(
calls[0][1],
[
"CoreMLExecutionProvider",
"CPUExecutionProvider",
],
)
def test_model_options_only_configure_coreml_provider(self):
runtime, calls = self.create_runtime(
["CoreMLExecutionProvider", "CPUExecutionProvider"]
)
with patch.dict(sys.modules, {"onnxruntime": runtime}):
onnx_runtime.create_session(
"model.onnx",
"coreml",
coreml_options={"RequireStaticInputShapes": "1"},
)
self.assertEqual(
calls[0][1],
[
(
"CoreMLExecutionProvider",
dict(onnx_runtime.COREML_PROVIDER_OPTIONS),
{"RequireStaticInputShapes": "1"},
),
"CPUExecutionProvider",
],
)
def test_coreml_device_selects_coreml_provider(self):
available = [
"CoreMLExecutionProvider",
"CPUExecutionProvider",
]
self.assertEqual(
onnx_runtime.select_providers("coreml", available),
["CoreMLExecutionProvider", "CPUExecutionProvider"],
)
self.assertEqual(
onnx_runtime.select_providers(
"coreml",
["CPUExecutionProvider"],
),
["CPUExecutionProvider"],
)
def test_old_mps_device_is_rejected(self):
with self.assertRaises(KeyError):
onnx_runtime.select_providers(
"mps",
["CoreMLExecutionProvider", "CPUExecutionProvider"],
)
def test_completing_cuda_run_shrinks_memory_arena(self):
configs = []