File size: 4,955 Bytes
bd95c9c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""
Smoke tests: SOMALayer forward pass for each identity model (soma, mhr, anny, smpl, garment)
as in tools/demo_soma_vis.py. Parametrized over CPU and CUDA; CUDA is skipped when unavailable.
Fails if assets/SOMA_neutral.npz is not present (e.g. run `git lfs pull` after clone).
Optional models (smpl, anny) are skipped when their assets or dependencies are missing.
"""
from pathlib import Path
import pytest
import torch
# Repo root = parent of tests/
REPO_ROOT = Path(__file__).resolve().parents[1]
ASSETS_DIR = REPO_ROOT / "assets"
CORE_ASSET = ASSETS_DIR / "SOMA_neutral.npz"
@pytest.fixture(scope="module")
def data_root():
if not ASSETS_DIR.is_dir():
pytest.fail(
f"Assets directory not found: {ASSETS_DIR}. "
"Clone the repo and run `git lfs pull` to fetch assets."
)
if not CORE_ASSET.is_file():
pytest.fail(
f"Required asset not found: {CORE_ASSET}. "
"Run `git lfs pull` (or `git lfs pull assets/`) to fetch LFS-tracked files."
)
return str(ASSETS_DIR)
def _make_layer(data_root, identity_model_type, device, low_lod=False):
"""Create SOMALayer; return (layer, skip_reason). skip_reason is non-None to skip the test."""
from soma import SOMALayer
try:
layer = SOMALayer(
data_root=data_root,
low_lod=low_lod,
device=device,
identity_model_type=identity_model_type,
mode="warp",
).to(device)
return layer, None
except (FileNotFoundError, ImportError) as e:
return None, f"Missing asset or dependency: {e}"
def _make_inputs(layer, identity_model_type, device, batch_size=1):
"""Build identity_coeffs and scale_params for the given identity model type."""
if identity_model_type == "anny":
ann = layer.identity_model.identity_model
identity_coeffs = {
k: torch.ones(batch_size, device=device) * 0.5 for k in ann.phenotype_labels
}
scale_params = {k: torch.zeros(batch_size, device=device) for k in ann.local_change_labels}
elif identity_model_type == "mhr":
n_id = layer.identity_model.num_identity_coeffs
n_scale = layer.identity_model.num_scale_params
identity_coeffs = torch.zeros(batch_size, n_id, device=device)
scale_params = torch.zeros(batch_size, n_scale, device=device)
else:
n_id = layer.identity_model.num_identity_coeffs
identity_coeffs = torch.zeros(batch_size, n_id, device=device)
scale_params = None
return identity_coeffs, scale_params
@pytest.mark.parametrize(
"apply_correctives",
[
pytest.param(False, id="no_correctives"),
pytest.param(True, id="correctives"),
],
)
@pytest.mark.parametrize(
"low_lod",
[
pytest.param(False, id="high_lod"),
pytest.param(True, id="low_lod"),
],
)
@pytest.mark.parametrize("identity_model_type", ["soma", "mhr", "anny", "smpl", "smplx", "garment"])
@pytest.mark.parametrize("device", ["cpu", "cuda"])
def test_soma_layer_forward(data_root, identity_model_type, device, low_lod, apply_correctives):
"""SOMALayer forward pass for each identity model, LOD, and corrective mode."""
if device == "cuda" and not torch.cuda.is_available():
pytest.skip("CUDA not available")
layer, skip_reason = _make_layer(data_root, identity_model_type, device, low_lod=low_lod)
if skip_reason is not None:
pytest.skip(skip_reason)
if apply_correctives and layer.correctives_model is None:
pytest.skip("Corrective model not available")
if low_lod:
assert layer.nv_lod_mid_to_low is not None
expected_num_verts = layer.nv_lod_mid_to_low.shape[0]
assert layer.bind_shape.shape[0] == expected_num_verts
batch_size = 1
num_pose_joints = 77
pose = torch.zeros(batch_size, num_pose_joints, 3, 3, device=device)
transl = torch.zeros(batch_size, 3, device=device)
identity_coeffs, scale_params = _make_inputs(layer, identity_model_type, device, batch_size)
with torch.no_grad():
out = layer(
pose,
identity_coeffs,
scale_params=scale_params,
transl=transl,
pose2rot=False,
apply_correctives=apply_correctives,
)
assert "vertices" in out
assert "joints" in out
verts = out["vertices"]
joints = out["joints"]
assert verts.dim() == 3 and verts.shape[0] == batch_size and verts.shape[2] == 3
if low_lod:
assert verts.shape[1] == expected_num_verts, (
f"Expected {expected_num_verts} low-LOD vertices, got {verts.shape[1]}"
)
assert joints.dim() == 3 and joints.shape[0] == batch_size and joints.shape[2] == 3
assert joints.shape[1] == num_pose_joints
|