Skip to content

scheduling_self_forcing_flow_match

Classes

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)

Bases: BaseScheduler, ConfigMixin, SchedulerMixin

Source code in fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py
@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)

Methods:

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 code in fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py
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 code in fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py
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_weight
training_weight(timestep)
Input
  • timestep: the timestep with shape [B*T]

Output: the corresponding weighting [B*T]

Source code in fastvideo/models/schedulers/scheduling_self_forcing_flow_match.py
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.SelfForcingFlowMatchSchedulerOutput

Bases: BaseOutput

Output class for the scheduler's step function output.

Parameters:

Name Type Description Default
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.

required

Functions: