Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

model/flow_matching/helpers/__init__.py

@@ -0,0 +1,20 @@
+"""Flow matching helpers package — direct reuse from ShapeR."""
+
+from model.flow_matching.helpers.path import AffineProbPath, PathSample
+from model.flow_matching.helpers.scheduler import (
+    CondOTScheduler,
+    FluxTimeSampler,
+    TimeSampler,
+)
+from model.flow_matching.helpers.solver import ODESolver
+from model.flow_matching.helpers.model_wrapper import ModelWrapper
+
+__all__ = [
+    "AffineProbPath",
+    "PathSample",
+    "CondOTScheduler",
+    "FluxTimeSampler",
+    "TimeSampler",
+    "ODESolver",
+    "ModelWrapper",
+]

model/flow_matching/helpers/model_wrapper.py

@@ -0,0 +1,46 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This source code is licensed under the CC BY-NC 4.0 license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# This file is derived from https://github.com/facebookresearch/flow_matching
+# Licensed under CC BY-NC 4.0: https://creativecommons.org/licenses/by-nc/4.0/
+
+# pyre-unsafe
+
+from abc import ABC
+
+from torch import nn, Tensor
+
+
+class ModelWrapper(ABC, nn.Module):
+    """
+    This class is used to wrap around another model, adding custom forward pass logic.
+    """
+
+    def __init__(self, model: nn.Module):
+        super().__init__()
+        self.model = model
+
+    def forward(self, x: Tensor, t: Tensor, **extras) -> Tensor:
+        r"""
+        This method defines how inputs should be passed through the wrapped model.
+        Here, we're assuming that the wrapped model takes both :math:`x` and :math:`t` as input,
+        along with any additional keyword arguments.
+
+        Optional things to do here:
+            - check that t is in the dimensions that the model is expecting.
+            - add a custom forward pass logic.
+            - call the wrapped model.
+
+        | given x, t
+        | returns the model output for input x at time t, with extra information `extra`.
+
+        Args:
+            x (Tensor): input data to the model (batch_size, ...).
+            t (Tensor): time (batch_size).
+            **extras: additional information forwarded to the model, e.g., text condition.
+
+        Returns:
+            Tensor: model output.
+        """
+        return self.model(x=x, t=t, **extras)

model/flow_matching/helpers/path.py

