scheduling_flow_match_euler_discrete ¶

def add_noise(
    self,
    clean_latent: torch.Tensor,
    noise: torch.Tensor,
    timestep: torch.IntTensor,
) -> torch.Tensor:

    """
    Args:
        clean_latent: the clean latent with shape [B, C, H, W],
            where B is batch_size or batch_size * num_frames
        noise: the noise with shape [B, C, H, W]
        timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames]

    Returns:
        the corrupted latent with shape [B, C, H, W]
    """
    # If timestep is [bs, num_frames]
    if timestep.ndim == 2:
        timestep = timestep.flatten(0, 1)
        assert timestep.numel() == clean_latent.shape[0]
    elif timestep.ndim == 1:
        # If timestep is [1]
        if timestep.shape[0] == 1:
            timestep = timestep.expand(clean_latent.shape[0])
        else:
            assert timestep.numel() == clean_latent.shape[0]
    else:
        raise ValueError(f"[add_noise] Invalid timestep shape: {timestep.shape}")
    # timestep shape should be [B]
    self.sigmas = self.sigmas.to(noise.device)
    self.timesteps = self.timesteps.to(noise.device)
    timestep_id = torch.argmin(
        (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1)
    sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1)
    sample = (1 - sigma) * clean_latent + sigma * noise
    return sample.type_as(noise)

def scale_noise(
    self,
    sample: torch.FloatTensor,
    timestep: float | torch.FloatTensor,
    noise: torch.FloatTensor | None = None,
) -> torch.FloatTensor:
    """
    Forward process in flow-matching

    Args:
        sample (`torch.FloatTensor`):
            The input sample.
        timestep (`int`, *optional*):
            The current timestep in the diffusion chain.

    Returns:
        `torch.FloatTensor`:
            A scaled input sample.
    """
    # Make sure sigmas and timesteps have the same device and dtype as original_samples
    sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype)

    if sample.device.type == "mps" and torch.is_floating_point(timestep):
        # mps does not support float64
        schedule_timesteps = self.timesteps.to(sample.device,
                                               dtype=torch.float32)
        assert isinstance(timestep, torch.Tensor)
        timestep = timestep.to(sample.device, dtype=torch.float32)
    else:
        schedule_timesteps = self.timesteps.to(sample.device)
        assert isinstance(timestep, torch.Tensor)
        timestep = timestep.to(sample.device)

    # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index
    if self.begin_index is None:
        step_indices = [
            self.index_for_timestep(t, schedule_timesteps) for t in timestep
        ]
    elif self.step_index is not None:
        # add_noise is called after first denoising step (for inpainting)
        step_indices = [self.step_index] * timestep.shape[0]
    else:
        # add noise is called before first denoising step to create initial latent(img2img)
        step_indices = [self.begin_index] * timestep.shape[0]

    sigma = sigmas[step_indices].flatten()
    while len(sigma.shape) < len(sample.shape):
        sigma = sigma.unsqueeze(-1)

    sample = sigma * noise + (1.0 - sigma) * sample

    return sample

def set_begin_index(self, begin_index: int = 0) -> None:
    """
    Sets the begin index for the scheduler. This function should be run from pipeline before the inference.

    Args:
        begin_index (`int`):
            The begin index for the scheduler.
    """
    self._begin_index = begin_index

