--- title: scheduling_self_forcing_flow_match ---

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.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_self_forcing_flow_match.BaseScheduler

BaseScheduler(*args, **kwargs)

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.BaseScheduler.set_shift

set_shift(shift: float) -> None

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.BaseScheduler.set_timesteps

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

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.BaseScheduler.scale_model_input

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

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.SelfForcingFlowMatchSchedulerOutput

SelfForcingFlowMatchSchedulerOutput()

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
class SelfForcingFlowMatchSchedulerOutput(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_self_forcing_flow_match.SelfForcingFlowMatchScheduler

SelfForcingFlowMatchScheduler(num_inference_steps = 100, num_train_timesteps = 1000, shift = 3.0, sigma_max = 1.0, sigma_min = 0.003 / 1.002, inverse_timesteps = False, extra_one_step = False, reverse_sigmas = False, training = False)
Source
class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):

    config_name = "scheduler_config.json"
    order = 1

    @register_to_config
    def __init__(self,
                 num_inference_steps=100,
                 num_train_timesteps=1000,
                 shift=3.0,
                 sigma_max=1.0,
                 sigma_min=0.003 / 1.002,
                 inverse_timesteps=False,
                 extra_one_step=False,
                 reverse_sigmas=False,
                 training=False):
        self.num_train_timesteps = num_train_timesteps
        self.shift = shift
        self.sigma_max = sigma_max
        self.sigma_min = sigma_min
        self.inverse_timesteps = inverse_timesteps
        self.extra_one_step = extra_one_step
        self.reverse_sigmas = reverse_sigmas
        self.set_timesteps(num_inference_steps, training=training)

    def set_timesteps(self,
                      num_inference_steps=100,
                      denoising_strength=1.0,
                      training=False,
                      return_dict=False,
                      **kwargs):
        sigma_start = self.sigma_min + \
            (self.sigma_max - self.sigma_min) * denoising_strength
        if self.extra_one_step:
            self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
        else:
            self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps)
        if self.inverse_timesteps:
            self.sigmas = torch.flip(self.sigmas, dims=[0])
        self.sigmas = self.shift * self.sigmas / \
            (1 + (self.shift - 1) * self.sigmas)
        if self.reverse_sigmas:
            self.sigmas = 1 - self.sigmas
        self.timesteps = self.sigmas * self.num_train_timesteps
        if training:
            x = self.timesteps
            y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps)**2)
            y_shifted = y - y.min()
            bsmntw_weighing = y_shifted * \
                (num_inference_steps / y_shifted.sum())
            self.linear_timesteps_weights = bsmntw_weighing

    def step(self,
             model_output: torch.FloatTensor,
             timestep: torch.FloatTensor,
             sample: torch.FloatTensor,
             to_final=False,
             return_dict=False,
             **kwargs):
        if timestep.ndim == 2:
            timestep = timestep.flatten(0, 1)
        elif timestep.ndim == 0:
            # handles the case where timestep is a scalar, this occurs when we
            # use this scheduler for ODE trajectory
            timestep = timestep.unsqueeze(0)

        self.sigmas = self.sigmas.to(model_output.device)
        self.timesteps = self.timesteps.to(model_output.device)
        timestep = timestep.to(model_output.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)
        if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
            sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0
        else:
            sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
        prev_sample = sample + model_output * (sigma_ - sigma)
        if isinstance(prev_sample, torch.Tensor | float) and not return_dict:
            return (prev_sample, )
        return SelfForcingFlowMatchSchedulerOutput(prev_sample=prev_sample)

    @staticmethod
    def calculate_alpha_beta_high(sigma, sigma_bound):
        alpha = (1 - sigma) / (1 - sigma_bound)
        beta = torch.sqrt(sigma**2 - (alpha * sigma_bound)**2)
        return alpha, beta

    def add_noise(self, original_samples, noise, timestep):
        """
        Diffusion forward corruption process.
        Input:
            - clean_latent: the clean latent with shape [B*T, C, H, W]
            - noise: the noise with shape [B*T, C, H, W]
            - timestep: the timestep with shape [B*T]
        Output: the corrupted latent with shape [B*T, C, H, W]
        """
        if timestep.ndim == 2:
            timestep = timestep.flatten(0, 1)
        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) * original_samples + sigma * noise
        return sample.type_as(noise)

    def add_noise_high(self, original_samples, noise, timestep, boundary_timestep):
        """
        Diffusion forward corruption process.
        Input:
            - clean_latent: the clean latent with shape [B*T, C, H, W]
            - noise: the noise with shape [B*T, C, H, W]
            - timestep: the timestep with shape [B*T]
        Output: the corrupted latent with shape [B*T, C, H, W]
        """
        if timestep.ndim == 2:
            timestep = timestep.flatten(0, 1)
        if boundary_timestep.ndim == 2:
            boundary_timestep = boundary_timestep.flatten(0, 1)
        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)

        boundary_timestep_id = torch.argmin((self.timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1)
        sigma_boundary = self.sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1)
        alpha, beta = self.calculate_alpha_beta_high(sigma, sigma_boundary)
        sample = alpha * original_samples + beta * noise
        return sample.type_as(noise)

    def training_target(self, sample, noise, timestep):
        target = noise - sample
        return target

    def training_weight(self, timestep):
        """
        Input:
            - timestep: the timestep with shape [B*T]
        Output: the corresponding weighting [B*T]
        """
        if timestep.ndim == 2:
            timestep = timestep.flatten(0, 1)
        self.linear_timesteps_weights = self.linear_timesteps_weights.to(timestep.device)
        timestep_id = torch.argmin((self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0)
        weights = self.linear_timesteps_weights[timestep_id]
        return weights

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

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

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.SelfForcingFlowMatchScheduler.set_timesteps

set_timesteps(num_inference_steps = 100, denoising_strength = 1.0, training = False, return_dict = False, **kwargs)
Source
def set_timesteps(self,
                  num_inference_steps=100,
                  denoising_strength=1.0,
                  training=False,
                  return_dict=False,
                  **kwargs):
    sigma_start = self.sigma_min + \
        (self.sigma_max - self.sigma_min) * denoising_strength
    if self.extra_one_step:
        self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
    else:
        self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps)
    if self.inverse_timesteps:
        self.sigmas = torch.flip(self.sigmas, dims=[0])
    self.sigmas = self.shift * self.sigmas / \
        (1 + (self.shift - 1) * self.sigmas)
    if self.reverse_sigmas:
        self.sigmas = 1 - self.sigmas
    self.timesteps = self.sigmas * self.num_train_timesteps
    if training:
        x = self.timesteps
        y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps)**2)
        y_shifted = y - y.min()
        bsmntw_weighing = y_shifted * \
            (num_inference_steps / y_shifted.sum())
        self.linear_timesteps_weights = bsmntw_weighing

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.SelfForcingFlowMatchScheduler.step