@@ -0,0 +1,134 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This source code is licensed under the CC BY-NC 4.0 license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# This file is derived from https://github.com/facebookresearch/flow_matching
+# Licensed under CC BY-NC 4.0: https://creativecommons.org/licenses/by-nc/4.0/
+
+# pyre-unsafe
+
+from dataclasses import dataclass, field
+
+from torch import Tensor
+
+
+def expand_tensor_like(input_tensor: Tensor, expand_to: Tensor) -> Tensor:
+    """`input_tensor` is a 1d vector of length equal to the batch size of `expand_to`,
+    expand `input_tensor` to have the same shape as `expand_to` along all remaining dimensions.
+
+    Args:
+        input_tensor (Tensor): (batch_size,).
+        expand_to (Tensor): (batch_size, ...).
+
+    Returns:
+        Tensor: (batch_size, ...).
+    """
+    assert input_tensor.ndim == 1, "Input tensor must be a 1d vector."
+    assert (
+        input_tensor.shape[0] == expand_to.shape[0]
+    ), f"The first (batch_size) dimension must match. Got shape {input_tensor.shape} and {expand_to.shape}."
+
+    dim_diff = expand_to.ndim - input_tensor.ndim
+
+    t_expanded = input_tensor.clone()
+    t_expanded = t_expanded.reshape(-1, *([1] * dim_diff))
+
+    return t_expanded.expand_as(expand_to)
+
+
+@dataclass
+class PathSample:
+    r"""Represents a sample of a conditional-flow generated probability path.
+
+    Attributes:
+        x_1 (Tensor): the target sample :math:`X_1`.
+        x_0 (Tensor): the source sample :math:`X_0`.
+        t (Tensor): the time sample :math:`t`.
+        x_t (Tensor): samples :math:`X_t \sim p_t(X_t)`, shape (batch_size, ...).
+        dx_t (Tensor): conditional target :math:`\frac{\partial X}{\partial t}`, shape: (batch_size, ...).
+
+    """
+
+    x_1: Tensor = field(metadata={"help": "target samples X_1 (batch_size, ...)."})
+    x_0: Tensor = field(metadata={"help": "source samples X_0 (batch_size, ...)."})
+    t: Tensor = field(metadata={"help": "time samples t (batch_size, ...)."})
+    x_t: Tensor = field(
+        metadata={"help": "samples x_t ~ p_t(X_t), shape (batch_size, ...)."}
+    )
+    dx_t: Tensor = field(
+        metadata={"help": "conditional target dX_t, shape: (batch_size, ...)."}
+    )
+
+
+class AffineProbPath:
+    r"""The ``AffineProbPath`` class represents a specific type of probability path where the transformation between distributions is affine.
+    An affine transformation can be represented as:
+
+    .. math::
+
+        X_t = \alpha_t X_1 + \sigma_t X_0,
+
+    where :math:`X_t` is the transformed data point at time `t`. :math:`X_0` and :math:`X_1` are the source and target data points, respectively. :math:`\alpha_t` and :math:`\sigma_t` are the parameters of the affine transformation at time `t`.
+
+    The scheduler is responsible for providing the time-dependent parameters :math:`\alpha_t` and :math:`\sigma_t`, as well as their derivatives, which define the affine transformation at any given time `t`.
+
+    Using ``AffineProbPath`` in the flow matching framework:
+
+    .. code-block:: python
+
+        # Instantiates a probability path
+        my_path = AffineProbPath(...)
+        mse_loss = torch.nn.MSELoss()
+
+        for x_1 in dataset:
+            # Sets x_0 to random noise
+            x_0 = torch.randn()
+
+            # Sets t to a random value in [0,1]
+            t = torch.rand()
+
+            # Samples the conditional path X_t ~ p_t(X_t|X_0,X_1)
+            path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t)
+
+            # Computes the MSE loss w.r.t. the velocity
+            loss = mse_loss(path_sample.dx_t, my_model(x_t, t))
+            loss.backward()
+
+    Args:
+        scheduler (Scheduler): An instance of a scheduler that provides the parameters :math:`\alpha_t`, :math:`\sigma_t`, and their derivatives over time.
+
+    """
+
+    def __init__(self, scheduler):
+        self.scheduler = scheduler
+
+    def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor):
+        r"""Sample from the affine probability path:
+
+        | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`(\alpha_t,\sigma_t)`.
+        | return :math:`X_0, X_1, X_t = \alpha_t X_1 + \sigma_t X_0`, and the conditional velocity at :math:`X_t, \dot{X}_t = \dot{\alpha}_t X_1 + \dot{\sigma}_t X_0`.
+
+        Args:
+            x_0 (Tensor): source data point, shape (batch_size, ...).
+            x_1 (Tensor): target data point, shape (batch_size, ...).
+            t (Tensor): times in [0,1], shape (batch_size).
+
+        Returns:
+            PathSample: a conditional sample at :math:`X_t \sim p_t`.
+        """
+        scheduler_output = self.scheduler(t)
+        alpha_t = expand_tensor_like(
+            input_tensor=scheduler_output.alpha_t, expand_to=x_1
+        )
+        sigma_t = expand_tensor_like(
+            input_tensor=scheduler_output.sigma_t, expand_to=x_1
+        )
+        d_alpha_t = expand_tensor_like(
+            input_tensor=scheduler_output.d_alpha_t, expand_to=x_1
+        )
+        d_sigma_t = expand_tensor_like(
+            input_tensor=scheduler_output.d_sigma_t, expand_to=x_1
+        )
+        x_t = sigma_t * x_0 + alpha_t * x_1
+        dx_t = d_sigma_t * x_0 + d_alpha_t * x_1
+        return PathSample(x_t=x_t, dx_t=dx_t, x_1=x_1, x_0=x_0, t=t)

model/flow_matching/helpers/scheduler.py

