fix(smplx): cpu fallback when mps unavailable
This commit is contained in:
@@ -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")
|
||||
Reference in New Issue
Block a user