"""Apache-2.0 body model for the StreamTalk demo. StreamTalk's reference pipeline drives an SMPL-X neutral body, but the official SMPL-X model file may not be redistributed by third parties, so it cannot be fetched from a community mirror. This module replaces it with NVIDIA's **SOMA-X** body (`nvidia/SOMA-X`, Apache-2.0, freely redistributable), whose unified skeleton contains a joint for every SMPL-X joint. `assets/soma_rig.npz` is a compact extract of the SOMA-X v0.3.0 neutral template rig (joint hierarchy, T-pose/bind transforms, skinning weights and the neutral mesh) redistributed here under the Apache-2.0 licence — see `LICENSE-SOMA-X`. Pose convention --------------- SOMA parameterises poses as rotations *relative to the T-pose*, composed down the kinematic chain exactly like SMPL-X: world[j] = world[parent] @ orient[parent]^T @ R_rel[j] @ orient[j] so that ``world[j] @ orient[j]^T = (world[parent] @ orient[parent]^T) @ R_rel[j]``. The accumulated term behaves identically to an SMPL-X global rotation, which means SMPL-X local joint rotations can be handed to SOMA joints unchanged; only the (slightly different) rest skeleton and mesh change. """ from __future__ import annotations from pathlib import Path import numpy as np import torch import torch.nn as nn # SMPL-X joint index -> SOMA joint index. # SMPL-X hand order is index / middle / pinky / ring / thumb. # # Fingers need care: SOMA's four fingers carry one joint more than SMPL-X's. # `1` sits ~3 cm from the wrist at the base of the metacarpal (the four # `*1` joints span only 4 cm across the palm), while `2` sits ~9 cm out # on the knuckle line (6.2 cm span) — i.e. SOMA `2/3/4` are the MCP / PIP / DIP # joints that SMPL-X calls `1/2/3`. The metacarpals are left un-rotated, which # is what SMPL-X assumes anyway. The thumb chain has no extra joint, so # `Thumb1/2/3` map straight across. SMPLX_TO_SOMA = [ 1, # 0 pelvis -> Hips 68, # 1 left_hip -> LeftLeg 73, # 2 right_hip -> RightLeg 2, # 3 spine1 -> Spine1 69, # 4 left_knee -> LeftShin 74, # 5 right_knee -> RightShin 3, # 6 spine2 -> Spine2 70, # 7 left_ankle -> LeftFoot 75, # 8 right_ankle -> RightFoot 4, # 9 spine3 -> Chest 71, # 10 left_foot -> LeftToeBase 76, # 11 right_foot -> RightToeBase 5, # 12 neck -> Neck1 12, # 13 left_collar -> LeftShoulder 40, # 14 right_collar -> RightShoulder 7, # 15 head -> Head 13, # 16 left_shoulder -> LeftArm 41, # 17 right_shoulder-> RightArm 14, # 18 left_elbow -> LeftForeArm 42, # 19 right_elbow -> RightForeArm 15, # 20 left_wrist -> LeftHand 43, # 21 right_wrist -> RightHand 9, # 22 jaw -> Jaw 10, # 23 left_eye -> LeftEye 11, # 24 right_eye -> RightEye 21, 22, 23, # left index -> LeftHandIndex2/3/4 26, 27, 28, # left middle -> LeftHandMiddle2/3/4 36, 37, 38, # left pinky -> LeftHandPinky2/3/4 31, 32, 33, # left ring -> LeftHandRing2/3/4 16, 17, 18, # left thumb -> LeftHandThumb1/2/3 49, 50, 51, # right index -> RightHandIndex2/3/4 54, 55, 56, # right middle -> RightHandMiddle2/3/4 64, 65, 66, # right pinky -> RightHandPinky2/3/4 59, 60, 61, # right ring -> RightHandRing2/3/4 44, 45, 46, # right thumb -> RightHandThumb1/2/3 ] assert len(SMPLX_TO_SOMA) == 55 and len(set(SMPLX_TO_SOMA)) == 55 CM_TO_M = 0.01 def _se3_inverse(T: torch.Tensor) -> torch.Tensor: R = T[..., :3, :3] t = T[..., :3, 3] out = torch.zeros_like(T) Rt = R.transpose(-1, -2) out[..., :3, :3] = Rt out[..., :3, 3] = -(Rt @ t.unsqueeze(-1)).squeeze(-1) out[..., 3, 3] = 1.0 return out class SomaBody(nn.Module): """Forward kinematics + linear blend skinning for the SOMA-X neutral body. Everything is registered as a buffer so ZeroGPU packs the model into VRAM together with the rest of the pipeline. """ def __init__(self, rig_path: str | Path): super().__init__() blob = np.load(rig_path, allow_pickle=False) parents = torch.from_numpy(blob["parents"].astype(np.int64)) bind = torch.from_numpy(blob["bind_pose_world"].astype(np.float32)).clone() bind[..., :3, 3] *= CM_TO_M t_pose = torch.from_numpy(blob["t_pose_world"].astype(np.float32)) orient = t_pose[:, :3, :3].contiguous() bind_local = _se3_inverse(bind)[parents] @ bind bind_local[0] = bind[0] self.num_joints = int(parents.shape[0]) self.parents = parents.tolist() self.register_buffer("orient", orient, persistent=False) self.register_buffer( "orient_parent_T", orient[parents].transpose(-1, -2).contiguous(), persistent=False, ) self.register_buffer("local_t", bind_local[:, :3, 3].contiguous(), persistent=False) self.register_buffer("inv_bind", _se3_inverse(bind), persistent=False) self.register_buffer( "bind_shape", torch.from_numpy(blob["bind_shape"].astype(np.float32)) * CM_TO_M, persistent=False, ) self.register_buffer( "skin_weights", torch.from_numpy(blob["skin_weights"].astype(np.float32)), persistent=False, ) self.register_buffer( "smplx_to_soma", torch.tensor(SMPLX_TO_SOMA, dtype=torch.long), persistent=False, ) self.faces = blob["faces"].astype(np.int64) # ------------------------------------------------------------------ # def _world_transforms( self, smplx_rotmats: torch.Tensor, transl: torch.Tensor | None = None ) -> torch.Tensor: """smplx_rotmats: (B, 55, 3, 3) SMPL-X *local* rotations -> (B, J, 4, 4).""" B = smplx_rotmats.shape[0] R = torch.eye(3, device=smplx_rotmats.device, dtype=smplx_rotmats.dtype) R = R.expand(B, self.num_joints, 3, 3).clone() R[:, self.smplx_to_soma] = smplx_rotmats local_R = self.orient_parent_T.unsqueeze(0) @ R @ self.orient.unsqueeze(0) T = torch.zeros( B, self.num_joints, 4, 4, device=R.device, dtype=R.dtype ) T[..., :3, :3] = local_R T[..., :3, 3] = self.local_t.unsqueeze(0) T[..., 3, 3] = 1.0 if transl is not None: T[:, 1, :3, 3] = T[:, 1, :3, 3] + transl world = [T[:, 0]] for j in range(1, self.num_joints): world.append(world[self.parents[j]] @ T[:, j]) return torch.stack(world, dim=1) def joints(self, smplx_rotmats: torch.Tensor) -> torch.Tensor: """SMPL-X-ordered joint positions, (B, 55, 3), in metres.""" world = self._world_transforms(smplx_rotmats) return world[:, self.smplx_to_soma, :3, 3] def vertices( self, smplx_rotmats: torch.Tensor, transl: torch.Tensor | None = None ) -> torch.Tensor: """Skinned mesh vertices, (B, V, 3), in metres.""" world = self._world_transforms(smplx_rotmats, transl) skin = world @ self.inv_bind.unsqueeze(0) mats = torch.einsum("vj,bjmn->bvmn", self.skin_weights, skin) return ( torch.einsum("bvmn,vn->bvm", mats[..., :3, :3], self.bind_shape) + mats[..., :3, 3] )