Spaces:
Running on Zero
Running on Zero
| # 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 | |
| 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] | |