@@ -0,0 +1,146 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This source code is licensed under the CC BY-NC 4.0 license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# This file is derived from https://github.com/facebookresearch/flow_matching
+# Licensed under CC BY-NC 4.0: https://creativecommons.org/licenses/by-nc/4.0/
+
+# pyre-unsafe
+
+import math
+from dataclasses import dataclass, field
+
+import torch
+from torch import Tensor
+from torch.distributions import LogisticNormal, Uniform
+
+
+def get_lin_function(
+    x1: float = 256, y1: float = 0.5, x2: float = 4096, y2: float = 1.15
+):
+    m = (y2 - y1) / (x2 - x1)
+    b = y1 - m * x1
+    return lambda x: m * x + b
+
+
+def time_shift(mu: float, sigma: float, t: Tensor):
+    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+
+@dataclass
+class SchedulerOutput:
+    r"""Represents a sample of a conditional-flow generated probability path.
+
+    Attributes:
+        alpha_t (Tensor): :math:`\alpha_t`, shape (...).
+        sigma_t (Tensor): :math:`\sigma_t`, shape (...).
+        d_alpha_t (Tensor): :math:`\frac{\partial}{\partial t}\alpha_t`, shape (...).
+        d_sigma_t (Tensor): :math:`\frac{\partial}{\partial t}\sigma_t`, shape (...).
+
+    """
+
+    alpha_t: Tensor = field(metadata={"help": "alpha_t"})
+    sigma_t: Tensor = field(metadata={"help": "sigma_t"})
+    d_alpha_t: Tensor = field(metadata={"help": "Derivative of alpha_t."})
+    d_sigma_t: Tensor = field(metadata={"help": "Derivative of sigma_t."})
+
+
+class CondOTScheduler:
+    """CondOT Scheduler."""
+
+    def __call__(self, t: Tensor) -> SchedulerOutput:
+        r"""
+        Args:
+            t (Tensor): times in [0,1], shape (...).
+
+        Returns:
+            SchedulerOutput: :math:`\alpha_t,\sigma_t,\frac{\partial}{\partial t}\alpha_t,\frac{\partial}{\partial t}\sigma_t`
+        """
+        return SchedulerOutput(
+            alpha_t=t,
+            sigma_t=1 - t,
+            d_alpha_t=torch.ones_like(t),
+            d_sigma_t=-torch.ones_like(t),
+        )
+
+    def snr_inverse(self, snr: Tensor) -> Tensor:
+        r"""
+        Computes :math:`t` from the signal-to-noise ratio :math:`\frac{\alpha_t}{\sigma_t}`.
+
+        Args:
+            snr (Tensor): The signal-to-noise, shape (...)
+
+        Returns:
+            Tensor: t, shape (...)
+        """
+        kappa_t = snr / (1.0 + snr)
+
+        return self.kappa_inverse(kappa=kappa_t)
+
+    def kappa_inverse(self, kappa: Tensor) -> Tensor:
+        """
+        Computes :math:`t` from :math:`kappa_t`.
+
+        Args:
+            kappa (Tensor): :math:`kappa`, shape (...)
+
+        Returns:
+            Tensor: t, shape (...)
+        """
+        return kappa
+
+
+class TimeSampler:
+    def __init__(self, dist):
+        self.dist_name = dist
+        if dist == "uniform":
+            self.distribution = Uniform(0.0, 1.0)
+        elif dist == "lognormal":
+            self.distribution = LogisticNormal(0.0, 1.0)
+        elif dist == "flux":
+            self.distribution = FluxTimeSampler(mode="train")
+        else:
+            raise ValueError(f"Unknown distribution {dist}")
+
+    def __call__(self, shape, num_tokens, device):
+        if self.dist_name == "uniform":
+            return self.distribution.sample(sample_shape=(shape,)).to(device).squeeze(0)
+        elif self.dist_name == "lognormal":
+            return self.distribution.sample(sample_shape=(shape,))[:, 0].to(device)
+        elif self.dist_name == "flux":
+            return self.distribution(shape, num_tokens, device)
+
+
+class FluxTimeSampler:
+    def __init__(self, **kwargs):
+        self.mode = kwargs.get("mode", "train")
+
+    def __call__(self, bs, num_tokens, device):
+        if self.mode == "train":
+            logits_norm = torch.randn(bs, device=device)
+            timesteps = logits_norm.sigmoid()
+            mu = get_lin_function(y1=0.5, y2=1.15)(num_tokens)
+            timesteps = time_shift(mu, 1.0, timesteps)
+            timesteps = 1 - timesteps
+        elif self.mode == "inference":
+            timesteps = torch.linspace(1, 0, bs + 1).to(device)
+            mu = get_lin_function(y1=0.5, y2=1.15)(num_tokens)
+            timesteps = time_shift(mu, 1.0, timesteps)
+            timesteps = 1 - timesteps
+        return timesteps
+
+
+def visualize_sampling():
+    import matplotlib.pyplot as plt
+
+    flux_sampler = FluxTimeSampler(mode="inference")
+    timesteps = flux_sampler(25, 2048, torch.device("cuda"))
+    x = torch.linspace(0, 1, 25 + 1)
+    plt.scatter(x, y=timesteps.cpu().numpy())
+    plt.savefig("scatter_plot.png")
+
+    flux_sampler = FluxTimeSampler(mode="train")
+    timesteps = flux_sampler(5000, 4096, torch.device("cuda"))
+    plt.hist(timesteps.cpu().numpy(), bins=50, edgecolor="black")
+    plt.savefig("scatter_plot_4096.png")
+    print(timesteps)

