Skip to content

scheduling_cosmos25_distilled

TrigFlow sampler for the distilled Cosmos Predict2.5 checkpoint.

Classes

fastvideo.models.schedulers.scheduling_cosmos25_distilled.Cosmos25DistilledScheduler 

Cosmos25DistilledScheduler(num_train_timesteps: int = 1000, sigma_data: float = 1.0)

Bases: SchedulerMixin, ConfigMixin, BaseScheduler

Official four-step TrigFlow/x0 sampler for Cosmos Predict2.5.

The distilled student predicts the rectified-flow network output. The scheduler applies the student's TrigFlow preconditioning, reconstructs x0, and re-noises it with the initial noise for the next student evaluation. This fixed-noise update is intentionally different from stochastic rCM sampling.

Source code in fastvideo/models/schedulers/scheduling_cosmos25_distilled.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
@register_to_config
def __init__(
    self,
    num_train_timesteps: int = 1000,
    sigma_data: float = 1.0,
) -> None:
    if sigma_data <= 0:
        raise ValueError(f"sigma_data must be positive, got {sigma_data}")

    self.num_train_timesteps = num_train_timesteps
    self.sigma_data = float(sigma_data)
    self.timesteps = torch.empty(0, dtype=torch.float64)
    self.trigflow_timesteps = torch.empty(0, dtype=torch.float64)
    self.sigmas = torch.empty(0, dtype=torch.float64)
    self._step_index: int | None = None
    self._begin_index: int | None = None
    self._initial_noise: torch.Tensor | None = None
    BaseScheduler.__init__(self)

Methods:

fastvideo.models.schedulers.scheduling_cosmos25_distilled.Cosmos25DistilledScheduler.set_shift 
set_shift(shift: float) -> None

The checkpoint's fixed distilled schedule does not use flow shift.

Source code in fastvideo/models/schedulers/scheduling_cosmos25_distilled.py 
77
78
def set_shift(self, shift: float) -> None:
    """The checkpoint's fixed distilled schedule does not use flow shift."""

fastvideo.models.schedulers.scheduling_cosmos25_distilled.Cosmos25DistilledSchedulerOutput dataclass 

Cosmos25DistilledSchedulerOutput(prev_sample: Tensor, pred_original_sample: Tensor)

Bases: BaseOutput

Output of one distilled Cosmos Predict2.5 sampling step.