scheduling_flow_match_euler_discrete
fastvideo.models.schedulers.scheduling_flow_match_euler_discrete.init_logger
init_logger(name: str) -> _FastvideoLoggerThe 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) -> Nonefastvideo.models.schedulers.scheduling_flow_match_euler_discrete.BaseScheduler.set_timesteps
set_timesteps(*args, **kwargs) -> Nonefastvideo.models.schedulers.scheduling_flow_match_euler_discrete.BaseScheduler.scale_model_input
scale_model_input(sample: torch.Tensor, timestep: int | None = None) -> torch.Tensorfastvideo.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 imagesComputed sample
(x_{t-1})of previous timestep.prev_sampleshould 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</code></pre></details>
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 1000The number of diffusion steps to train the model.
shift: `float`, defaults to 1.0The shift value for the timestep schedule.
use_dynamic_shifting: `bool`, defaults to FalseWhether to apply timestep shifting on-the-fly based on the image resolution.
base_shift: `float`, defaults to 0.5Value to stabilize image generation. Increasing base_shift reduces variation and image is more consistent
with desired output.
max_shift: `float`, defaults to 1.15Value 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 256The base image sequence length.
max_image_seq_len: `int`, defaults to 4096The maximum image sequence length.
invert_sigmas: `bool`, defaults to FalseWhether to invert the sigmas.
shift_terminal: `float`, defaults to NoneThe end value of the shifted timestep schedule.
use_karras_sigmas: `bool`, defaults to FalseWhether to use Karras sigmas for step sizes in the noise schedule during sampling.
use_exponential_sigmas: `bool`, defaults to FalseWhether to use exponential sigmas for step sizes in the noise schedule during sampling.
use_beta_sigmas: `bool`, defaults to FalseWhether 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 FalseWhether 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 FalseSome 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</code></pre></details>
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</code></pre></details>
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.FloatTensortorch.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</code></pre></details>
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.Tensortorch.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</code></pre></details>
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</code></pre></details>
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()</code></pre></details>
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.0Scaling 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)</code></pre></details>
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.Tensorthe clean latent with shape [B, C, H, W],
where B is batch_size or batch_size * num_frames
noise: torch.Tensorthe noise with shape [B, C, H, W]
timestep: torch.IntTensorthe timestep with shape [1] or [bs * num_frames] or [bs, num_frames]
Returns
torch.Tensorthe 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)</code></pre></details>
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