step(model_output: torch.FloatTensor, timestep: torch.FloatTensor, sample: torch.FloatTensor, to_final = False, return_dict = False, **kwargs)
Source
def step(self,
         model_output: torch.FloatTensor,
         timestep: torch.FloatTensor,
         sample: torch.FloatTensor,
         to_final=False,
         return_dict=False,
         **kwargs):
    if timestep.ndim == 2:
        timestep = timestep.flatten(0, 1)
    elif timestep.ndim == 0:
        # handles the case where timestep is a scalar, this occurs when we
        # use this scheduler for ODE trajectory
        timestep = timestep.unsqueeze(0)

    self.sigmas = self.sigmas.to(model_output.device)
    self.timesteps = self.timesteps.to(model_output.device)
    timestep = timestep.to(model_output.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)
    if to_final or (timestep_id + 1 >= len(self.timesteps)).any():
        sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0
    else:
        sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1)
    prev_sample = sample + model_output * (sigma_ - sigma)
    if isinstance(prev_sample, torch.Tensor | float) and not return_dict:
        return (prev_sample, )
    return SelfForcingFlowMatchSchedulerOutput(prev_sample=prev_sample)

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.SelfForcingFlowMatchScheduler.calculate_alpha_beta_high

calculate_alpha_beta_high(sigma, sigma_bound)
Source
@staticmethod
def calculate_alpha_beta_high(sigma, sigma_bound):
    alpha = (1 - sigma) / (1 - sigma_bound)
    beta = torch.sqrt(sigma**2 - (alpha * sigma_bound)**2)
    return alpha, beta

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.SelfForcingFlowMatchScheduler.add_noise

add_noise(original_samples, noise, timestep)

Diffusion forward corruption process. Input: - clean_latent: the clean latent with shape [BT, C, H, W] - noise: the noise with shape [BT, C, H, W] - timestep: the timestep with shape [BT] Output: the corrupted latent with shape [BT, C, H, W]

Source
def add_noise(self, original_samples, noise, timestep):
    """
    Diffusion forward corruption process.
    Input:
        - clean_latent: the clean latent with shape [B*T, C, H, W]
        - noise: the noise with shape [B*T, C, H, W]
        - timestep: the timestep with shape [B*T]
    Output: the corrupted latent with shape [B*T, C, H, W]
    """
    if timestep.ndim == 2:
        timestep = timestep.flatten(0, 1)
    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) * original_samples + sigma * noise
    return sample.type_as(noise)

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.SelfForcingFlowMatchScheduler.add_noise_high

add_noise_high(original_samples, noise, timestep, boundary_timestep)

Diffusion forward corruption process. Input: - clean_latent: the clean latent with shape [BT, C, H, W] - noise: the noise with shape [BT, C, H, W] - timestep: the timestep with shape [BT] Output: the corrupted latent with shape [BT, C, H, W]

Source
def add_noise_high(self, original_samples, noise, timestep, boundary_timestep):
    """
    Diffusion forward corruption process.
    Input:
        - clean_latent: the clean latent with shape [B*T, C, H, W]
        - noise: the noise with shape [B*T, C, H, W]
        - timestep: the timestep with shape [B*T]
    Output: the corrupted latent with shape [B*T, C, H, W]
    """
    if timestep.ndim == 2:
        timestep = timestep.flatten(0, 1)
    if boundary_timestep.ndim == 2:
        boundary_timestep = boundary_timestep.flatten(0, 1)
    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)

    boundary_timestep_id = torch.argmin((self.timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1)
    sigma_boundary = self.sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1)
    alpha, beta = self.calculate_alpha_beta_high(sigma, sigma_boundary)
    sample = alpha * original_samples + beta * noise
    return sample.type_as(noise)

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.SelfForcingFlowMatchScheduler.training_target

training_target(sample, noise, timestep)
Source
def training_target(self, sample, noise, timestep):
    target = noise - sample
    return target

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.SelfForcingFlowMatchScheduler.training_weight

training_weight(timestep)

Input

Output: the corresponding weighting [B*T]

Source
def training_weight(self, timestep):
    """
    Input:
        - timestep: the timestep with shape [B*T]
    Output: the corresponding weighting [B*T]
    """
    if timestep.ndim == 2:
        timestep = timestep.flatten(0, 1)
    self.linear_timesteps_weights = self.linear_timesteps_weights.to(timestep.device)
    timestep_id = torch.argmin((self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0)
    weights = self.linear_timesteps_weights[timestep_id]
    return weights

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.SelfForcingFlowMatchScheduler.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

fastvideo.models.schedulers.scheduling_self_forcing_flow_match.SelfForcingFlowMatchScheduler.set_shift

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