model/flow_matching/helpers/solver.py

@@ -0,0 +1,78 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This source code is licensed under the CC BY-NC 4.0 license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# This file is derived from https://github.com/facebookresearch/flow_matching
+# Licensed under CC BY-NC 4.0: https://creativecommons.org/licenses/by-nc/4.0/
+
+# pyre-unsafe
+
+import torch
+from torch import Tensor
+from torchdiffeq import odeint
+
+
+class ODESolver:
+    """A class to solve ordinary differential equations (ODEs) using a specified velocity model.
+
+    This class utilizes a velocity field model to solve ODEs over a given time grid using numerical ode solvers.
+
+    Args:
+        velocity_model (Union[ModelWrapper, Callable]): a velocity field model receiving :math:`(x,t)` and returning :math:`u_t(x)`
+    """
+
+    def __init__(self, velocity_model):
+        super().__init__()
+        self.velocity_model = velocity_model
+
+    def sample(
+        self,
+        x_init: Tensor,
+        # pyre-fixme[9]: step_size has type `float`; used as `None`.
+        step_size: float = None,
+        method: str = "euler",
+        atol: float = 1e-5,
+        rtol: float = 1e-5,
+        time_grid: Tensor = torch.tensor([0.0, 1.0]),
+        return_intermediates: bool = False,
+        enable_grad: bool = False,
+        **model_extras,
+    ):
+        r"""Solve the ODE with the velocity field.
+
+        Args:
+            x_init (Tensor): initial conditions (e.g., source samples :math:`X_0 \sim p`). Shape: [batch_size, ...].
+            step_size (Optional[float]): The step size. Must be None for adaptive step solvers.
+            method (str): A method supported by torchdiffeq. Defaults to "euler". Other commonly used solvers are "dopri5", "midpoint" and "heun3". For a complete list, see torchdiffeq.
+            atol (float): Absolute tolerance, used for adaptive step solvers.
+            rtol (float): Relative tolerance, used for adaptive step solvers.
+            time_grid (Tensor): The process is solved in the interval [min(time_grid, max(time_grid)] and if step_size is None then time discretization is set by the time grid. May specify a descending time_grid to solve in the reverse direction. Defaults to torch.tensor([0.0, 1.0]).
+            return_intermediates (bool, optional): If True then return intermediate time steps according to time_grid. Defaults to False.
+            enable_grad (bool, optional): Whether to compute gradients during sampling. Defaults to False.
+            **model_extras: Additional input for the model.
+
+        Returns:
+          The last timestep when return_intermediates=False, otherwise all values specified in time_grid.
+        """
+
+        def ode_func(t, x):
+            return self.velocity_model(x=x, t=t, **model_extras)
+
+        ode_opts = {"step_size": step_size} if step_size is not None else {}
+
+        with torch.set_grad_enabled(enable_grad):
+            # Approximate ODE solution with numerical ODE solver
+            sol = odeint(
+                ode_func,
+                x_init,
+                time_grid,
+                method=method,
+                options=ode_opts,
+                atol=atol,
+                rtol=rtol,
+            )
+
+        if return_intermediates:
+            return sol
+        else:
+            return sol[-1]