44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
"""Sampling utilities for Rectified Flow inference."""
|
|
|
|
import math
|
|
|
|
import torch
|
|
|
|
|
|
def time_shift(mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Apply time shift to timesteps for flow matching schedule.
|
|
|
|
Args:
|
|
mu: Time shift parameter (controls schedule steepness)
|
|
sigma: Sigma parameter (typically 1.0)
|
|
t: Timestep tensor with values in (0, 1]
|
|
|
|
Returns:
|
|
Shifted timesteps
|
|
"""
|
|
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
|
|
|
|
|
def get_rf_schedule(num_steps: int, mu: float = 1.5, reverse: bool = True) -> list[float]:
|
|
"""
|
|
Generate timestep schedule for Rectified Flow sampling.
|
|
|
|
Creates a shifted linear schedule that provides better sample quality
|
|
by spending more time at higher noise levels.
|
|
|
|
Args:
|
|
num_steps: Number of sampling steps
|
|
mu: Time shift parameter (higher = more time at high noise)
|
|
reverse: If True, returns schedule from t=0 to t=1 (for denoising)
|
|
|
|
Returns:
|
|
List of timesteps of length num_steps + 1
|
|
"""
|
|
if reverse:
|
|
mu = -mu
|
|
timesteps = torch.linspace(1, 0, num_steps + 1)
|
|
timesteps = time_shift(mu, 1.0, timesteps)
|
|
timesteps = timesteps.tolist()
|
|
return timesteps[::-1] if reverse else timesteps
|