ardy/model/latent_utils.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Shared utility functions for latent motion models.

These functions were originally in latent_ardy_train.py and are used by auto_latent training and
test-time code.
"""

import logging
from functools import cached_property

import torch

log = logging.getLogger(__name__)


class HybridMotionConverter:
    """Converts between explicit motion and the hybrid (root + latent token) rep.

    Bundles the ``denoiser``, ``autoencoder`` and ``motion_rep`` that every conversion needs so call
    sites don't have to thread them through manually. Derived attributes (``num_frames_per_token``,
    ``motion_rep``) are resolved lazily and raise a clear error if the required object is missing,
    instead of half-initializing.
    """

    def __init__(self, *, denoiser=None, autoencoder=None, motion_rep=None, gen_horizon_len=None):
        if denoiser is None and autoencoder is None:
            raise ValueError("HybridMotionConverter needs a denoiser and/or an autoencoder")
        self.denoiser = denoiser
        self.autoencoder = autoencoder
        self._motion_rep = motion_rep
        self.gen_horizon_len = gen_horizon_len

    @classmethod
    def from_model(cls, ardy_model):
        """Build from an object exposing ``denoiser``/``autoencoder``/``motion_rep``."""
        return cls(
            denoiser=ardy_model.denoiser,
            autoencoder=ardy_model.autoencoder,
            motion_rep=ardy_model.motion_rep,
            gen_horizon_len=getattr(ardy_model, "gen_horizon_len", None),
        )

    @cached_property
    def motion_rep(self):
        rep = self._motion_rep
        if rep is None and self.denoiser is not None:
            rep = self.denoiser.motion_rep
        if rep is None and self.autoencoder is not None:
            rep = self.autoencoder.motion_rep
        if rep is None:
            raise ValueError("No motion_rep available on denoiser or autoencoder")
        return rep

    @cached_property
    def num_frames_per_token(self):
        if self.autoencoder is not None:
            return self.autoencoder.num_frames_per_token
        return self.denoiser.num_frames_per_token

    def get_num_frames_from_hybrid(self, hybrid_motion):
        num_tokens = hybrid_motion.shape[1]
        num_frames = self.num_frames_per_token * num_tokens
        return num_frames

    def get_root_and_latent_body_motion_from_hybrid(self, hybrid_motion):
        num_frames = self.get_num_frames_from_hybrid(hybrid_motion)
        bs = hybrid_motion.shape[0]

        root_motion = hybrid_motion[
            :, :, : self.denoiser.nframe_root_dim
        ]  # (B, num_frames // num_frames_per_token, dim_root * num_frames_per_token)
        root_motion = root_motion.reshape(bs, num_frames, self.motion_rep.motion_root_dim)  # (B, num_frames, dim_root)
        latent_body_motion = hybrid_motion[:, :, self.denoiser.nframe_root_dim :]  # (B, num_tokens, dim_latent_body)
        return root_motion, latent_body_motion

    def get_explicit_motion_from_hybrid(
        self,
        hybrid_motion,
        motion_pad_mask,
        motion_len,
        motion_mask=None,
    ):
        # Separate the root and the latents
        global_root_motion, latent_body_motion = self.get_root_and_latent_body_motion_from_hybrid(hybrid_motion)
        # Compute the local root as a condition
        local_root_motion = self.motion_rep.global_root_to_local_root(
            global_root_motion, normalized=True, lengths=motion_len
        )

        # Detokenize the tokens
        output = self.autoencoder.detokenize(
            latent_body_motion,
            external_cond=local_root_motion,
            motion_pad_mask=motion_pad_mask,
        )
        decoded_body_motion = output["body"]
        # Construct back the motion features
        motion = self.motion_rep.concat_root_body(global_root_motion, decoded_body_motion)
        return motion

    def get_explicit_motion_from_hybrid_autoregressive(
        self,
        hybrid_motion,
        motion_pad_mask,
        motion_len,
        motion_mask=None,
        crop_history_length=None,
    ):
        if self.gen_horizon_len is None:
            raise ValueError("gen_horizon_len is required for autoregressive decoding (see from_model)")
        gen_horizon_len = self.gen_horizon_len
        num_frames_per_token = self.num_frames_per_token
        num_generation_tokens = gen_horizon_len // num_frames_per_token
        token_len = hybrid_motion.shape[1]
        results = []
        for token_idx in range(0, token_len, num_generation_tokens):
            generation_token_end = min(token_idx + num_generation_tokens, token_len)
            generation_token_start = token_idx
            history_token_start = (
                max(0, token_idx - crop_history_length // num_frames_per_token)
                if crop_history_length is not None
                else 0
            )
            generation_frame_start = generation_token_start * num_frames_per_token
            generation_frame_end = generation_token_end * num_frames_per_token
            history_frame_start = history_token_start * num_frames_per_token
            explicit_motion = self.get_explicit_motion_from_hybrid(
                hybrid_motion[:, history_token_start:generation_token_end],
                motion_pad_mask[:, history_frame_start:generation_frame_end],
                motion_pad_mask[:, history_frame_start:generation_frame_end].sum(dim=-1),
                motion_mask=motion_mask[:, history_frame_start:generation_frame_end]
                if motion_mask is not None
                else None,
            )
            results.append(
                explicit_motion[
                    :,
                    generation_frame_start - history_frame_start : generation_frame_end - history_frame_start,
                ]
            )
        motion = torch.cat(results, dim=1)
        return motion

    def get_hybrid_motion_from_root_and_latent_body_motion(
        self,
        root_motion,
        latent_body_motion,
    ):
        bs = root_motion.shape[0]
        nframe_root_dim = self.denoiser.nframe_root_dim
        num_tokens = latent_body_motion.shape[1]
        hybrid_motion = torch.cat(
            [root_motion.reshape(bs, num_tokens, nframe_root_dim), latent_body_motion],
            dim=-1,
        )
        return hybrid_motion

    def get_hybrid_motion_from_explicit(
        self,
        motion,
        motion_len,
        motion_pad_mask,
    ):
        token_embedding = self.autoencoder.tokenize(motion, motion_pad_mask=motion_pad_mask)  # (B, T, D)

        # ensure masked tokens are set to 0
        bs, num_frames = motion_pad_mask.shape
        num_frames_per_token = self.num_frames_per_token

        motion_token_mask = motion_pad_mask.reshape(bs, num_frames // num_frames_per_token, num_frames_per_token).all(
            dim=-1
        )  # (B, num_frames // num_frames_per_token)
        token_embedding = token_embedding * motion_token_mask[:, :, None]

        root = self.motion_rep.extract_root(motion)
        root_reshape = root.reshape(
            bs,
            num_frames // num_frames_per_token,
            self.motion_rep.motion_root_dim * num_frames_per_token,
        )  # (B, num_frames // num_frames_per_token, dim_root * num_frames_per_token)
        hybrid_motion = torch.cat(
            [root_reshape, token_embedding], dim=-1
        )  # (B, num_frames // num_frames_per_token, dim_hybrid)
        hybrid_motion_pad_mask = motion_pad_mask.reshape(
            bs, num_frames // num_frames_per_token, num_frames_per_token
        ).all(dim=-1)  # (B, num_frames // num_frames_per_token)

        return hybrid_motion, hybrid_motion_pad_mask

    def convert_frame_mask_to_token_mask(
        self,
        history_mask,
        generation_mask,
        future_mask,
        motion_mask=None,
    ):
        num_frames_per_token = self.num_frames_per_token
        bs, num_frames = history_mask.shape[0], history_mask.shape[1]
        history_token_mask = history_mask.reshape(bs, num_frames // num_frames_per_token, num_frames_per_token).all(
            dim=-1
        )  # (B, num_frames // num_frames_per_token)
        generation_token_mask = generation_mask.reshape(
            bs, num_frames // num_frames_per_token, num_frames_per_token
        ).all(dim=-1)  # (B, num_frames // num_frames_per_token)
        future_token_mask = torch.zeros_like(history_token_mask)
        if motion_mask is not None:
            future_token_mask = future_mask.reshape(bs, num_frames // num_frames_per_token, num_frames_per_token).all(
                dim=-1
            )  # (B, num_frames // num_frames_per_token)
            has_observation = motion_mask.any(dim=-1)  # if any feature is observed, the token is valid
            has_observation = has_observation.reshape(
                bs, num_frames // num_frames_per_token, num_frames_per_token
            ).any(
                dim=-1
            )  # (B, num_frames // num_frames_per_token), future constraints are sparse, any frame with constraints make the token valid
            future_token_mask = future_token_mask & has_observation

        return history_token_mask, generation_token_mask, future_token_mask