Diffusers

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.41.0).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

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.

PiflowSchedulerOutput

class diffusers.schedulers.scheduling_piflow.PiflowSchedulerOutput

< source >

( prev_sample: Tensor )

Parameters

prev_sample (torch.Tensor) — Computed sample at the next PiFlow grid point. Should be used as the next denoising input.

Output class for the scheduler’s step function output.

Update on GitHub

←MiniMaxH3Scheduler PNDMScheduler→