def set_timesteps(
    self,
    num_inference_steps: int | None = None,
    device: str | torch.device = None,
    sigmas: list[float] | None = None,
    mu: float | None = None,
    timesteps: list[float] | None = None,
) -> None:
    """
    Sets the discrete timesteps used for the diffusion chain (to be run before inference).

    Args:
        num_inference_steps (`int`, *optional*):
            The number of diffusion steps used when generating samples with a pre-trained model.
        device (`str` or `torch.device`, *optional*):
            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
        sigmas (`List[float]`, *optional*):
            Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed
            automatically.
        mu (`float`, *optional*):
            Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep
            shifting.
        timesteps (`List[float]`, *optional*):
            Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed
            automatically.
    """
    if self.config.use_dynamic_shifting and mu is None:
        raise ValueError(
            "`mu` must be passed when `use_dynamic_shifting` is set to be `True`"
        )

    if sigmas is not None and timesteps is not None and len(sigmas) != len(
            timesteps):
        raise ValueError(
            "`sigmas` and `timesteps` should have the same length")

    if num_inference_steps is not None:
        if (sigmas is not None and len(sigmas) != num_inference_steps) or (
                timesteps is not None
                and len(timesteps) != num_inference_steps):
            raise ValueError(
                "`sigmas` and `timesteps` should have the same length as num_inference_steps, if `num_inference_steps` is provided"
            )
    else:
        if sigmas is not None:
            num_inference_steps = len(sigmas)
        elif timesteps is not None:
            num_inference_steps = len(timesteps)
        else:
            raise ValueError(
                "Either num_inference_steps, sigmas, or timesteps must be provided"
            )

    self.num_inference_steps = num_inference_steps

    # 1. Prepare default sigmas
    is_timesteps_provided = timesteps is not None

    timesteps_array: np.ndarray | None = None
    if is_timesteps_provided:
        assert timesteps is not None
        timesteps_array = np.array(timesteps).astype(np.float32)

    sigmas_array: np.ndarray
    if sigmas is None:
        if timesteps_array is None:
            t_max = self._sigma_to_t(self.sigma_max)
            t_min = self._sigma_to_t(self.sigma_min)
            if self.config.use_reference_discrete_timesteps:
                # Some reference schedulers (for example Z-Image) build a
                # float64 num_steps+1 linspace and drop the terminal point.
                timesteps_array = np.linspace(
                    t_max,
                    t_min,
                    num_inference_steps + 1,
                )[:-1]
            else:
                # Preserve the original numpy default (float64) here —
                # casting to float32 silently shifts rounded timestep
                # values for every existing model that uses this branch.
                timesteps_array = np.linspace(t_max, t_min, num_inference_steps)
        sigmas_array = timesteps_array / self.config.num_train_timesteps
    else:
        sigmas_array = np.array(sigmas).astype(np.float32)
        num_inference_steps = len(sigmas_array)

    # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of
    #    "exponential" or "linear" type is applied
    if self.config.use_dynamic_shifting:
        assert mu is not None, "mu cannot be None when use_dynamic_shifting is True"
        sigmas_array = self.time_shift(mu, 1.0, sigmas_array)
    else:
        sigmas_array = self.shift * sigmas_array / (
            1 + (self.shift - 1) * sigmas_array)

    # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value
    if self.config.shift_terminal:
        sigmas_tensor = torch.from_numpy(sigmas_array).to(
            dtype=torch.float32)
        sigmas_tensor = self.stretch_shift_to_terminal(sigmas_tensor)
        sigmas_array = sigmas_tensor.numpy()

    # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules
    if self.config.use_karras_sigmas:
        sigmas_tensor = torch.from_numpy(sigmas_array).to(
            dtype=torch.float32)
        sigmas_tensor = self._convert_to_karras(
            in_sigmas=sigmas_tensor,
            num_inference_steps=num_inference_steps)
        sigmas_array = sigmas_tensor.numpy()
    elif self.config.use_exponential_sigmas:
        sigmas_tensor = torch.from_numpy(sigmas_array).to(
            dtype=torch.float32)
        sigmas_tensor = self._convert_to_exponential(
            in_sigmas=sigmas_tensor,
            num_inference_steps=num_inference_steps)
        sigmas_array = sigmas_tensor.numpy()
    elif self.config.use_beta_sigmas:
        sigmas_tensor = torch.from_numpy(sigmas_array).to(
            dtype=torch.float32)
        sigmas_tensor = self._convert_to_beta(
            in_sigmas=sigmas_tensor,
            num_inference_steps=num_inference_steps)
        sigmas_array = sigmas_tensor.numpy()

    # 5. Convert sigmas and timesteps to tensors and move to specified device
    sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32,
                                                      device=device)
    if not is_timesteps_provided:
        timesteps_tensor = sigmas_tensor * self.config.num_train_timesteps
    else:
        assert timesteps_array is not None
        timesteps_tensor = torch.from_numpy(timesteps_array).to(
            dtype=torch.float32, device=device)

    # 6. Append the terminal sigma value.
    #    If a model requires inverted sigma schedule for denoising but timesteps without inversion, the
    #    `invert_sigmas` flag can be set to `True`. This case is only required in Mochi
    if self.config.invert_sigmas:
        sigmas_tensor = 1.0 - sigmas_tensor
        timesteps_tensor = sigmas_tensor * self.config.num_train_timesteps
        sigmas_tensor = torch.cat(
            [sigmas_tensor,
             torch.ones(1, device=sigmas_tensor.device)])
    else:
        sigmas_tensor = torch.cat([sigmas_tensor, torch.zeros(1, device=sigmas_tensor.device)])

    self.timesteps = timesteps_tensor
    self.sigmas = sigmas_tensor
    self._step_index = None
    self._begin_index = None

