PiflowScheduler

PiflowScheduler is the few-step scheduler of the distilled Kandinsky 6 checkpoints, both the text/image-to-video-and-audio model and the video super-resolution model. It implements π-Flow: the transformer predicts n_grid denoised estimates per latent channel at a small number of grid points, and the scheduler integrates a network-free policy between them.

The reference implementation can be found at Lakonik/LakonLab.

PiflowScheduler

class diffusers.PiflowScheduler

< source >

( num_train_timesteps: int = 1000shift: float = 5.0n_grid: int = 10eps: float = 1e-06final_step_size_scale: float = 0.5num_policy_substeps: int = 128 )

Parameters

num_train_timesteps (int, optional, defaults to 1000) — Number of training diffusion steps.
shift (float, optional, defaults to 5.0) — Flow-matching timestep shift.
n_grid (int, optional, defaults to 10) — Number of predictions in the widened model output.
eps (float, optional, defaults to 1e-6) — Minimum timestep and policy denominator.
final_step_size_scale (float, optional, defaults to 0.5) — Relative size of the final raw-timestep segment.
num_policy_substeps (int, optional, defaults to 128) — Maximum policy integration substeps per raw-timestep unit.

Few-step PiFlow scheduler for widened-output diffusion transformers.

PiFlow evaluates the denoising model at a small number of grid points and integrates a network-free policy between those evaluations. The scheduler is intended for distilled Kandinsky 6 checkpoints, including the main video/audio model and the video super-resolution model. Their model output contains n_grid predictions per sample channel.

This scheduler inherits from SchedulerMixin and ConfigMixin. Check the superclass documentation for the generic methods implemented for all schedulers (loading, saving, etc.).

set_timesteps

< source >

( num_inference_steps: int | None = Nonedevice: typing.Union[str, torch.device, NoneType] = Nonesigmas: list[float] | None = Nonemu: float | None = Nonetimesteps: list[float] | None = None )

Parameters

num_inference_steps (int) — Number of model evaluations.
device (str or torch.device, optional) — Device for the schedule.
sigmas (list[float], optional) — Unsupported custom sigma schedule.
mu (float, optional) — Unsupported dynamic-shift parameter.
timesteps (list[float], optional) — Unsupported custom timestep schedule.

Set the distilled PiFlow timestep schedule.

step

< source >

( model_output: FloatTensortimestep: typing.Union[float, torch.FloatTensor]sample: FloatTensorreturn_dict: bool = True ) → [PiflowSchedulerOutput] or tuple

Parameters

model_output (torch.FloatTensor) — Widened model output containing n_grid predictions per sample channel.
timestep (float or torch.FloatTensor) — Current scheduler timestep.
sample (torch.FloatTensor) — Current noisy sample.
return_dict (bool, optional, defaults to True) — Whether to return a [PiflowSchedulerOutput].

Returns

[PiflowSchedulerOutput] or tuple

Updated sample.

Advance one step by integrating the PiFlow policy.