# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from abc import ABC, abstractmethod
from typing import Literal, Tuple
import einops
import torch
from nemo.collections.common.parts.utils import mask_sequence_tensor
from nemo.utils import logging
ESTIMATOR_TARGET = Literal['conditional_vector_field', 'data']
[docs]
class ConditionalFlow(ABC):
"""
Abstract class for different conditional flow-matching (CFM) classes
Time horizon is [time_min, time_max (should be 1)]
every path is "conditioned" on endpoints of the path
endpoints are just our paired data samples
subclasses need to implement mean, std, and vector_field
"""
def __init__(self, time_min: float = 1e-8, time_max: float = 1.0):
self.time_min = time_min
self.time_max = time_max
[docs]
@abstractmethod
def mean(self, *, time: torch.Tensor, x_start: torch.Tensor, x_end: torch.Tensor) -> torch.Tensor:
"""
Return the mean of p_t(x | x_start, x_end) at time t
"""
pass
[docs]
@abstractmethod
def std(self, *, time: torch.Tensor, x_start: torch.Tensor, x_end: torch.Tensor) -> torch.Tensor:
"""
Return the standard deviation of p_t(x | x_start, x_end) at time t
"""
pass
[docs]
@abstractmethod
def vector_field(
self, *, time: torch.Tensor, x_start: torch.Tensor, x_end: torch.Tensor, point: torch.Tensor
) -> torch.Tensor:
"""
Compute the conditional vector field v_t( point | x_start, x_end)
"""
pass
@staticmethod
def _broadcast_time(time: torch.Tensor, n_dim: int) -> torch.Tensor:
"""
Broadcast time tensor to the desired number of dimensions
"""
if time.ndim == 1:
target_shape = ' '.join(['B'] + ['1'] * (n_dim - 1))
time = einops.rearrange(time, f'B -> {target_shape}')
return time
[docs]
def generate_time(self, batch_size: int, rng: torch.random.Generator = None) -> torch.Tensor:
"""
Randomly sample a batchsize of time_steps from U[self.time_min, self.time_max]
Supports an external random number generator for better reproducibility
"""
return torch.rand((batch_size,), generator=rng) * (self.time_max - self.time_min) + self.time_min
[docs]
def sample(self, *, time: torch.Tensor, x_start: torch.Tensor, x_end: torch.Tensor) -> torch.Tensor:
"""
Generate a sample from p_t(x | x_start, x_end) at time t.
Note that this implementation assumes all path marginals are normally distributed.
"""
time = self._broadcast_time(time, n_dim=x_start.ndim)
mean = self.mean(time=time, x_start=x_start, x_end=x_end)
std = self.std(time=time, x_start=x_start, x_end=x_end)
return mean + std * torch.randn_like(mean)
[docs]
def flow(
self, *, time: torch.Tensor, x_start: torch.Tensor, x_end: torch.Tensor, point: torch.Tensor
) -> torch.Tensor:
"""
Compute the conditional flow phi_t( point | x_start, x_end).
This is an affine flow.
"""
mean = self.mean(time=time, x_start=x_start, x_end=x_end)
std = self.std(time=time, x_start=x_start, x_end=x_end)
return mean + std * (point - x_start)
[docs]
class OptimalTransportFlow(ConditionalFlow):
"""The OT-CFM model from [Lipman et at, 2023]
Every conditional path the following holds:
p_0 = N(x_start, sigma_start)
p_1 = N(x_end, sigma_end),
mean(x, t) = (time_max - t) * x_start + t * x_end
(linear interpolation between x_start and x_end)
std(x, t) = (time_max - t) * sigma_start + t * sigma_end
Every conditional path is optimal transport map from p_0(x_start, x_end) to p_1(x_start, x_end)
Marginal path is not guaranteed to be an optimal transport map from p_0 to p_1
To get the OT-CFM model from [Lipman et at, 2023] just pass zeroes for x_start
To get the I-CFM model, set sigma_min=sigma_max
To get the rectified flow model, set sigma_min=sigma_max=0
Args:
time_min: minimum time value used in the process
time_max: maximum time value used in the process
sigma_start: the standard deviation of the initial distribution
sigma_end: the standard deviation of the target distribution
"""
def __init__(
self, time_min: float = 1e-8, time_max: float = 1.0, sigma_start: float = 1.0, sigma_end: float = 1e-4
):
super().__init__(time_min=time_min, time_max=time_max)
self.sigma_start = sigma_start
self.sigma_end = sigma_end
logging.debug('Initialized %s with', self.__class__.__name__)
logging.debug('\ttime_min: %s', self.time_min)
logging.debug('\ttime_max: %s', self.time_max)
logging.debug('\tsgima_start: %s', self.sigma_start)
logging.debug('\tsigma_end: %s', self.sigma_end)
[docs]
def mean(self, *, x_start: torch.Tensor, x_end: torch.Tensor, time: torch.Tensor) -> torch.Tensor:
return (self.time_max - time) * x_start + time * x_end
[docs]
def std(self, *, x_start: torch.Tensor, x_end: torch.Tensor, time: torch.Tensor) -> torch.Tensor:
return (self.time_max - time) * self.sigma_start + time * self.sigma_end
[docs]
def vector_field(
self,
*,
x_start: torch.Tensor,
x_end: torch.Tensor,
time: torch.Tensor,
point: torch.Tensor,
eps: float = 1e-6,
) -> torch.Tensor:
time = self._broadcast_time(time, n_dim=x_start.ndim)
if self.sigma_start == self.sigma_end:
return x_end - x_start
num = self.sigma_end * (point - x_start) - self.sigma_start * (point - x_end)
denom = (1 - time) * self.sigma_start + time * self.sigma_end
return num / (denom + eps)
[docs]
class ConditionalFlowMatchingSampler(ABC):
"""
Abstract class for different sampler to solve the ODE in CFM
Args:
estimator: the NN-based conditional vector field estimator
num_steps: How many time steps to iterate in the process
time_min: minimum time value used in the process
time_max: maximum time value used in the process
"""
def __init__(
self,
estimator: torch.nn.Module,
num_steps: int = 5,
time_min: float = 1e-8,
time_max: float = 1.0,
):
self.estimator = estimator
self.num_steps = num_steps
self.time_min = time_min
self.time_max = time_max
@property
def time_step(self):
return (self.time_max - self.time_min) / self.num_steps
[docs]
@abstractmethod
def forward(
self, state: torch.Tensor, estimator_condition: torch.Tensor, state_length: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
pass
[docs]
class ConditionalFlowMatchingEulerSampler(ConditionalFlowMatchingSampler):
"""
The Euler Sampler for solving the ODE in CFM on a uniform time grid
"""
def __init__(
self,
estimator: torch.nn.Module,
num_steps: int = 5,
time_min: float = 1e-8,
time_max: float = 1.0,
estimator_target: ESTIMATOR_TARGET = 'conditional_vector_field',
flow: ConditionalFlow = None,
):
super().__init__(
estimator=estimator,
num_steps=num_steps,
time_min=time_min,
time_max=time_max,
)
self.estimator_target = estimator_target
if self.estimator_target == 'data' and flow is None:
raise ValueError('Flow is required for estimator_target=data')
self.flow = flow
logging.debug('Initialized %s with', self.__class__.__name__)
logging.debug('\tnum_steps: %s', self.num_steps)
logging.debug('\ttime_min: %s', self.time_min)
logging.debug('\ttime_max: %s', self.time_max)
logging.debug('\testimator_target: %s', self.estimator_target)
logging.debug('\tflow: %s', self.flow)
def __call__(self, *args, **kwargs):
return self.forward(*args, **kwargs)
[docs]
@torch.inference_mode()
def forward(
self, state: torch.Tensor, estimator_condition: torch.Tensor, state_length: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
time_steps = torch.linspace(self.time_min, self.time_max, self.num_steps + 1)
if state_length is not None:
state = mask_sequence_tensor(state, state_length)
init_state = state.clone()
for t in time_steps[:-1]:
time = t * torch.ones(state.shape[0], device=state.device)
if estimator_condition is None:
estimator_input = state
else:
estimator_input = torch.cat([state, estimator_condition], dim=1)
if self.estimator_target == 'conditional_vector_field':
vector_field, _ = self.estimator(input=estimator_input, input_length=state_length, condition=time)
state = state + vector_field * self.time_step
elif self.estimator_target == 'data':
endpoint, _ = self.estimator(input=estimator_input, input_length=state_length, condition=time)
vector_field = self.flow.vector_field(time=time, x_start=init_state, x_end=endpoint, point=state)
state = state + vector_field * self.time_step
if state_length is not None:
state = mask_sequence_tensor(state, state_length)
return state, state_length
