From 576bb2f6f2eacbfb97fb7fddb41b18ed04deb069 Mon Sep 17 00:00:00 2001 From: Nik Nichols Date: Wed, 24 Jun 2026 16:30:29 -0500 Subject: [PATCH] test(gpu): add tests for Apple Silicon MPS device resolution and fallback --- tests/tools/test_mps_device.py | 337 +++++++++++++++++++++++++++++++++ 1 file changed, 337 insertions(+) create mode 100644 tests/tools/test_mps_device.py diff --git a/tests/tools/test_mps_device.py b/tests/tools/test_mps_device.py new file mode 100644 index 00000000..844d633e --- /dev/null +++ b/tests/tools/test_mps_device.py @@ -0,0 +1,337 @@ +"""Tests for Apple Silicon MPS / device resolution helper. + +These tests mock at sys.modules level so they run without torch installed. +""" + +from __future__ import annotations + +import sys +from unittest.mock import MagicMock + +import pytest + + +@pytest.fixture() +def mock_torch(): + """Inject a fake torch module into sys.modules for the duration of a test.""" + fake = MagicMock() + # Restore whatever was there before (real torch, or nothing) + previous = sys.modules.get("torch", None) + sys.modules["torch"] = fake + yield fake + if previous is None: + sys.modules.pop("torch", None) + else: + sys.modules["torch"] = previous + + +# ------------------------------------------------------------------ +# get_torch_device — basic routing +# ------------------------------------------------------------------ + +def test_get_torch_device_returns_mps_when_cuda_absent(mock_torch): + """mps is used when cuda is unavailable but mps is available.""" + mock_torch.cuda.is_available.return_value = False + mock_torch.backends.mps.is_built.return_value = True + mock_torch.backends.mps.is_available.return_value = True + + from tools.video._shared import get_torch_device + assert get_torch_device() == "mps" + + +def test_get_torch_device_returns_cpu_as_fallback(mock_torch): + """cpu is the final fallback when neither cuda nor mps is available.""" + mock_torch.cuda.is_available.return_value = False + mock_torch.backends.mps.is_built.return_value = True + mock_torch.backends.mps.is_available.return_value = False + + from tools.video._shared import get_torch_device + assert get_torch_device() == "cpu" + + +def test_get_torch_device_returns_cuda_when_available(mock_torch): + """cuda takes priority over mps when both are present.""" + mock_torch.cuda.is_available.return_value = True + mock_torch.backends.mps.is_built.return_value = True + mock_torch.backends.mps.is_available.return_value = True + + from tools.video._shared import get_torch_device + assert get_torch_device() == "cuda" + + +def test_get_torch_device_returns_cpu_when_torch_not_installed(): + """cpu is returned safely when torch cannot be imported.""" + previous = sys.modules.pop("torch", None) + try: + # Make torch unimportable + sys.modules["torch"] = None # type: ignore[assignment] + from tools.video._shared import get_torch_device + assert get_torch_device() == "cpu" + finally: + if previous is None: + sys.modules.pop("torch", None) + else: + sys.modules["torch"] = previous + + +# ------------------------------------------------------------------ +# get_torch_device — MPS guard (torch.backends.mps missing) +# ------------------------------------------------------------------ + +def test_get_torch_device_cpu_when_mps_backend_missing(mock_torch): + """Falls back to cpu when torch.backends.mps does not exist (e.g. Linux wheels).""" + mock_torch.cuda.is_available.return_value = False + # Simulate a torch build without mps backend + del mock_torch.backends.mps + + from tools.video._shared import get_torch_device + assert get_torch_device() == "cpu" + + +def test_get_torch_device_cpu_when_mps_not_built(mock_torch): + """Falls back to cpu when MPS is not built into torch.""" + mock_torch.cuda.is_available.return_value = False + mock_torch.backends.mps.is_built.return_value = False + mock_torch.backends.mps.is_available.return_value = False + + from tools.video._shared import get_torch_device + assert get_torch_device() == "cpu" + + +# ------------------------------------------------------------------ +# load_diffusers_pipeline — device and dtype routing +# ------------------------------------------------------------------ + +def _make_pipeline_mocks(monkeypatch, *, cuda=False, mps=False, bf16=False): + """Helper: inject fake torch + diffusers and return the pipeline mock.""" + import importlib + + fake_torch = MagicMock() + fake_torch.cuda.is_available.return_value = cuda + fake_torch.cuda.is_bf16_supported.return_value = bf16 + fake_torch.backends.mps.is_built.return_value = mps + fake_torch.backends.mps.is_available.return_value = mps + fake_torch.float16 = "float16" + fake_torch.float32 = "float32" + fake_torch.bfloat16 = "bfloat16" + + fake_pipeline_instance = MagicMock() + fake_pipeline_class = MagicMock() + fake_pipeline_class.from_pretrained = MagicMock(return_value=fake_pipeline_instance) + fake_diffusers = MagicMock() + fake_diffusers.LTXPipeline = fake_pipeline_class + + monkeypatch.setitem(sys.modules, "torch", fake_torch) + monkeypatch.setitem(sys.modules, "diffusers", fake_diffusers) + + from tools.video import _shared + importlib.reload(_shared) + + return _shared, fake_pipeline_class, fake_pipeline_instance + + +def test_load_diffusers_pipeline_routes_to_mps(monkeypatch): + """load_diffusers_pipeline must call .to('mps') on Apple Silicon when offload is off.""" + _shared, fake_cls, fake_inst = _make_pipeline_mocks(monkeypatch, mps=True) + + _shared.load_diffusers_pipeline("LTXPipeline", "Lightricks/LTX-Video", enable_offload=False) + + fake_inst.to.assert_called_once_with("mps") + # MPS should use float16 (not bfloat16, not float32) + fake_cls.from_pretrained.assert_called_once_with( + "Lightricks/LTX-Video", torch_dtype="float16" + ) + + +def test_load_diffusers_pipeline_offload_falls_back_on_mps(monkeypatch): + """enable_offload=True on MPS must NOT call enable_model_cpu_offload() — it's CUDA-only.""" + _shared, _, fake_inst = _make_pipeline_mocks(monkeypatch, mps=True) + + _shared.load_diffusers_pipeline("LTXPipeline", "Lightricks/LTX-Video", enable_offload=True) + + # Must NOT call enable_model_cpu_offload() on MPS + fake_inst.enable_model_cpu_offload.assert_not_called() + # Must fall back to .to("mps") + fake_inst.to.assert_called_once_with("mps") + + +def test_load_diffusers_pipeline_cpu_uses_float32(monkeypatch): + """CPU fallback must use float32 — float16 is emulated and unreliable on CPU.""" + _shared, fake_cls, fake_inst = _make_pipeline_mocks(monkeypatch) # neither cuda nor mps + + _shared.load_diffusers_pipeline("LTXPipeline", "Lightricks/LTX-Video", enable_offload=False) + + fake_inst.to.assert_called_once_with("cpu") + fake_cls.from_pretrained.assert_called_once_with( + "Lightricks/LTX-Video", torch_dtype="float32" + ) + + +# ------------------------------------------------------------------ +# install_instructions — Apple Silicon guidance +# ------------------------------------------------------------------ + +def test_upscale_install_instructions_mentions_apple_silicon(): + """install_instructions must not tell M-series users they need CUDA.""" + from tools.enhancement.upscale import Upscale + tool = Upscale() + inst = tool.install_instructions + assert "MPS" in inst or "Apple" in inst or "macOS" in inst, ( + f"install_instructions is CUDA-only, misleads Apple Silicon users: {inst!r}" + ) + + +@pytest.mark.parametrize("module_path,class_name", [ + ("tools.enhancement.face_restore", "FaceRestore"), + ("tools.video.ltx_video_local", "LTXVideoLocal"), +]) +def test_local_gpu_tool_mentions_apple_silicon(module_path, class_name): + """Every LOCAL_GPU tool must tell users MPS/Apple Silicon works.""" + import importlib + mod = importlib.import_module(module_path) + cls = getattr(mod, class_name) + tool = cls() + inst = tool.install_instructions + assert "MPS" in inst or "Apple" in inst or "macOS" in inst, ( + f"{class_name}.install_instructions has no Apple Silicon guidance: {inst!r}" + ) + + +# ------------------------------------------------------------------ +# upscale.py — signature guard for device= +# ------------------------------------------------------------------ + +def test_upscale_build_upsampler_uses_signature_guard(monkeypatch): + """_build_upsampler must use inspect to check if RealESRGANer accepts device=.""" + import importlib + import inspect + + fake_torch = MagicMock() + fake_torch.device = lambda x: f"device({x})" + + fake_rrdbnet = MagicMock() + + # Build a fake RealESRGANer whose __init__ DOES accept device= + class FakeRealESRGANer: + def __init__(self, *, scale, model_path, model, dni_weight, half, device=None): + self.called_with_device = device + fake_realesrganer_cls = FakeRealESRGANer + + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + basicsr_mock = MagicMock() + basicsr_mock.archs.rrdbnet_arch.RRDBNet = fake_rrdbnet + monkeypatch.setitem(sys.modules, "basicsr", basicsr_mock) + monkeypatch.setitem(sys.modules, "basicsr.archs", basicsr_mock.archs) + monkeypatch.setitem(sys.modules, "basicsr.archs.rrdbnet_arch", basicsr_mock.archs.rrdbnet_arch) + + realesrgan_mock = MagicMock() + realesrgan_mock.RealESRGANer = fake_realesrganer_cls + monkeypatch.setitem(sys.modules, "realesrgan", realesrgan_mock) + + fake_shared = MagicMock() + monkeypatch.setitem(sys.modules, "tools.video._shared", fake_shared) + fake_shared.get_torch_device.return_value = "mps" + + from tools.enhancement import upscale + importlib.reload(upscale) + + tool = upscale.Upscale() + result = tool._build_upsampler(scale=4, model_name="RealESRGAN_x4plus", denoise_strength=0.5, face_enhance=False) + assert result.called_with_device == "device(mps)" + + +def test_upscale_build_upsampler_skips_device_when_unsupported(monkeypatch): + """_build_upsampler must NOT pass device= if the installed version doesn't accept it.""" + import importlib + + fake_torch = MagicMock() + fake_torch.device = lambda x: f"device({x})" + + fake_rrdbnet = MagicMock() + + # Build a fake RealESRGANer whose __init__ does NOT accept device= + class FakeRealESRGANerNoDevice: + def __init__(self, *, scale, model_path, model, dni_weight, half): + self.called_with_device = None # no device param + fake_realesrganer_cls = FakeRealESRGANerNoDevice + + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + basicsr_mock = MagicMock() + basicsr_mock.archs.rrdbnet_arch.RRDBNet = fake_rrdbnet + monkeypatch.setitem(sys.modules, "basicsr", basicsr_mock) + monkeypatch.setitem(sys.modules, "basicsr.archs", basicsr_mock.archs) + monkeypatch.setitem(sys.modules, "basicsr.archs.rrdbnet_arch", basicsr_mock.archs.rrdbnet_arch) + + realesrgan_mock = MagicMock() + realesrgan_mock.RealESRGANer = fake_realesrganer_cls + monkeypatch.setitem(sys.modules, "realesrgan", realesrgan_mock) + + fake_shared = MagicMock() + monkeypatch.setitem(sys.modules, "tools.video._shared", fake_shared) + fake_shared.get_torch_device.return_value = "mps" + + from tools.enhancement import upscale + importlib.reload(upscale) + + tool = upscale.Upscale() + # Should NOT raise TypeError about unexpected keyword argument 'device' + result = tool._build_upsampler(scale=4, model_name="RealESRGAN_x4plus", denoise_strength=0.5, face_enhance=False) + assert result.called_with_device is None + + +# ------------------------------------------------------------------ +# face_restore.py — signature guard and device routing +# ------------------------------------------------------------------ + +def test_face_restore_uses_signature_guard(monkeypatch): + """FaceRestore must use inspect to check if GFPGANer accepts device=.""" + import importlib + + fake_cv2 = MagicMock() + fake_cv2.imread.return_value = "fake_image_data" + fake_cv2.imwrite.return_value = True + + fake_torch = MagicMock() + fake_torch.device = lambda x: f"device({x})" + + fake_rrdbnet = MagicMock() + + # GFPGANer that DOES accept device= + class FakeGFPGANer: + def __init__(self, *, model_path, upscale, arch, bg_upsampler=None, device=None): + self.called_with_device = device + def enhance(self, img, **kwargs): + return (None, [1], "fake_restored") + + monkeypatch.setitem(sys.modules, "cv2", fake_cv2) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + basicsr_mock = MagicMock() + basicsr_mock.archs.rrdbnet_arch.RRDBNet = fake_rrdbnet + monkeypatch.setitem(sys.modules, "basicsr", basicsr_mock) + monkeypatch.setitem(sys.modules, "basicsr.archs", basicsr_mock.archs) + monkeypatch.setitem(sys.modules, "basicsr.archs.rrdbnet_arch", basicsr_mock.archs.rrdbnet_arch) + + realesrgan_mock = MagicMock() + monkeypatch.setitem(sys.modules, "realesrgan", realesrgan_mock) + + gfpgan_mock = MagicMock() + gfpgan_mock.GFPGANer = FakeGFPGANer + monkeypatch.setitem(sys.modules, "gfpgan", gfpgan_mock) + + fake_shared = MagicMock() + fake_shared.get_torch_device.return_value = "mps" + monkeypatch.setitem(sys.modules, "tools.video._shared", fake_shared) + + from tools.enhancement import face_restore + importlib.reload(face_restore) + + from pathlib import Path + monkeypatch.setattr(Path, "exists", lambda self: True) + monkeypatch.setattr(Path, "mkdir", lambda self, *a, **kw: None) + + tool = face_restore.FaceRestore() + result = tool.execute({"input_path": "dummy.png"}) + assert result.success is True