soma_body.py

"""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.
# `<Finger>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 `<Finger>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]
        )