def step(
    self,
    model_output: torch.FloatTensor,
    timestep: int | torch.Tensor,
    sample: torch.FloatTensor,
    s_churn: float = 0.0,
    s_tmin: float = 0.0,
    s_tmax: float = float("inf"),
    s_noise: float = 1.0,
    generator: torch.Generator | None = None,
    per_token_timesteps: torch.Tensor | None = None,
    return_dict: bool = True,
) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor, ...]:
    """
    Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
    process from the learned model outputs (most often the predicted noise).

    Args:
        model_output (`torch.FloatTensor`):
            The direct output from learned diffusion model.
        timestep (`int` or `torch.Tensor`):
            The current discrete timestep in the diffusion chain.
        sample (`torch.FloatTensor`):
            A current instance of a sample created by the diffusion process.
        s_churn (`float`):
        s_tmin  (`float`):
        s_tmax  (`float`):
        s_noise (`float`, defaults to 1.0):
            Scaling factor for noise added to the sample.
        generator (`torch.Generator`, *optional*):
            A random number generator.
        per_token_timesteps (`torch.Tensor`, *optional*):
            The timesteps for each token in the sample.
        return_dict (`bool`):
            Whether or not to return a
            [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple.

    Returns:
        [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`:
            If return_dict is `True`,
            [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned,
            otherwise a tuple is returned where the first element is the sample tensor.
    """

    if (isinstance(timestep, int | torch.IntTensor | torch.LongTensor)):
        raise ValueError((
            "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"
            " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass"
            " one of the `scheduler.timesteps` as a timestep."), )

    if self.step_index is None:
        self._init_step_index(timestep)

    # Upcast to avoid precision issues when computing prev_sample
    sample = sample.to(torch.float32)

    if per_token_timesteps is not None:
        per_token_sigmas = per_token_timesteps / self.config.num_train_timesteps

        sigmas = self.sigmas[:, None, None]
        lower_mask = sigmas < per_token_sigmas[None] - 1e-6
        lower_sigmas = lower_mask * sigmas
        lower_sigmas, _ = lower_sigmas.max(dim=0)

        current_sigma = per_token_sigmas[..., None]
        next_sigma = lower_sigmas[..., None]
        dt = current_sigma - next_sigma
    else:
        if self.step_index is None:
            self._init_step_index(timestep)

        sigma_idx = self.step_index
        sigma = self.sigmas[sigma_idx]
        sigma_next = self.sigmas[sigma_idx + 1]

        current_sigma = sigma
        next_sigma = sigma_next
        dt = sigma_next - sigma

    if self.config.stochastic_sampling:
        x0 = sample - current_sigma * model_output
        noise = torch.randn_like(sample)
        prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise
    else:
        prev_sample = sample + dt * model_output

    # upon completion increase step index by one
    self._step_index += 1
    if per_token_timesteps is None:
        # Cast sample back to model compatible dtype
        prev_sample = prev_sample.to(model_output.dtype)

    if isinstance(prev_sample, torch.Tensor | float) and not return_dict:
        return (prev_sample, )

    return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)

def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor:
    r"""
    Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config
    value.

    Reference:
    https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51

    Args:
        t (`torch.Tensor`):
            A tensor of timesteps to be stretched and shifted.

    Returns:
        `torch.Tensor`:
            A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`.
    """
    one_minus_z = 1 - t
    scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal)
    stretched_t = 1 - (one_minus_z / scale_factor)
    return stretched_t