Spaces:
Running on Zero
Running on Zero
File size: 10,060 Bytes
a1833c8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
from typing import Callable
import torch
import torch.distributed
from diffusers import FlowMatchEulerDiscreteScheduler
from cosmos_framework.model.generator.algorithm.loss.time_weight import TrainTimeWeight
class TrainTimeSampler:
_WAVER_MODE_S = 1.29
# 99.9th and 0.5th percentiles of the standard normal, used for ltx2 stretching.
_LTX2_NORMAL_999_PCTILE = 3.0902
_LTX2_NORMAL_005_PCTILE = -2.5758
_LTX2_UNIFORM_PROB = 0.1
def __init__(
self,
distribution: str = "uniform",
):
self.distribution = distribution
@torch.no_grad()
def __call__(
self,
batch_size: int,
device: torch.device = torch.device("cpu"),
dtype: torch.dtype = torch.float32,
generator: torch.Generator | None = None,
shifts: torch.Tensor | None = None,
) -> torch.Tensor:
"""
Sample sigma ∈ [0, 1] for training.
Args:
batch_size: Number of samples.
device: Target device.
dtype: Target dtype.
shifts: Optional 1-D per-sample shift values, shape ``(batch_size,)``. For non-ltx2
distributions, the raw sample ``t`` is warped through
``sigma = shift * t / (1 + (shift-1) * t)``. For ``ltx2``, ``shifts`` is
required and used as the per-sample logit-normal mean.
Returns:
torch.Tensor: sigma ∈ [0, 1], shape (batch_size,).
"""
if self.distribution == "uniform":
t = torch.rand((batch_size,), generator=generator).to(device=device, dtype=dtype) # [B]
elif self.distribution == "logitnormal":
t = torch.sigmoid(torch.randn((batch_size,), generator=generator)).to(device=device, dtype=dtype) # [B]
elif self.distribution == "waver":
u = torch.rand((batch_size,), dtype=torch.float32, generator=generator) # [B]
t = 1.0 - u - self._WAVER_MODE_S * (torch.cos(torch.pi / 2.0 * u) ** 2 - 1 + u) # [B]
t = t.to(device=device, dtype=dtype) # [B]
elif self.distribution == "ltx2":
# Shifted logit-normal with percentile-based stretching and 10% uniform fallback.
assert shifts is not None, "'ltx2' distribution requires per-sample shifts."
# shift(sigmoid(t), s) = sigmoid(t + ln(s))
mu = torch.log(shifts.to(device=torch.device("cpu"), dtype=torch.float32)) # [B]
std = 1.0
eps = 1e-3
normal_samples = torch.randn((batch_size,), dtype=torch.float32, generator=generator) * std + mu # [B]
logitnormal_samples = torch.sigmoid(normal_samples) # [B]
percentile_999 = torch.sigmoid(mu + self._LTX2_NORMAL_999_PCTILE * std) # [B]
percentile_005 = torch.sigmoid(mu + self._LTX2_NORMAL_005_PCTILE * std) # [B]
zero_terminal_raw = (logitnormal_samples - percentile_005) / (percentile_999 - percentile_005)
stretched = torch.where(
zero_terminal_raw >= eps,
zero_terminal_raw,
2 * eps - zero_terminal_raw,
)
stretched = torch.clamp(stretched, 0, 1)
uniform = (1 - eps) * torch.rand((batch_size,), dtype=torch.float32, generator=generator) + eps
prob = torch.rand((batch_size,), dtype=torch.float32, generator=generator)
t = torch.where(prob > self._LTX2_UNIFORM_PROB, stretched, uniform).to(device=device, dtype=dtype)
return t # skip post-shift
else:
raise NotImplementedError(f"Time distribution '{self.distribution}' is not implemented.")
if shifts is not None:
shifts = shifts.to(device=device, dtype=dtype) # [B]
t = shifts * t / (1 + (shifts - 1) * t) # [B], sigma ∈ [0,1]
return t # [B]
class RectifiedFlow:
def __init__(
self,
velocity_field: Callable,
train_time_distribution: TrainTimeSampler | str = "uniform",
train_time_weight_method: str = "uniform",
use_dynamic_shift: bool = False,
shift: int = 3,
device: torch.device = torch.device("cpu"),
dtype: torch.dtype = torch.float32,
):
r"""Initialize the RectifiedFlow class.
Args:
velocity_field (`Callable`):
A function that predicts the velocity given the current state and time.
train_time_distribution (`TrainTimeSampler` or `str`, *optional*, defaults to `"uniform"`):
Distribution for sampling training times.
Can be an instance of `TrainTimeSampler` or a string specifying the distribution type.
train_time_weight (`TrainTimeWeight` or `str`, *optional*, defaults to `"uniform"`):
Weight applied to training times.
Can be an instance of `TrainTimeWeight` or a string specifying the weight type.
"""
self.velocity_field = velocity_field
self.train_time_sampler: TrainTimeSampler = (
train_time_distribution
if isinstance(train_time_distribution, TrainTimeSampler)
else TrainTimeSampler(train_time_distribution)
)
if use_dynamic_shift:
self.noise_scheduler = FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=use_dynamic_shift)
else:
self.noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=shift)
self.train_time_weight = TrainTimeWeight(self.noise_scheduler, train_time_weight_method)
self.device = torch.device(device) if isinstance(device, str) else device
self.dtype = torch.dtype(dtype) if isinstance(dtype, str) else dtype
def sample_train_time(self, batch_size: int, iteration: int | None = None, shifts: torch.Tensor | None = None):
r"""This method calls the `TrainTimeSampler` to sample training times.
Args:
batch_size: Number of samples.
iteration: When provided, sampling uses a local generator seeded from
``(iteration, rank)`` so results are identical across independent runs
regardless of prior global RNG state.
shifts: Optional 1-D shift tensor, shape ``(batch_size,)``. Forwarded to
``TrainTimeSampler.__call__``; see that docstring for details.
Returns:
t (`torch.Tensor`):
A tensor of sampled sigmas with shape `(batch_size,)`,
matching the class specified `device` and `dtype`.
"""
generator = None
if iteration is not None and torch.are_deterministic_algorithms_enabled():
rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
generator = torch.Generator()
generator.manual_seed(iteration * 65536 + rank)
time = self.train_time_sampler(
batch_size, device=self.device, dtype=self.dtype, generator=generator, shifts=shifts
)
return time
def get_discrete_timestamp(self, u, tensor_kwargs):
r"""This method map time from 0,1 to discrete steps"""
indices = (u.squeeze() * self.noise_scheduler.config.num_train_timesteps).long() # [B]
timesteps = self.noise_scheduler.timesteps.to(**tensor_kwargs)[indices] # [B]
return timesteps.unsqueeze(0) if timesteps.ndim == 0 else timesteps # [B]
def get_sigmas(self, timesteps, tensor_kwargs): # timesteps: [B], returns [B]
sigmas = self.noise_scheduler.sigmas.to(**tensor_kwargs) # [N_timesteps+1]
schedule_timesteps = self.noise_scheduler.timesteps.to(**tensor_kwargs) # [N_timesteps]
step_indices = [(schedule_timesteps == t).nonzero().squeeze().tolist() for t in timesteps]
assert len(step_indices) == timesteps.shape[0], "Number of indices do not match the given timesteps."
sigma = sigmas[step_indices].flatten() # [B]
return sigma # [B]
def get_interpolation(
self,
x_0: list[torch.Tensor], # each element: [B,C,T,H,W] or [B,D1,...,Dn]
x_1: list[torch.Tensor], # each element: [B,C,T,H,W] or [B,D1,...,Dn]
t: list[torch.Tensor], # each element: [B] or [B,1,1,1,1]
):
r"""
This method computes interpolation `X_t` and their time derivatives `dotX_t` at the specified time points `t`.
Note that `x_0` is the noise, and `x_1` is the clean data. This is aligned with the notation in the recified flow community,
but different from the notation in the diffusion community.
Args:
x_0 (`torch.Tensor`):
noise, shape `(B, D1, D2, ..., Dn)`, where `B` is the batch size, and `D1, D2, ..., Dn` are the data dimensions.
x_1 (`torch.Tensor`):
clean data, with the same shape as `x_0`
t (`torch.Tensor`):
A tensor of time steps with values in `[0, 1]`. Can be shape `(B,)` or
pre-broadcast to `(B, 1, T, ..., 1)` matching `x_1`'s dimensionality along batch and temporal dimension.
Returns:
(x_t, dot_x_t) (`Tuple[torch.Tensor, torch.Tensor]`):
- x_t (`torch.Tensor`): The interpolated state, with shape `(B, D1, D2, ..., Dn)`.
- dot_x_t (torch.Tensor): The time derivative of the interpolated state, with the same shape as `x_t`.
"""
assert len(x_0) == len(x_1), "x_0 and x_1 must have the same length."
assert len(x_0) == len(t), "Batch size of x_0 and x_1 must match."
assert len(t) == len(x_1), "Batch size of t must match x_1."
x_t = []
dot_x_t = []
for i in range(len(x_0)):
x_t.append(x_0[i] * t[i] + x_1[i] * (1 - t[i])) # [B,C,T,H,W]; t[i] broadcasts [B] or [B,1,1,1,1]
dot_x_t.append(x_0[i] - x_1[i]) # [B,C,T,H,W]
return x_t, dot_x_t # each list element: [B,C,T,H,W]
|