--- title: scheduling_flow_match_euler_discrete ---

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.init_logger

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.BaseScheduler

BaseScheduler(*args, **kwargs)

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.BaseScheduler.set_shift

set_shift(shift: float) -> None

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.BaseScheduler.set_timesteps

set_timesteps(*args, **kwargs) -> None

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.BaseScheduler.scale_model_input

scale_model_input(sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput

FlowMatchEulerDiscreteSchedulerOutput(prev_sample: torch.FloatTensor)

Output class for the scheduler's step function output.

Parameters

prev_sample: `torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images

Computed sample (x_{t-1}) of previous timestep. prev_sample should be used as next model input in the denoising loop.

Source
@dataclass
class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput):
    """
    Output class for the scheduler's `step` function output.

    Args:
        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
            denoising loop.
    """

    prev_sample: torch.FloatTensor

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler

FlowMatchEulerDiscreteScheduler(num_train_timesteps: int = 1000, shift: float = 1.0, use_dynamic_shifting: bool = False, base_shift: float | None = 0.5, max_shift: float | None = 1.15, base_image_seq_len: int | None = 256, max_image_seq_len: int | None = 4096, invert_sigmas: bool = False, shift_terminal: float | None = None, use_karras_sigmas: bool | None = False, use_exponential_sigmas: bool | None = False, use_beta_sigmas: bool | None = False, time_shift_type: str = 'exponential', stochastic_sampling: bool = False, final_sigmas_type: str = 'sigma_min', sigma_max: float | None = None, sigma_min: float | None = None, sigma_data: float | None = None, use_reference_discrete_timesteps: bool = False)

Euler scheduler.

This model inherits from [SchedulerMixin] and [ConfigMixin]. Check the superclass documentation for the generic methods the library implements for all schedulers such as loading and saving.

Parameters

num_train_timesteps: `int`, defaults to 1000

The number of diffusion steps to train the model.

shift: `float`, defaults to 1.0

The shift value for the timestep schedule.

use_dynamic_shifting: `bool`, defaults to False

Whether to apply timestep shifting on-the-fly based on the image resolution.

base_shift: `float`, defaults to 0.5

Value to stabilize image generation. Increasing base_shift reduces variation and image is more consistent with desired output.

max_shift: `float`, defaults to 1.15

Value change allowed to latent vectors. Increasing max_shift encourages more variation and image may be more exaggerated or stylized.

base_image_seq_len: `int`, defaults to 256

The base image sequence length.

max_image_seq_len: `int`, defaults to 4096

The maximum image sequence length.

invert_sigmas: `bool`, defaults to False

Whether to invert the sigmas.

shift_terminal: `float`, defaults to None

The end value of the shifted timestep schedule.

use_karras_sigmas: `bool`, defaults to False

Whether to use Karras sigmas for step sizes in the noise schedule during sampling.

use_exponential_sigmas: `bool`, defaults to False

Whether to use exponential sigmas for step sizes in the noise schedule during sampling.

use_beta_sigmas: `bool`, defaults to False

Whether to use beta sigmas for step sizes in the noise schedule during sampling.

time_shift_type: `str`, defaults to "exponential"

The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear".

stochastic_sampling: `bool`, defaults to False

Whether to use stochastic sampling.

final_sigmas_type: `str`, defaults to "sigma_min"

The type of final sigmas to use. Either "sigma_min" or "zero".

sigma_max: `float`, *optional*

The maximum sigma value for the noise schedule.

sigma_min: `float`, *optional*

The minimum sigma value for the noise schedule.

sigma_data: `float`, *optional*

The sigma data value for scaling.

use_reference_discrete_timesteps: `bool`, defaults to False

Some reference schedulers (e.g. MMAudio) construct the timestep schedule by linspacing num_inference_steps + 1 points from t_max to t_min and dropping the terminal point. Default (False) preserves the original np.linspace(t_max, t_min, num_inference_steps) (float64) behaviour used by every existing model. Enable this flag only when matching a reference scheduler that expects the +1 + drop-terminal construction.

Source
class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
    """
    Euler scheduler.

    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
    methods the library implements for all schedulers such as loading and saving.

    Args:
        num_train_timesteps (`int`, defaults to 1000):
            The number of diffusion steps to train the model.
        shift (`float`, defaults to 1.0):
            The shift value for the timestep schedule.
        use_dynamic_shifting (`bool`, defaults to False):
            Whether to apply timestep shifting on-the-fly based on the image resolution.
        base_shift (`float`, defaults to 0.5):
            Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent
            with desired output.
        max_shift (`float`, defaults to 1.15):
            Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be
            more exaggerated or stylized.
        base_image_seq_len (`int`, defaults to 256):
            The base image sequence length.
        max_image_seq_len (`int`, defaults to 4096):
            The maximum image sequence length.
        invert_sigmas (`bool`, defaults to False):
            Whether to invert the sigmas.
        shift_terminal (`float`, defaults to None):
            The end value of the shifted timestep schedule.
        use_karras_sigmas (`bool`, defaults to False):
            Whether to use Karras sigmas for step sizes in the noise schedule during sampling.
        use_exponential_sigmas (`bool`, defaults to False):
            Whether to use exponential sigmas for step sizes in the noise schedule during sampling.
        use_beta_sigmas (`bool`, defaults to False):
            Whether to use beta sigmas for step sizes in the noise schedule during sampling.
        time_shift_type (`str`, defaults to "exponential"):
            The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear".
        stochastic_sampling (`bool`, defaults to False):
            Whether to use stochastic sampling.
        final_sigmas_type (`str`, defaults to "sigma_min"):
            The type of final sigmas to use. Either "sigma_min" or "zero".
        sigma_max (`float`, *optional*):
            The maximum sigma value for the noise schedule.
        sigma_min (`float`, *optional*):
            The minimum sigma value for the noise schedule.
        sigma_data (`float`, *optional*):
            The sigma data value for scaling.
        use_reference_discrete_timesteps (`bool`, defaults to False):
            Some reference schedulers (e.g. MMAudio) construct the timestep
            schedule by linspacing `num_inference_steps + 1` points from
            `t_max` to `t_min` and dropping the terminal point. Default
            (`False`) preserves the original `np.linspace(t_max, t_min,
            num_inference_steps)` (float64) behaviour used by every existing
            model. Enable this flag only when matching a reference scheduler
            that expects the +1 + drop-terminal construction.
    """

    _compatibles: list[Any] = []
    order = 1

    @register_to_config
    def __init__(
        self,
        num_train_timesteps: int = 1000,
        shift: float = 1.0,
        use_dynamic_shifting: bool = False,
        base_shift: float | None = 0.5,
        max_shift: float | None = 1.15,
        base_image_seq_len: int | None = 256,
        max_image_seq_len: int | None = 4096,
        invert_sigmas: bool = False,
        shift_terminal: float | None = None,
        use_karras_sigmas: bool | None = False,
        use_exponential_sigmas: bool | None = False,
        use_beta_sigmas: bool | None = False,
        time_shift_type: str = "exponential",
        stochastic_sampling: bool = False,
        final_sigmas_type: str = "sigma_min",
        sigma_max: float | None = None,
        sigma_min: float | None = None,
        sigma_data: float | None = None,
        use_reference_discrete_timesteps: bool = False,
    ):
        if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1:
            raise ValueError(
                "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used."
            )
        if time_shift_type not in {"exponential", "linear"}:
            raise ValueError("`time_shift_type` must either be 'exponential' or 'linear'.")

        timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy()
        timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32)

        sigmas = timesteps / num_train_timesteps
        if not use_dynamic_shifting:
            # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
            sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)

        self.timesteps = sigmas * num_train_timesteps
        self.num_train_timesteps = num_train_timesteps

        self._step_index: int | None = None
        self._begin_index: int | None = None

        self._shift = shift

        self.sigmas = sigmas.to("cpu")  # to avoid too much CPU/GPU communication
        self.sigma_min = sigma_min if sigma_min is not None else self.sigmas[-1].item()
        self.sigma_max = self.sigmas[0].item()

        BaseScheduler.__init__(self)

    @property
    def shift(self) -> float:
        """
        The value used for shifting.
        """
        return self._shift

    @property
    def step_index(self) -> int | None:
        """
        The index counter for current timestep. It will increase 1 after each scheduler step.
        """
        return self._step_index

    @property
    def begin_index(self) -> int | None:
        """
        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
        """
        return self._begin_index

    # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
    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_shift(self, shift: float) -> None:
        self._shift = shift

    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 _sigma_to_t(self, sigma: float) -> float:
        return sigma * self.config.num_train_timesteps

    def time_shift(self, mu: float, sigma: float, t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray:
        if self.config.time_shift_type == "exponential":
            return self._time_shift_exponential(mu, sigma, t)
        elif self.config.time_shift_type == "linear":
            return self._time_shift_linear(mu, sigma, t)
        else:
            raise ValueError(f"Unknown time_shift_type: {self.config.time_shift_type}")

    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

    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 MMAudio) 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 index_for_timestep(self,
                           timestep: float | torch.FloatTensor,
                           schedule_timesteps: torch.Tensor | None = None) -> int:
        if schedule_timesteps is None:
            schedule_timesteps = self.timesteps

        indices = (schedule_timesteps == timestep).nonzero()

        # The sigma index that is taken for the **very** first `step`
        # is always the second index (or the last index if there is only 1)
        # This way we can ensure we don't accidentally skip a sigma in
        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
        pos = 1 if len(indices) > 1 else 0

        return indices[pos].item()

    def _init_step_index(self, timestep: float | torch.FloatTensor) -> None:
        if self.begin_index is None:
            if isinstance(timestep, torch.Tensor):
                timestep = timestep.to(self.timesteps.device)
            self._step_index = self.index_for_timestep(timestep)
        else:
            self._step_index = self._begin_index

    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)

    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
    def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
        """Constructs the noise schedule of Karras et al. (2022)."""

        # Hack to make sure that other schedulers which copy this function don't break
        # TODO: Add this logic to the other schedulers
        sigma_min = self.config.sigma_min if hasattr(self.config, "sigma_min") else None

        sigma_max = self.config.sigma_max if hasattr(self.config, "sigma_max") else None

        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()

        rho = 7.0  # 7.0 is the value used in the paper
        ramp = np.linspace(0, 1, num_inference_steps)
        min_inv_rho = sigma_min**(1 / rho)
        max_inv_rho = sigma_max**(1 / rho)
        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho))**rho
        return torch.from_numpy(sigmas).to(dtype=in_sigmas.dtype, device=in_sigmas.device)

    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential
    def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor:
        """Constructs an exponential noise schedule."""

        # Hack to make sure that other schedulers which copy this function don't break
        # TODO: Add this logic to the other schedulers
        sigma_min = self.config.sigma_min if hasattr(self.config, "sigma_min") else None

        sigma_max = self.config.sigma_max if hasattr(self.config, "sigma_max") else None

        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()

        sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps))
        return torch.from_numpy(sigmas).to(dtype=in_sigmas.dtype, device=in_sigmas.device)

    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta
    def _convert_to_beta(self,
                         in_sigmas: torch.Tensor,
                         num_inference_steps: int,
                         alpha: float = 0.6,
                         beta: float = 0.6) -> torch.Tensor:
        """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)"""

        # Hack to make sure that other schedulers which copy this function don't break
        # TODO: Add this logic to the other schedulers
        sigma_min = self.config.sigma_min if hasattr(self.config, "sigma_min") else None

        sigma_max = self.config.sigma_max if hasattr(self.config, "sigma_max") else None

        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()

        sigmas = np.array([
            sigma_min + (ppf * (sigma_max - sigma_min)) for ppf in
            [scipy.stats.beta.ppf(timestep, alpha, beta) for timestep in 1 - np.linspace(0, 1, num_inference_steps)]
        ])
        return torch.from_numpy(sigmas).to(dtype=in_sigmas.dtype, device=in_sigmas.device)

    def _time_shift_exponential(self, mu: float, sigma: float,
                                t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray:
        if isinstance(t, np.ndarray):
            return np.exp(mu) / (np.exp(mu) + (1 / t - 1)**sigma)
        else:
            return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma)

    def _time_shift_linear(self, mu: float, sigma: float, t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray:
        return mu / (mu + (1 / t - 1)**sigma)

    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_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:
        return sample

    def __len__(self) -> int:
        return self.config.num_train_timesteps

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.shift

shift: float

The value used for shifting.

Source
def shift(self) -> float:
    """
    The value used for shifting.
    """
    return self._shift

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.step_index

step_index: int | None

The index counter for current timestep. It will increase 1 after each scheduler step.

Source
def step_index(self) -> int | None:
    """
    The index counter for current timestep. It will increase 1 after each scheduler step.
    """
    return self._step_index

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.begin_index

begin_index: int | None

The index for the first timestep. It should be set from pipeline with set_begin_index method.

Source
def begin_index(self) -> int | None:
    """
    The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
    """
    return self._begin_index

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_begin_index

set_begin_index(begin_index: int = 0) -> None

Sets the begin index for the scheduler. This function should be run from pipeline before the inference.

Parameters

begin_index: `int`

The begin index for the scheduler.

Source
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

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_shift

set_shift(shift: float) -> None
Source
def set_shift(self, shift: float) -> None:
    self._shift = shift

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.scale_noise

scale_noise(sample: torch.FloatTensor, timestep: float | torch.FloatTensor, noise: torch.FloatTensor | None = None) -> torch.FloatTensor

Forward process in flow-matching

Parameters

sample: `torch.FloatTensor`

The input sample.

timestep: `int`, *optional*

The current timestep in the diffusion chain.

Returns

torch.FloatTensor

torch.FloatTensor: A scaled input sample.

Source
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

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.time_shift

time_shift(mu: float, sigma: float, t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray
Source
def time_shift(self, mu: float, sigma: float, t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray:
    if self.config.time_shift_type == "exponential":
        return self._time_shift_exponential(mu, sigma, t)
    elif self.config.time_shift_type == "linear":
        return self._time_shift_linear(mu, sigma, t)
    else:
        raise ValueError(f"Unknown time_shift_type: {self.config.time_shift_type}")

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.stretch_shift_to_terminal

stretch_shift_to_terminal(t: torch.Tensor) -> torch.Tensor

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

Parameters

t: `torch.Tensor`

A tensor of timesteps to be stretched and shifted.

Returns

torch.Tensor

torch.Tensor: A tensor of adjusted timesteps such that the final value equals self.config.shift_terminal.

Source
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

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps

set_timesteps(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).

Parameters

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.

Source
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 MMAudio) 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

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.index_for_timestep

index_for_timestep(timestep: float | torch.FloatTensor, schedule_timesteps: torch.Tensor | None = None) -> int
Source
def index_for_timestep(self,
                       timestep: float | torch.FloatTensor,
                       schedule_timesteps: torch.Tensor | None = None) -> int:
    if schedule_timesteps is None:
        schedule_timesteps = self.timesteps

    indices = (schedule_timesteps == timestep).nonzero()

    # The sigma index that is taken for the **very** first `step`
    # is always the second index (or the last index if there is only 1)
    # This way we can ensure we don't accidentally skip a sigma in
    # case we start in the middle of the denoising schedule (e.g. for image-to-image)
    pos = 1 if len(indices) > 1 else 0

    return indices[pos].item()

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.step

step(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).

Parameters

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

FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor, ...]

[~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.

Source
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)

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.add_noise

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

Parameters

clean_latent: torch.Tensor

the clean latent with shape [B, C, H, W], where B is batch_size or batch_size * num_frames

noise: torch.Tensor

the noise with shape [B, C, H, W]

timestep: torch.IntTensor

the timestep with shape [1] or [bs * num_frames] or [bs, num_frames]

Returns

torch.Tensor

the corrupted latent with shape [B, C, H, W]

Source
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)

fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.scale_model_input

scale_model_input(sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor
Source
def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor:
    return sample