scheduling_rcm ¶

def add_noise(
    self,
    original_samples: torch.Tensor,
    noise: torch.Tensor,
    timesteps: torch.IntTensor,
) -> torch.Tensor:
    """
    Add noise to samples (forward diffusion process).

    Not typically used for rCM inference, but provided for API compatibility.
    """
    raise NotImplementedError(
        "add_noise is not implemented for RCMScheduler. "
        "Use scale_noise for initializing the sampling process."
    )

def scale_model_input(
    self,
    sample: torch.Tensor,
    timestep: int | None = None,
) -> torch.Tensor:
    """
    rCM doesn't scale model input, returns sample as-is.
    """
    return sample

def scale_noise(
    self,
    sample: torch.FloatTensor,
    timestep: torch.FloatTensor | None = None,
    noise: torch.FloatTensor | None = None,
) -> torch.FloatTensor:
    """
    Scale initial noise for rCM sampling.

    In rCM, initial noise is scaled by the first timestep (raw sigma):
        x_0 = noise * t_steps[0]

    Args:
        sample: Not used (for API compatibility)
        timestep: Not used (for API compatibility)
        noise: The noise tensor to scale

    Returns:
        Scaled noise tensor ready for sampling
    """
    if noise is None:
        raise ValueError("noise must be provided for rCM scale_noise")

    # Use raw sigma (not scaled timestep) for initial noise scaling
    t_initial = self.sigmas[0]
    return noise.to(torch.float64) * t_initial

def set_begin_index(self, begin_index: int = 0) -> None:
    """
    Sets the begin index for the scheduler.

    Args:
        begin_index (`int`):
            The begin index for the scheduler.
    """
    self._begin_index = begin_index

def set_shift(self, shift: float) -> None:
    """
    rCM doesn't use shift parameter, but required by BaseScheduler.
    """
    pass

def set_timesteps(
    self,
    num_inference_steps: int,
    device: str | torch.device | None = None,
    sigma_max: float | None = None,
) -> None:
    """
    Sets the discrete timesteps used for the rCM sampling process.

    The timesteps are computed using TrigFlow → RectifiedFlow conversion:
    1. Start with atan(sigma_max) and intermediate values
    2. Convert via: t = sin(t) / (cos(t) + sin(t))

    Args:
        num_inference_steps (`int`):
            The number of diffusion steps (1-4 for rCM).
        device (`str` or `torch.device`, *optional*):
            The device to move timesteps to.
        sigma_max (`float`, *optional*):
            Override the initial sigma value.
    """
    if num_inference_steps < 1 or num_inference_steps > 4:
        logger.warning(
            "rCM is optimized for 1-4 steps, got %d steps. "
            "Performance may be suboptimal.", num_inference_steps
        )

    self.num_inference_steps = num_inference_steps

    if sigma_max is not None:
        self.sigma_max = sigma_max

    # Build timestep schedule
    mid_t = self._mid_timesteps[:num_inference_steps - 1]

    # TrigFlow timesteps: [atan(sigma_max), mid_t..., 0]
    t_steps = torch.tensor(
        [math.atan(self.sigma_max), *mid_t, 0],
        dtype=torch.float64,
        device=device,
    )

    # Convert TrigFlow → RectifiedFlow: t = sin(t) / (cos(t) + sin(t))
    t_steps = torch.sin(t_steps) / (torch.cos(t_steps) + torch.sin(t_steps))

    # Store raw sigmas for use in step() formula
    self.sigmas = t_steps.clone()

    # Scale timesteps by 1000 for model input (as per TurboDiffusion)
    self.timesteps = t_steps * 1000

    self._step_index = None
    self._begin_index = None

    logger.debug("rCM timesteps (scaled): %s", self.timesteps.tolist())
    logger.debug("rCM sigmas (raw): %s", self.sigmas.tolist())

def step(
    self,
    model_output: torch.FloatTensor,
    timestep: int | torch.Tensor,
    sample: torch.FloatTensor,
    generator: torch.Generator | None = None,
    return_dict: bool = True,
) -> RCMSchedulerOutput | tuple[torch.FloatTensor, ...]:
    """
    Predict the sample from the previous timestep using rCM update rule.

    The rCM update formula is:
        x_{t+1} = (1 - t_next) * (x_t - t_cur * v_pred) + t_next * noise

    Args:
        model_output (`torch.FloatTensor`):
            The velocity prediction from the model (v_pred).
        timestep (`int` or `torch.Tensor`):
            The current timestep index (not the actual timestep value).
            For rCM, this should be the index into self.timesteps.
        sample (`torch.FloatTensor`):
            Current sample x_t.
        generator (`torch.Generator`, *optional*):
            Random number generator for noise.
        return_dict (`bool`):
            Whether to return RCMSchedulerOutput or tuple.

    Returns:
        `RCMSchedulerOutput` or `tuple`:
            The denoised sample for the next step.
    """
    if self._step_index is None:
        self._init_step_index()

    assert self._step_index is not None

    # Get current and next sigma values (raw, unscaled) for rCM formula
    # Note: self.timesteps is scaled by 1000 for model input,
    # but we need raw values for the step formula
    t_cur = self.sigmas[self._step_index]
    # On the final step, t_next should be 0 (fully denoised)
    if self._step_index + 1 < len(self.sigmas):
        t_next = self.sigmas[self._step_index + 1]
    else:
        t_next = torch.tensor(0.0, device=sample.device, dtype=torch.float64)

    # Ensure we're working in float64 for precision
    sample = sample.to(torch.float64)
    model_output = model_output.to(torch.float64)

    # rCM update: x = (1 - t_next) * (x - t_cur * v_pred) + t_next * noise
    x_denoised = sample - t_cur * model_output

    # Generate noise for SDE sampling
    if isinstance(generator, list):
        generator = generator[0]
    noise = torch.randn(
        sample.shape,
        dtype=torch.float32,
        device="cpu",
        generator=generator,
    ).to(sample.device).to(torch.float64)

    prev_sample = (1 - t_next) * x_denoised + t_next * noise

    # Increment step counter
    self._step_index += 1

    # Cast back to model output dtype
    prev_sample = prev_sample.to(model_output.dtype)

    if not return_dict:
        return (prev_sample,)

    return RCMSchedulerOutput(prev_sample=prev_sample)