fix(smplx): cpu fallback when mps unavailable

This commit is contained in:
L'électron rare
2026-05-13 13:02:48 +02:00
parent 722208e0eb
commit b0d3ec15aa
2 changed files with 58 additions and 4 deletions
+9 -4
View File
@@ -26,9 +26,14 @@ class SMPLXDecoder:
"""Charge SMPL-X NEUTRAL et expose decode(params) -> (verts, joints)."""
def __init__(self, model_path: str, device: str = "mps") -> None:
_require_torch()
import smplx
torch = _require_torch()
# Demote unsupported devices to CPU (mirrors MultiHMRWorker pattern)
if device == "mps" and not torch.backends.mps.is_available():
device = "cpu"
elif device.startswith("cuda") and not torch.cuda.is_available():
device = "cpu"
self.device = device
import smplx
model_path_p = Path(model_path)
if model_path_p.is_file():
# smplx.SMPLXLayer attend le dossier contenant SMPLX_<GENDER>.<ext>
@@ -43,8 +48,8 @@ class SMPLXDecoder:
num_betas=10,
num_expression_coeffs=10,
ext=ext,
).to(device).eval()
LOG.info("SMPL-X loaded from %s (device=%s)", model_folder, device)
).to(self.device).eval()
LOG.info("SMPL-X loaded from %s (device=%s)", model_folder, self.device)
def decode(
self,
@@ -0,0 +1,49 @@
"""SMPLXDecoder must validate device and fall back to CPU when MPS unavailable."""
import sys
from types import ModuleType
from unittest.mock import MagicMock, patch
def _make_smplx_mock():
"""Return a fake smplx module whose SMPLXLayer is chainable."""
mock_smplx = ModuleType("smplx")
layer = MagicMock()
layer.to.return_value = layer
layer.eval.return_value = layer
mock_smplx.SMPLXLayer = MagicMock(return_value=layer)
return mock_smplx, layer
def test_decoder_falls_back_to_cpu_when_mps_unavailable(tmp_path):
mock_smplx, layer = _make_smplx_mock()
with patch.dict(sys.modules, {"smplx": mock_smplx}):
from data_only_viz import smplx_decoder
with patch.object(smplx_decoder, "_require_torch") as mock_require:
fake_torch = mock_require.return_value
fake_torch.backends.mps.is_available.return_value = False
fake_torch.cuda.is_available.return_value = False
dec = smplx_decoder.SMPLXDecoder(model_path=tmp_path, device="mps")
# The decoder must NOT have called .to("mps") — it must have demoted to "cpu":
layer.to.assert_called_with("cpu")
assert dec.device == "cpu"
def test_decoder_keeps_cpu_device(tmp_path):
mock_smplx, layer = _make_smplx_mock()
with patch.dict(sys.modules, {"smplx": mock_smplx}):
from data_only_viz import smplx_decoder
with patch.object(smplx_decoder, "_require_torch") as mock_require:
fake_torch = mock_require.return_value
fake_torch.backends.mps.is_available.return_value = False
fake_torch.cuda.is_available.return_value = False
dec = smplx_decoder.SMPLXDecoder(model_path=tmp_path, device="cpu")
assert dec.device == "cpu"
layer.to.assert_called_with("cpu")