--- title: kandinsky5 ---

fastvideo.pipelines.stages.kandinsky5.LocalAttention

LocalAttention(num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, default_backend: AttentionBackendEnum | None = None, **extra_impl_args)

Attention layer.

fastvideo.pipelines.stages.kandinsky5.LocalAttention.forward

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor

Apply local attention between query, key and value tensors.

Parameters

q: torch.Tensor

Query tensor of shape [batch_size, seq_len, num_heads, head_dim]

k: torch.Tensor

Key tensor of shape [batch_size, seq_len, num_heads, head_dim]

v: torch.Tensor

Value tensor of shape [batch_size, seq_len, num_heads, head_dim]

Returns

torch.Tensor

torch.Tensor: Output tensor after local attention

fastvideo.pipelines.stages.kandinsky5.NablaAttentionMetadataBuilder

NablaAttentionMetadataBuilder()

fastvideo.pipelines.stages.kandinsky5.NablaAttentionMetadataBuilder.prepare

prepare() -> None

fastvideo.pipelines.stages.kandinsky5.NablaAttentionMetadataBuilder.build

build(current_timestep: int, sta_mask: torch.Tensor, P: float, visual_shape: tuple[int, int, int], **kwargs: Any) -> NablaAttentionMetadata

fastvideo.pipelines.stages.kandinsky5.component_attention_backend

component_attention_backend(component: object) -> AttentionBackendEnum | _NoRequest

Read back the decision :func:record_resolved_attention_backend wrote.

Returns NO_REQUEST unless the component recorded a concrete backend, so a caller that passes this through only overrides the ambient fallback when there is a real decision to override it with.

NO_REQUEST rather than None for the no-decision case is deliberate, and cannot be derived from the attribute's presence: ModelConfig declares _resolved_attention_backend as a field defaulting to None, so the attribute always exists and a getattr default can never fire. Worse, record_resolved_attention_backend writes None whenever no scope is active, so "resolved to automatic selection" and "never recorded" are the same stored value. Neither state should suppress the environment variable at a call site that previously honoured it, and collapsing both to NO_REQUEST keeps that behavior identical.

fastvideo.pipelines.stages.kandinsky5.get_local_torch_device

get_local_torch_device() -> torch.device

Return the torch device for the current rank.

fastvideo.pipelines.stages.kandinsky5.FastVideoArgs

FastVideoArgs(model_path: str, mode: ExecutionMode = ExecutionMode.INFERENCE, workload_type: WorkloadType = WorkloadType.T2V, distributed_executor_backend: str = 'mp', enable_streaming_ipc_queues: bool = False, ray_placement_group: PlacementGroup | None = None, ray_runtime_env: RuntimeEnv | None = None, inference_mode: bool = True, trust_remote_code: bool = False, revision: str | None = None, num_gpus: int = 1, tp_size: int = -1, sp_size: int = -1, hsdp_replicate_dim: int = 1, hsdp_shard_dim: int = -1, dist_timeout: int | None = None, pipeline_config: PipelineConfig = PipelineConfig(), preprocess_config: PreprocessConfig | None = None, lora_path: str | None = None, lora_nickname: str = 'default', lora_strength: float = 1.0, lora_target_modules: list[str] | None = None, output_type: str = 'pil', attention_backend: str | None = None, dit_cpu_offload: bool = True, use_fsdp_inference: bool = False, dit_layerwise_offload: bool = True, text_encoder_cpu_offload: bool = True, image_encoder_cpu_offload: bool = True, vae_cpu_offload: bool = True, pin_cpu_memory: bool = True, h3_sequential_load: bool | None = None, h3_encoder_split: bool | None = None, h3_encoder_nodes: int = -1, h3_encoder_workers: int = 0, video_decode_backend: str = 'h3-vae', taeh3_checkpoint: str | None = None, taeh3_chunk_size: int = 5, h3_vae_nvfp4_decoder: str | None = None, lazy_module_load: bool | None = None, vae_parallel_decode: bool = False, vae_parallel_encode: bool = False, vae_parallel_decode_strategy: str | None = None, enable_torch_compile: bool = False, regional_compile: bool = False, enable_torch_compile_text_encoder: bool = False, enable_torch_compile_vae: bool = False, enable_torch_compile_audio_vae: bool = False, torch_compile_kwargs: dict[str, Any] = dict(), torch_compile_kwargs_dit: dict[str, Any] = dict(), torch_compile_kwargs_text_encoder: dict[str, Any] = dict(), torch_compile_kwargs_vae: dict[str, Any] = dict(), torch_compile_kwargs_audio_vae: dict[str, Any] = dict(), inference_torch_compile: bool = False, h3_vsa_fp4: dict[str, Any] = dict(), disable_autocast: bool = False, VSA_sparsity: float | None = None, VSA_tile_size: int | None = None, moba_config_path: str | None = None, moba_config: dict[str, Any] = dict(), master_port: int | None = None, enable_stage_verification: bool = True, prompt_txt: str | None = None, ltx2_vae_tiling: bool | None = None, ltx2_vae_spatial_tile_size_in_pixels: int | None = None, ltx2_vae_spatial_tile_overlap_in_pixels: int | None = None, ltx2_vae_temporal_tile_size_in_frames: int | None = None, ltx2_vae_temporal_tile_overlap_in_frames: int | None = None, ltx2_initial_latent_path: str | None = None, ltx2_audio_latent_path: str | None = None, refine_enabled: bool | None = None, refine_upsampler_path: str | None = None, refine_transformer_path: str | None = None, refine_lora_path: str | None = None, refine_num_inference_steps: int | None = None, refine_guidance_scale: float | None = None, refine_add_noise: bool | None = None, refine_noise_path: str | None = None, refine_audio_noise_path: str | None = None, ltx2_refine_enabled: bool = False, ltx2_refine_upsampler_path: str | None = None, ltx2_refine_transformer_path: str | None = None, ltx2_refine_lora_path: str | None = None, ltx2_refine_num_inference_steps: int = 3, ltx2_refine_guidance_scale: float = 1.0, ltx2_refine_add_noise: bool = True, ltx2_refine_noise_path: str | None = None, ltx2_refine_audio_noise_path: str | None = None, ltx2_legacy_native_noise_order: bool = False, ltx2_use_distilled_sigmas: bool = True, model_paths: dict[str, str] = dict(), model_loaded: dict[str, bool] = (lambda: {'transformer': True, 'vae': True, 'upsampler': True})(), override_text_encoder_safetensors: str | None = None, override_text_encoder_quant: QuantizationMethods = None, transformer_quant: Any | None = None, override_transformer_cls_name: str | None = None, init_weights_from_safetensors: str = '', init_weights_from_safetensors_2: str = '', override_pipeline_cls_name: str | None = None, boundary_ratio: float | None = 0.875)

fastvideo.pipelines.stages.kandinsky5.FastVideoArgs.is_output_rank

is_output_rank: bool

Whether this process may materialize user-facing outputs.

This is runtime state assigned by an SPMD executor, not a public configuration field, so dataclass serialization and typed config parsing deliberately ignore it.

fastvideo.pipelines.stages.kandinsky5.FastVideoArgs.add_cli_args

add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser

fastvideo.pipelines.stages.kandinsky5.FastVideoArgs.from_cli_args

from_cli_args(args: argparse.Namespace) -> FastVideoArgs

fastvideo.pipelines.stages.kandinsky5.FastVideoArgs.from_kwargs

from_kwargs(**kwargs: Any) -> FastVideoArgs

fastvideo.pipelines.stages.kandinsky5.FastVideoArgs.check_fastvideo_args

check_fastvideo_args() -> None

Validate inference arguments for consistency

fastvideo.pipelines.stages.kandinsky5.FastVideoArgs.finalize_device_offload_policy

finalize_device_offload_policy(device_id: int = 0) -> bool

Apply device-local memory policy, then resolve incompatible modes.

fastvideo.pipelines.stages.kandinsky5.FastVideoArgs.disable_offload_on_unified_memory

disable_offload_on_unified_memory(device_id: int = 0, *, offload_flag: str | None = None) -> bool

Disable host offload after a worker has selected its device.

CUDA's unified-memory probe reads runtime device properties and may initialize a CUDA context. Callers must therefore use this only inside a device-owning process, after selecting and binding device_id. Returning the classification lets direct component-loader callers apply the same policy to explicit per-call overrides. When offload_flag is given, the return value says whether this policy covers that component role.

fastvideo.pipelines.stages.kandinsky5.set_forward_context

set_forward_context(current_timestep, attn_metadata, forward_batch: Optional[ForwardBatch] = None)

A context manager that stores the current forward context, can be attention metadata, etc. Here we can inject common logic for every model forward pass.

fastvideo.pipelines.stages.kandinsky5.init_logger

init_logger(name: str) -> _FastvideoLogger

The main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.

fastvideo.pipelines.stages.kandinsky5.TransformerLoader

TransformerLoader(device = None)

Loader for transformer.

fastvideo.pipelines.stages.kandinsky5.TransformerLoader.load

load(model_path: str, fastvideo_args: FastVideoArgs)

Load the transformer based on the model path, and inference args.

fastvideo.pipelines.stages.kandinsky5.VAELoader

VAELoader(device = None)

Loader for VAE.

fastvideo.pipelines.stages.kandinsky5.VAELoader.load

load(model_path: str, fastvideo_args: FastVideoArgs)

Load the VAE based on the model path, and inference args.

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler

FlowMatchEulerDiscreteScheduler(num_train_timesteps: int = 1000, shift: float = 1.0, use_dynamic_shifting: bool = False, base_shift: float | None = 0.5, max_shift: float | None = 1.15, base_image_seq_len: int | None = 256, max_image_seq_len: int | None = 4096, invert_sigmas: bool = False, shift_terminal: float | None = None, use_karras_sigmas: bool | None = False, use_exponential_sigmas: bool | None = False, use_beta_sigmas: bool | None = False, time_shift_type: str = 'exponential', stochastic_sampling: bool = False, final_sigmas_type: str = 'sigma_min', sigma_max: float | None = None, sigma_min: float | None = None, sigma_data: float | None = None, use_reference_discrete_timesteps: bool = False)

Euler scheduler.

This model inherits from [SchedulerMixin] and [ConfigMixin]. Check the superclass documentation for the generic methods the library implements for all schedulers such as loading and saving.

Parameters

num_train_timesteps: `int`, defaults to 1000

The number of diffusion steps to train the model.

shift: `float`, defaults to 1.0

The shift value for the timestep schedule.

use_dynamic_shifting: `bool`, defaults to False

Whether to apply timestep shifting on-the-fly based on the image resolution.

base_shift: `float`, defaults to 0.5

Value to stabilize image generation. Increasing base_shift reduces variation and image is more consistent with desired output.

max_shift: `float`, defaults to 1.15

Value change allowed to latent vectors. Increasing max_shift encourages more variation and image may be more exaggerated or stylized.

base_image_seq_len: `int`, defaults to 256

The base image sequence length.

max_image_seq_len: `int`, defaults to 4096

The maximum image sequence length.

invert_sigmas: `bool`, defaults to False

Whether to invert the sigmas.

shift_terminal: `float`, defaults to None

The end value of the shifted timestep schedule.

use_karras_sigmas: `bool`, defaults to False

Whether to use Karras sigmas for step sizes in the noise schedule during sampling.

use_exponential_sigmas: `bool`, defaults to False

Whether to use exponential sigmas for step sizes in the noise schedule during sampling.

use_beta_sigmas: `bool`, defaults to False

Whether to use beta sigmas for step sizes in the noise schedule during sampling.

time_shift_type: `str`, defaults to "exponential"

The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear".

stochastic_sampling: `bool`, defaults to False

Whether to use stochastic sampling.

final_sigmas_type: `str`, defaults to "sigma_min"

The type of final sigmas to use. Either "sigma_min" or "zero".

sigma_max: `float`, *optional*

The maximum sigma value for the noise schedule.

sigma_min: `float`, *optional*

The minimum sigma value for the noise schedule.

sigma_data: `float`, *optional*

The sigma data value for scaling.

use_reference_discrete_timesteps: `bool`, defaults to False

Some reference schedulers (e.g. MMAudio) construct the timestep schedule by linspacing num_inference_steps + 1 points from t_max to t_min and dropping the terminal point. Default (False) preserves the original np.linspace(t_max, t_min, num_inference_steps) (float64) behaviour used by every existing model. Enable this flag only when matching a reference scheduler that expects the +1 + drop-terminal construction.

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler.shift

shift: float

The value used for shifting.

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler.step_index

step_index: int | None

The index counter for current timestep. It will increase 1 after each scheduler step.

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler.begin_index

begin_index: int | None

The index for the first timestep. It should be set from pipeline with set_begin_index method.

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler.set_begin_index

set_begin_index(begin_index: int = 0) -> None

Sets the begin index for the scheduler. This function should be run from pipeline before the inference.

Parameters

begin_index: `int`

The begin index for the scheduler.

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler.set_shift

set_shift(shift: float) -> None

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler.scale_noise

scale_noise(sample: torch.FloatTensor, timestep: float | torch.FloatTensor, noise: torch.FloatTensor | None = None) -> torch.FloatTensor

Forward process in flow-matching

Parameters

sample: `torch.FloatTensor`

The input sample.

timestep: `int`, *optional*

The current timestep in the diffusion chain.

Returns

torch.FloatTensor

torch.FloatTensor: A scaled input sample.

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler.time_shift

time_shift(mu: float, sigma: float, t: torch.Tensor | np.ndarray) -> torch.Tensor | np.ndarray

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler.stretch_shift_to_terminal

stretch_shift_to_terminal(t: torch.Tensor) -> torch.Tensor

Stretches and shifts the timestep schedule to ensure it terminates at the configured shift_terminal config value.

Reference: https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51

Parameters

t: `torch.Tensor`

A tensor of timesteps to be stretched and shifted.

Returns

torch.Tensor

torch.Tensor: A tensor of adjusted timesteps such that the final value equals self.config.shift_terminal.

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler.set_timesteps

set_timesteps(num_inference_steps: int | None = None, device: str | torch.device = None, sigmas: list[float] | None = None, mu: float | None = None, timesteps: list[float] | None = None) -> None

Sets the discrete timesteps used for the diffusion chain (to be run before inference).

Parameters

num_inference_steps: `int`, *optional*

The number of diffusion steps used when generating samples with a pre-trained model.

device: `str` or `torch.device`, *optional*

The device to which the timesteps should be moved to. If None, the timesteps are not moved.

sigmas: `List[float]`, *optional*

Custom values for sigmas to be used for each diffusion step. If None, the sigmas are computed automatically.

mu: `float`, *optional*

Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep shifting.

timesteps: `List[float]`, *optional*

Custom values for timesteps to be used for each diffusion step. If None, the timesteps are computed automatically.

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler.index_for_timestep

index_for_timestep(timestep: float | torch.FloatTensor, schedule_timesteps: torch.Tensor | None = None) -> int

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler.step

step(model_output: torch.FloatTensor, timestep: int | torch.Tensor, sample: torch.FloatTensor, s_churn: float = 0.0, s_tmin: float = 0.0, s_tmax: float = float('inf'), s_noise: float = 1.0, generator: torch.Generator | None = None, per_token_timesteps: torch.Tensor | None = None, return_dict: bool = True) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor, ...]

Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned model outputs (most often the predicted noise).

Parameters

model_output: `torch.FloatTensor`

The direct output from learned diffusion model.

timestep: `int` or `torch.Tensor`

The current discrete timestep in the diffusion chain.

sample: `torch.FloatTensor`

A current instance of a sample created by the diffusion process.

s_churn: `float`
s_tmin: `float`
s_tmax: `float`
s_noise: `float`, defaults to 1.0

Scaling factor for noise added to the sample.

generator: `torch.Generator`, *optional*

A random number generator.

per_token_timesteps: `torch.Tensor`, *optional*

The timesteps for each token in the sample.

return_dict: `bool`

Whether or not to return a [~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput] or tuple.

Returns

FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor, ...]

[~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput] or tuple: If return_dict is True, [~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput] is returned, otherwise a tuple is returned where the first element is the sample tensor.

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler.add_noise

add_noise(clean_latent: torch.Tensor, noise: torch.Tensor, timestep: torch.IntTensor) -> torch.Tensor

Parameters

clean_latent: torch.Tensor

the clean latent with shape [B, C, H, W], where B is batch_size or batch_size * num_frames

noise: torch.Tensor

the noise with shape [B, C, H, W]

timestep: torch.IntTensor

the timestep with shape [1] or [bs * num_frames] or [bs, num_frames]

Returns

torch.Tensor

the corrupted latent with shape [B, C, H, W]

fastvideo.pipelines.stages.kandinsky5.FlowMatchEulerDiscreteScheduler.scale_model_input

scale_model_input(sample: torch.Tensor, timestep: int | None = None) -> torch.Tensor

fastvideo.pipelines.stages.kandinsky5.pred_noise_to_pred_video

pred_noise_to_pred_video(pred_noise: torch.Tensor, noise_input_latent: torch.Tensor, timestep: torch.Tensor, scheduler: Any) -> torch.Tensor

Convert predicted noise to clean latent.

pred_noise: the predicted noise with shape [B, C, H, W] where B is batch_size or batch_size * num_frames noise_input_latent: the noisy latent with shape [B, C, H, W], timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames] scheduler: the scheduler

Returns

torch.Tensor

the predicted video with shape [B, C, H, W]

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE

ParallelTiledVAE(config: VAEConfig, **kwargs)

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE.to

to(device) -> ParallelTiledVAE

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE.encode

encode(x: torch.Tensor) -> torch.Tensor

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE.decode

decode(z: torch.Tensor) -> torch.Tensor

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE.blend_v

blend_v(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE.blend_h

blend_h(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE.blend_t

blend_t(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE.spatial_tiled_encode

spatial_tiled_encode(x: torch.Tensor) -> torch.Tensor

Encode a batch of images using a tiled encoder.

Parameters

x: `torch.Tensor`

Input batch of videos.

Returns

torch.Tensor

torch.Tensor: The latent representation of the encoded videos.

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE.parallel_tiled_decode

parallel_tiled_decode(z: torch.FloatTensor) -> torch.FloatTensor

Parallel version of tiled_decode that distributes both temporal and spatial computation across GPUs

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE.spatial_tiled_decode

spatial_tiled_decode(z: torch.Tensor) -> torch.Tensor

Decode a batch of images using a tiled decoder.

Parameters

z: `torch.Tensor`

Input batch of latent vectors.

Returns

torch.Tensor

torch.Tensor: The decoded images.

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE.tiled_encode

tiled_encode(x: torch.Tensor) -> torch.Tensor

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE.tiled_decode

tiled_decode(z: torch.Tensor) -> torch.Tensor

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE.enable_tiling

enable_tiling(tile_sample_min_height: int | None = None, tile_sample_min_width: int | None = None, tile_sample_min_num_frames: int | None = None, tile_sample_stride_height: int | None = None, tile_sample_stride_width: int | None = None, tile_sample_stride_num_frames: int | None = None, blend_num_frames: int | None = None, use_tiling: bool | None = None, use_temporal_tiling: bool | None = None, use_parallel_tiling: bool | None = None) -> None

Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow processing larger images.

Parameters

tile_sample_min_height: `int`, *optional*

The minimum height required for a sample to be separated into tiles across the height dimension.

tile_sample_min_width: `int`, *optional*

The minimum width required for a sample to be separated into tiles across the width dimension.

tile_sample_min_num_frames: `int`, *optional*

The minimum number of frames required for a sample to be separated into tiles across the frame dimension.

tile_sample_stride_height: `int`, *optional*

The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are no tiling artifacts produced across the height dimension.

tile_sample_stride_width: `int`, *optional*

The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling artifacts produced across the width dimension.

tile_sample_stride_num_frames: `int`, *optional*

The stride between two consecutive frame tiles. This is to ensure that there are no tiling artifacts produced across the frame dimension.

fastvideo.pipelines.stages.kandinsky5.ParallelTiledVAE.disable_tiling

disable_tiling() -> None

Disable tiled VAE decoding. If enable_tiling was previously enabled, this method will go back to computing decoding in one step.

fastvideo.pipelines.stages.kandinsky5.normalize

normalize(images: np.ndarray | torch.Tensor) -> np.ndarray | torch.Tensor

Normalize an image array to [-1,1].

Parameters

images: `np.ndarray` or `torch.Tensor`

The image array to normalize.

Returns

np.ndarray | torch.Tensor

np.ndarray or torch.Tensor: The normalized image array.

fastvideo.pipelines.stages.kandinsky5.numpy_to_pt

numpy_to_pt(images: np.ndarray) -> torch.Tensor

Convert a NumPy image to a PyTorch tensor.

Parameters

images: `np.ndarray`

The NumPy image array to convert to PyTorch format.

Returns

torch.Tensor

torch.Tensor: A PyTorch tensor representation of the images.

fastvideo.pipelines.stages.kandinsky5.pil_to_numpy

pil_to_numpy(images: list[PIL.Image.Image] | PIL.Image.Image) -> np.ndarray

Convert a PIL image or a list of PIL images to NumPy arrays.

Parameters

images: `PIL.Image.Image` or `List[PIL.Image.Image]`

The PIL image or list of images to convert to NumPy format.

Returns

np.ndarray

np.ndarray: A NumPy array representation of the images.

fastvideo.pipelines.stages.kandinsky5.resize

resize(image: PIL.Image.Image | np.ndarray | torch.Tensor, height: int, width: int, resize_mode: str = 'default', resample: str = 'lanczos') -> PIL.Image.Image | np.ndarray | torch.Tensor

Resize image.

Parameters

image: `PIL.Image.Image`, `np.ndarray` or `torch.Tensor`

The image input, can be a PIL image, numpy array or pytorch tensor.

height: `int`

The height to resize to.

width: `int`

The width to resize to.

resize_mode: `str`, *optional*, defaults to `default`

The resize mode to use, can be one of default or fill. If default, will resize the image to fit within the specified width and height, and it may not maintaining the original aspect ratio. If fill, will resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, filling empty with data from image. If crop, will resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, cropping the excess. Note that resize_mode fill and crop are only supported for PIL image input.

Returns

PIL.Image.Image | np.ndarray | torch.Tensor

PIL.Image.Image, np.ndarray or torch.Tensor: The resized image.

fastvideo.pipelines.stages.kandinsky5.ForwardBatch

ForwardBatch(data_type: str, generator: torch.Generator | list[torch.Generator] | None = None, image_path: str | None = None, image_embeds: list[torch.Tensor] = list(), pil_image: torch.Tensor | PIL.Image.Image | None = None, last_image: torch.Tensor | PIL.Image.Image | None = None, references: list[Any] | None = None, preprocessed_image: torch.Tensor | None = None, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = None, prompt_path: str | None = None, output_path: str = 'outputs/', output_video_name: str | None = None, video_path: str | None = None, video_latent: torch.Tensor | None = None, pose_video_path: str | None = None, face_video_path: str | None = None, background_video_path: str | None = None, mask_video_path: str | None = None, animate_mode: str | None = None, pose_latents: torch.Tensor | None = None, face_pixel_values: torch.Tensor | None = None, audio_path: str | None = None, audio_embeds: torch.Tensor | None = None, t_thresh: float = 0.5, num_cond_frames: int = 0, prompt_embeds: list[torch.Tensor] = list(), negative_prompt_embeds: list[torch.Tensor] | None = None, prompt_attention_mask: list[torch.Tensor] | None = None, negative_attention_mask: list[torch.Tensor] | None = None, clip_embedding_pos: list[torch.Tensor] | None = None, clip_embedding_neg: list[torch.Tensor] | None = None, max_sequence_length: int | None = None, prompt_template: dict[str, Any] | None = None, do_classifier_free_guidance: bool = False, use_embedded_guidance: bool = False, true_cfg_scale: float = 1.0, batch_size: int | None = None, num_videos_per_prompt: int = 1, seed: int | None = None, seeds: list[int] | None = None, is_prompt_processed: bool = False, latents: torch.Tensor | None = None, audio_latents: torch.Tensor | None = None, lq_latents: torch.Tensor | None = None, raw_latent_shape: tuple[int, ...] | None = None, noise_pred: torch.Tensor | None = None, image_latent: torch.Tensor | None = None, first_frame_latent: torch.Tensor | None = None, helios_latent_chunks: list[torch.Tensor] | None = None, mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, grid_sizes: torch.Tensor | None = None, num_iterations: int | None = None, use_base_model: bool = False, pose: str | None = None, c2ws_plucker_emb: torch.Tensor | None = None, action_path: str | None = None, trajectory_type: str | None = None, movement_distance: float | None = None, camera_rotation: str | None = None, height_latents: list[int] | int | None = None, width_latents: list[int] | int | None = None, num_frames: list[int] | int = 1, height: list[int] | int | None = None, width: list[int] | int | None = None, height_sr: list[int] | int | None = None, width_sr: list[int] | int | None = None, fps: list[int] | int | None = None, timesteps: torch.Tensor | None = None, timestep: torch.Tensor | float | int | None = None, step_index: int | None = None, boundary_ratio: float | None = None, num_inference_steps: int = 50, num_inference_steps_sr: int = 50, guidance_scale: float = 1.0, batch_cfg: bool = False, guidance_scale_2: float | None = None, guidance_rescale: float = 0.0, eta: float = 0.0, sigmas: list[float] | None = None, pyramid_num_inference_steps_list: list[int] | None = None, history_sizes: list[int] | None = None, num_latent_frames_per_chunk: int = 9, keep_first_frame: bool = True, is_skip_first_chunk: bool = False, use_zero_init: bool = True, zero_steps: int = 1, is_amplify_first_chunk: bool = False, enable_teacache: bool = False, ltx2_cfg_scale_video: float = 1.0, ltx2_cfg_scale_audio: float = 1.0, ltx2_modality_scale_video: float = 1.0, ltx2_modality_scale_audio: float = 1.0, ltx2_rescale_scale: float = 0.0, ltx2_stg_scale_video: float = 0.0, ltx2_stg_scale_audio: float = 0.0, ltx2_stg_blocks_video: list[int] = list(), ltx2_stg_blocks_audio: list[int] = list(), ltx2_use_ancestral_sampler: bool = False, ltx2_images: list[tuple[str, int, float]] | None = None, ltx2_image_crf: float = 33.0, ltx2_conditioning_latent_stage1: torch.Tensor | None = None, ltx2_conditioning_latent_stage2: torch.Tensor | None = None, ltx2_video_conditions: list[tuple[list[str], int, float]] | None = None, audio_start_in_s: float | None = None, audio_end_in_s: float | None = None, init_audio: Any = None, init_audio_strength: float | None = None, init_noise_level: float | None = None, inpaint_audio: Any = None, inpaint_mask: Any = None, n_tokens: int | None = None, extra_step_kwargs: dict[str, Any] = dict(), modules: dict[str, Any] = dict(), output: torch.Tensor | None = None, return_trajectory_latents: bool = False, return_trajectory_decoded: bool = False, trajectory_timesteps: list[torch.Tensor] | None = None, trajectory_latents: torch.Tensor | None = None, trajectory_decoded: list[torch.Tensor] | None = None, continuation_state: ContinuationState | None = None, return_continuation_state: bool = False, extra: dict[str, Any] = dict(), save_video: bool = True, return_frames: bool = False, is_cfg_negative: bool = False, VSA_sparsity: float = 0.0, logging_info: PipelineLoggingInfo = PipelineLoggingInfo(), *, return_samples: bool = False)

Complete state passed through the pipeline execution.

This dataclass contains all information needed during the diffusion pipeline execution, allowing methods to update specific components without needing to manage numerous individual parameters.

fastvideo.pipelines.stages.kandinsky5.PipelineStage

PipelineStage()

Abstract base class for all pipeline stages.

A pipeline stage represents a discrete step in the diffusion process that can be composed with other stages to create a complete pipeline. Each stage is responsible for a specific part of the process, such as prompt encoding, latent preparation, etc.

fastvideo.pipelines.stages.kandinsky5.PipelineStage.verify_input

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify the input for the stage.

Example

from fastvideo.pipelines.stages.validators import V, VerificationResult

def verify_input(self, batch, fastvideo_args): result = VerificationResult() result.add_check("height", batch.height, V.positive_int_divisible(8)) result.add_check("width", batch.width, V.positive_int_divisible(8)) result.add_check("image_latent", batch.image_latent, V.is_tensor) return result

Parameters

batch: ForwardBatch

The current batch information.

fastvideo_args: FastVideoArgs

The inference arguments.

Returns

VerificationResult

A VerificationResult containing the verification status.

fastvideo.pipelines.stages.kandinsky5.PipelineStage.verify_output

verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify the output for the stage.

Parameters

batch: ForwardBatch

The current batch information.

fastvideo_args: FastVideoArgs

The inference arguments.

Returns

VerificationResult

A VerificationResult containing the verification status.

fastvideo.pipelines.stages.kandinsky5.PipelineStage.device

device: torch.device

Get the device for this stage.

fastvideo.pipelines.stages.kandinsky5.PipelineStage.set_logging

set_logging(enable: bool)

Enable or disable logging for this stage.

Parameters

enable: bool

Whether to enable logging.

fastvideo.pipelines.stages.kandinsky5.PipelineStage.forward

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch

Forward pass of the stage's processing.

This method should be implemented by subclasses to provide the forward processing logic for the stage.

Parameters

batch: ForwardBatch

The current batch information.

fastvideo_args: FastVideoArgs

The inference arguments.

Returns

ForwardBatch

The updated batch information after this stage's processing.

fastvideo.pipelines.stages.kandinsky5.DecodingStage

DecodingStage(vae, pipeline = None)

Stage for decoding latent representations into pixel space.

This stage handles the decoding of latent representations into the final output format (e.g., pixel values).

fastvideo.pipelines.stages.kandinsky5.DecodingStage.verify_input

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify decoding stage inputs.

fastvideo.pipelines.stages.kandinsky5.DecodingStage.verify_output

verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify decoding stage outputs.

fastvideo.pipelines.stages.kandinsky5.DecodingStage.decode

decode(latents: torch.Tensor, fastvideo_args: FastVideoArgs) -> torch.Tensor

Decode latent representations into pixel space using VAE.

Parameters

latents: torch.Tensor

Input latent tensor with shape (batch, channels, frames, height_latents, width_latents)

fastvideo_args: FastVideoArgs

Configuration containing: - disable_autocast: Whether to disable automatic mixed precision (default: False) - pipeline_config.vae_precision: VAE computation precision ("fp32", "fp16", "bf16") - pipeline_config.vae_tiling: Whether to enable VAE tiling for memory efficiency

Returns

torch.Tensor

Decoded video tensor with shape (batch, channels, frames, height, width),

torch.Tensor

normalized to [0, 1] range and moved to CPU as float32

fastvideo.pipelines.stages.kandinsky5.DecodingStage.streaming_decode

streaming_decode(latents: torch.Tensor, fastvideo_args: FastVideoArgs, cache: list[torch.Tensor | None] | None = None, is_first_chunk: bool = False) -> tuple[torch.Tensor, list[torch.Tensor | None]]

Decode latent representations into pixel space using VAE with streaming cache.

Parameters

latents: torch.Tensor

Input latent tensor with shape (batch, channels, frames, height_latents, width_latents)

fastvideo_args: FastVideoArgs

Configuration object.

cache: list[torch.Tensor | None] | None

VAE cache from previous call, or None to initialize a new cache.

is_first_chunk: bool

Whether this is the first chunk.

Returns

tuple[torch.Tensor, list[torch.Tensor | None]]

A tuple of (decoded_frames, updated_cache).

fastvideo.pipelines.stages.kandinsky5.DecodingStage.forward

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch

Decode latent representations into pixel space.

This method processes the batch through the VAE decoder, converting latent representations to pixel-space video/images. It also optionally decodes trajectory latents for visualization purposes.

Parameters

batch: ForwardBatch

The current batch containing: - latents: Tensor to decode (batch, channels, frames, height_latents, width_latents) - return_trajectory_decoded (optional): Flag to decode trajectory latents - trajectory_latents (optional): Latents at different timesteps - trajectory_timesteps (optional): Corresponding timesteps

fastvideo_args: FastVideoArgs

Configuration containing: - output_type: "latent" to skip decoding, otherwise decode to pixels - vae_cpu_offload: Whether to offload VAE to CPU after decoding - model_loaded: Track VAE loading state - model_paths: Path to VAE model if loading needed

Returns

ForwardBatch

Modified batch with: - output: Decoded frames (batch, channels, frames, height, width) as CPU float32 - trajectory_decoded (if requested): List of decoded frames per timestep

fastvideo.pipelines.stages.kandinsky5.EncodingStage

EncodingStage(vae: ParallelTiledVAE)

Stage for encoding pixel space representations into latent space.

This stage handles the encoding of pixel-space video/images into latent representations for further processing in the diffusion pipeline.

fastvideo.pipelines.stages.kandinsky5.EncodingStage.verify_input

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify encoding stage inputs.

fastvideo.pipelines.stages.kandinsky5.EncodingStage.verify_output

verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify encoding stage outputs.

fastvideo.pipelines.stages.kandinsky5.EncodingStage.forward

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch

Encode pixel space representations into latent space.

Parameters

batch: ForwardBatch

The current batch information.

fastvideo_args: FastVideoArgs

The inference arguments.

Returns

ForwardBatch

The batch with encoded latents.

fastvideo.pipelines.stages.kandinsky5.V

V()

Common validators for pipeline stages.

fastvideo.pipelines.stages.kandinsky5.V.not_none

not_none(value: Any) -> bool

Check if value is not None.

fastvideo.pipelines.stages.kandinsky5.V.positive_int

positive_int(value: Any) -> bool

Check if value is a positive integer.

fastvideo.pipelines.stages.kandinsky5.V.non_negative_int

non_negative_int(value: Any) -> bool

Check if value is a non-negative integer (allows 0).

fastvideo.pipelines.stages.kandinsky5.V.positive_float

positive_float(value: Any) -> bool

Check if value is a positive float.

fastvideo.pipelines.stages.kandinsky5.V.non_negative_float

non_negative_float(value: Any) -> bool

Check if value is a non-negative float.

fastvideo.pipelines.stages.kandinsky5.V.divisible_by

divisible_by(value: Any, divisor: int) -> bool

Check if value is divisible by divisor.

fastvideo.pipelines.stages.kandinsky5.V.is_tensor

is_tensor(value: Any) -> bool

Check if value is a torch tensor and doesn't contain NaN values.

fastvideo.pipelines.stages.kandinsky5.V.tensor_with_dims

tensor_with_dims(value: Any, dims: int) -> bool

Check if value is a tensor with specific dimensions and no NaN values.

fastvideo.pipelines.stages.kandinsky5.V.tensor_min_dims

tensor_min_dims(value: Any, min_dims: int) -> bool

Check if value is a tensor with at least min_dims dimensions and no NaN values.

fastvideo.pipelines.stages.kandinsky5.V.tensor_shape_matches

tensor_shape_matches(value: Any, expected_shape: tuple) -> bool

Check if tensor shape matches expected shape (None for any size) and no NaN values.

fastvideo.pipelines.stages.kandinsky5.V.list_not_empty

list_not_empty(value: Any) -> bool

Check if value is a non-empty list.

fastvideo.pipelines.stages.kandinsky5.V.list_length

list_length(value: Any, length: int) -> bool

Check if list has specific length.

fastvideo.pipelines.stages.kandinsky5.V.list_min_length

list_min_length(value: Any, min_length: int) -> bool

Check if list has at least min_length items.

fastvideo.pipelines.stages.kandinsky5.V.string_not_empty

string_not_empty(value: Any) -> bool

Check if value is a non-empty string.

fastvideo.pipelines.stages.kandinsky5.V.string_or_list_strings

string_or_list_strings(value: Any) -> bool

Check if value is a string or list of strings.

fastvideo.pipelines.stages.kandinsky5.V.bool_value

bool_value(value: Any) -> bool

Check if value is a boolean.

fastvideo.pipelines.stages.kandinsky5.V.generator_or_list_generators

generator_or_list_generators(value: Any) -> bool

Check if value is a Generator or list of Generators.

fastvideo.pipelines.stages.kandinsky5.V.is_list

is_list(value: Any) -> bool

Check if value is a list (can be empty).

fastvideo.pipelines.stages.kandinsky5.V.is_tuple

is_tuple(value: Any) -> bool

Check if value is a tuple.

fastvideo.pipelines.stages.kandinsky5.V.none_or_tensor

none_or_tensor(value: Any) -> bool

Check if value is None or a tensor without NaN values.

fastvideo.pipelines.stages.kandinsky5.V.list_of_tensors_with_dims

list_of_tensors_with_dims(value: Any, dims: int) -> bool

Check if value is a non-empty list where all items are tensors with specific dimensions and no NaN values.

fastvideo.pipelines.stages.kandinsky5.V.list_of_tensors

list_of_tensors(value: Any) -> bool

Check if value is a non-empty list where all items are tensors without NaN values.

fastvideo.pipelines.stages.kandinsky5.V.list_of_tensors_with_min_dims

list_of_tensors_with_min_dims(value: Any, min_dims: int) -> bool

Check if value is a non-empty list where all items are tensors with at least min_dims dimensions and no NaN values.

fastvideo.pipelines.stages.kandinsky5.V.none_or_tensor_with_dims

none_or_tensor_with_dims(dims: int) -> Callable[[Any], bool]

Return a validator that checks if value is None or a tensor with specific dimensions and no NaN values.

fastvideo.pipelines.stages.kandinsky5.V.none_or_list

none_or_list(value: Any) -> bool

Check if value is None or a list.

fastvideo.pipelines.stages.kandinsky5.V.none_or_positive_int

none_or_positive_int(value: Any) -> bool

Check if value is None or a positive integer.

fastvideo.pipelines.stages.kandinsky5.V.with_dims

with_dims(dims: int) -> Callable[[Any], bool]

Return a validator that checks if tensor has specific dimensions and no NaN values.

fastvideo.pipelines.stages.kandinsky5.V.min_list_length

min_list_length(min_length: int) -> Callable[[Any], bool]

Return a validator that checks if a list has at least min_length items.

fastvideo.pipelines.stages.kandinsky5.V.min_dims

min_dims(min_dims: int) -> Callable[[Any], bool]

Return a validator that checks if tensor has at least min_dims dimensions and no NaN values.

fastvideo.pipelines.stages.kandinsky5.V.divisible

divisible(divisor: int) -> Callable[[Any], bool]

Return a validator that checks if value is divisible by divisor.

fastvideo.pipelines.stages.kandinsky5.V.positive_int_divisible

positive_int_divisible(divisor: int) -> Callable[[Any], bool]

Return a validator that checks if value is a positive integer divisible by divisor.

fastvideo.pipelines.stages.kandinsky5.V.list_of_tensors_dims

list_of_tensors_dims(dims: int) -> Callable[[Any], bool]

Return a validator that checks if value is a list of tensors with specific dimensions and no NaN values.

fastvideo.pipelines.stages.kandinsky5.V.list_of_tensors_min_dims

list_of_tensors_min_dims(min_dims: int) -> Callable[[Any], bool]

Return a validator that checks if value is a list of tensors with at least min_dims dimensions and no NaN values.

fastvideo.pipelines.stages.kandinsky5.VerificationResult

VerificationResult()

Wrapper class for stage verification results.

fastvideo.pipelines.stages.kandinsky5.VerificationResult.add_check

add_check(field_name: str, value: Any, validators: Callable[[Any], bool] | list[Callable[[Any], bool]]) -> VerificationResult

Add a validation check for a field.

Parameters

field_name: str

Name of the field being checked

value: Any

The actual value to validate

validators: Callable[[Any], bool] | list[Callable[[Any], bool]]

Single validation function or list of validation functions. Each function will be called with the value as its first argument.

Returns

VerificationResult

Self for method chaining

Examples

Single validator

result.add_check("tensor", my_tensor, V.is_tensor)

Multiple validators (all must pass)

result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])

Using partial functions for parameters

result.add_check("height", batch.height, [V.not_none, V.divisible(8)])

fastvideo.pipelines.stages.kandinsky5.VerificationResult.is_valid

is_valid() -> bool

Check if all validations passed.

fastvideo.pipelines.stages.kandinsky5.VerificationResult.get_failed_fields

get_failed_fields() -> list[str]

Get list of fields that failed validation.

fastvideo.pipelines.stages.kandinsky5.VerificationResult.get_detailed_failures

get_detailed_failures() -> dict[str, list[ValidationFailure]]

Get detailed failure information for each failed field.

fastvideo.pipelines.stages.kandinsky5.VerificationResult.get_failure_summary

get_failure_summary() -> str

Get a comprehensive summary of all validation failures.

fastvideo.pipelines.stages.kandinsky5.VerificationResult.to_dict

to_dict() -> dict

Convert to dictionary for backward compatibility.

fastvideo.pipelines.stages.kandinsky5.AttentionBackendEnum

AttentionBackendEnum()

fastvideo.pipelines.stages.kandinsky5.get_mixed_precision_state

get_mixed_precision_state() -> MixedPrecisionState

Get the current mixed precision state.

fastvideo.pipelines.stages.kandinsky5.Kandinsky5LatentPreparationStage

Kandinsky5LatentPreparationStage(scheduler, transformer)
Source
class Kandinsky5LatentPreparationStage(PipelineStage):

    def __init__(self, scheduler, transformer) -> None:
        super().__init__()
        self.scheduler = scheduler
        self.transformer = transformer

    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
        if batch.height is None or batch.width is None:
            raise ValueError("height and width must be provided for Kandinsky5.")
        height = int(batch.height)
        width = int(batch.width)
        num_frames = int(batch.num_frames)

        temporal_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
        spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
        patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size

        if num_frames % temporal_ratio != 1:
            num_frames = num_frames // temporal_ratio * temporal_ratio + 1
            batch.num_frames = num_frames

        required_divisor_h = spatial_ratio * patch_size[1]
        required_divisor_w = spatial_ratio * patch_size[2]
        if height % required_divisor_h != 0 or width % required_divisor_w != 0:
            raise ValueError(f"Kandinsky5 height must be divisible by {required_divisor_h} and width by "
                             f"{required_divisor_w}; "
                             f"got height={height}, width={width}.")

        # NABLA sparse attention (Pro checkpoints) reshapes the post-patch grid
        # into 8x8 blocks; validate here instead of crashing mid-denoise after
        # all the encoding work is done.
        arch_cfg = getattr(self.transformer, "config", None) or fastvideo_args.pipeline_config.dit_config.arch_config
        if getattr(arch_cfg, "attention_type", "regular") == "nabla":
            nabla_divisor_h = required_divisor_h * 8
            nabla_divisor_w = required_divisor_w * 8
            if height % nabla_divisor_h != 0 or width % nabla_divisor_w != 0:
                raise ValueError(f"Kandinsky5 NABLA checkpoints require height divisible by {nabla_divisor_h} and "
                                 f"width divisible by {nabla_divisor_w}; "
                                 f"got height={height}, width={width}.")

        if isinstance(batch.prompt, list):
            batch_size = len(batch.prompt)
        elif batch.prompt is not None:
            batch_size = 1
        else:
            batch_size = batch.prompt_embeds[0].shape[0]
        batch_size *= batch.num_videos_per_prompt

        dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
        device = get_local_torch_device()
        num_latent_frames = (num_frames - 1) // temporal_ratio + 1
        num_channels = getattr(
            self.transformer,
            "in_visual_dim",
            fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
        )
        shape = (
            batch_size,
            num_latent_frames,
            height // spatial_ratio,
            width // spatial_ratio,
            num_channels,
        )

        if isinstance(batch.generator, list) and len(batch.generator) != batch_size:
            raise ValueError(f"generator list length {len(batch.generator)} does not match batch size {batch_size}.")

        visual_cond = getattr(self.transformer, "visual_cond", False)

        if batch.latents is None:
            latents = randn_tensor(shape, generator=batch.generator, device=device, dtype=dtype)
            if hasattr(self.scheduler, "init_noise_sigma"):
                latents = latents * self.scheduler.init_noise_sigma
        else:
            valid_shapes = [shape]
            if visual_cond:
                valid_shapes.append((*shape[:-1], num_channels * 2 + 1))
            if tuple(batch.latents.shape) not in valid_shapes:
                raise ValueError(f"Provided latents shape {list(batch.latents.shape)} does not match expected "
                                 f"Kandinsky5 latent shape(s): {[list(s) for s in valid_shapes]}.")
            latents = batch.latents.to(device=device, dtype=dtype)

        if visual_cond and latents.shape[-1] == num_channels:
            cond = torch.zeros_like(latents)
            cond_mask = torch.zeros(
                (*latents.shape[:-1], 1),
                device=latents.device,
                dtype=latents.dtype,
            )
            latents = torch.cat([latents, cond, cond_mask], dim=-1)

        # I2V image conditioning is placed by Kandinsky5ImageEncodingStage,
        # which runs AFTER this stage so the initial noise is the generator's
        # first draw (matching the official kandinskylab/kandinsky-5 order).
        batch.latents = latents
        batch.raw_latent_shape = (
            batch_size,
            num_channels,
            num_latent_frames,
            height // spatial_ratio,
            width // spatial_ratio,
        )
        return batch

    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
        result = VerificationResult()
        result.add_check("num_frames", batch.num_frames, V.positive_int)
        result.add_check("height", batch.height, V.positive_int)
        result.add_check("width", batch.width, V.positive_int)
        result.add_check("latents", batch.latents, V.none_or_tensor)
        return result

    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
        result = VerificationResult()
        result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
        return result

fastvideo.pipelines.stages.kandinsky5.Kandinsky5LatentPreparationStage.forward

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    if batch.height is None or batch.width is None:
        raise ValueError("height and width must be provided for Kandinsky5.")
    height = int(batch.height)
    width = int(batch.width)
    num_frames = int(batch.num_frames)

    temporal_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
    spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
    patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size

    if num_frames % temporal_ratio != 1:
        num_frames = num_frames // temporal_ratio * temporal_ratio + 1
        batch.num_frames = num_frames

    required_divisor_h = spatial_ratio * patch_size[1]
    required_divisor_w = spatial_ratio * patch_size[2]
    if height % required_divisor_h != 0 or width % required_divisor_w != 0:
        raise ValueError(f"Kandinsky5 height must be divisible by {required_divisor_h} and width by "
                         f"{required_divisor_w}; "
                         f"got height={height}, width={width}.")

    # NABLA sparse attention (Pro checkpoints) reshapes the post-patch grid
    # into 8x8 blocks; validate here instead of crashing mid-denoise after
    # all the encoding work is done.
    arch_cfg = getattr(self.transformer, "config", None) or fastvideo_args.pipeline_config.dit_config.arch_config
    if getattr(arch_cfg, "attention_type", "regular") == "nabla":
        nabla_divisor_h = required_divisor_h * 8
        nabla_divisor_w = required_divisor_w * 8
        if height % nabla_divisor_h != 0 or width % nabla_divisor_w != 0:
            raise ValueError(f"Kandinsky5 NABLA checkpoints require height divisible by {nabla_divisor_h} and "
                             f"width divisible by {nabla_divisor_w}; "
                             f"got height={height}, width={width}.")

    if isinstance(batch.prompt, list):
        batch_size = len(batch.prompt)
    elif batch.prompt is not None:
        batch_size = 1
    else:
        batch_size = batch.prompt_embeds[0].shape[0]
    batch_size *= batch.num_videos_per_prompt

    dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
    device = get_local_torch_device()
    num_latent_frames = (num_frames - 1) // temporal_ratio + 1
    num_channels = getattr(
        self.transformer,
        "in_visual_dim",
        fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
    )
    shape = (
        batch_size,
        num_latent_frames,
        height // spatial_ratio,
        width // spatial_ratio,
        num_channels,
    )

    if isinstance(batch.generator, list) and len(batch.generator) != batch_size:
        raise ValueError(f"generator list length {len(batch.generator)} does not match batch size {batch_size}.")

    visual_cond = getattr(self.transformer, "visual_cond", False)

    if batch.latents is None:
        latents = randn_tensor(shape, generator=batch.generator, device=device, dtype=dtype)
        if hasattr(self.scheduler, "init_noise_sigma"):
            latents = latents * self.scheduler.init_noise_sigma
    else:
        valid_shapes = [shape]
        if visual_cond:
            valid_shapes.append((*shape[:-1], num_channels * 2 + 1))
        if tuple(batch.latents.shape) not in valid_shapes:
            raise ValueError(f"Provided latents shape {list(batch.latents.shape)} does not match expected "
                             f"Kandinsky5 latent shape(s): {[list(s) for s in valid_shapes]}.")
        latents = batch.latents.to(device=device, dtype=dtype)

    if visual_cond and latents.shape[-1] == num_channels:
        cond = torch.zeros_like(latents)
        cond_mask = torch.zeros(
            (*latents.shape[:-1], 1),
            device=latents.device,
            dtype=latents.dtype,
        )
        latents = torch.cat([latents, cond, cond_mask], dim=-1)

    # I2V image conditioning is placed by Kandinsky5ImageEncodingStage,
    # which runs AFTER this stage so the initial noise is the generator's
    # first draw (matching the official kandinskylab/kandinsky-5 order).
    batch.latents = latents
    batch.raw_latent_shape = (
        batch_size,
        num_channels,
        num_latent_frames,
        height // spatial_ratio,
        width // spatial_ratio,
    )
    return batch

fastvideo.pipelines.stages.kandinsky5.Kandinsky5LatentPreparationStage.verify_input

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult
Source
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
    result = VerificationResult()
    result.add_check("num_frames", batch.num_frames, V.positive_int)
    result.add_check("height", batch.height, V.positive_int)
    result.add_check("width", batch.width, V.positive_int)
    result.add_check("latents", batch.latents, V.none_or_tensor)
    return result

fastvideo.pipelines.stages.kandinsky5.Kandinsky5LatentPreparationStage.verify_output

verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult
Source
def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
    result = VerificationResult()
    result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
    return result

fastvideo.pipelines.stages.kandinsky5.Kandinsky5DenoisingStage

Kandinsky5DenoisingStage(transformer, scheduler)
Source
class Kandinsky5DenoisingStage(PipelineStage):

    def __init__(self, transformer, scheduler) -> None:
        super().__init__()
        self.transformer = transformer
        self.scheduler = scheduler

    @staticmethod
    def _scale_factor(height: int, width: int) -> tuple[float, float, float]:
        if 480 <= height <= 854 and 480 <= width <= 854:
            return (1.0, 2.0, 2.0)
        return (1.0, 3.16, 3.16)

    @staticmethod
    def _text_rope_pos(mask: torch.Tensor, device: torch.device) -> torch.Tensor:
        seq_len = int(mask.sum(1).max().item())
        return torch.arange(seq_len, device=device)

    # Backends whose selection is all-or-nothing: fastvideo.platforms.cuda
    # raises ImportError immediately if the ATTN_QAT_TRAIN kernel isn't
    # available (training must never silently de-quantize), rather than the
    # generic "requested backend unsupported by this layer, falling back to
    # automatic selection" path _cached_get_attn_backend takes for most
    # other backend names (see fastvideo/attention/selector.py).
    # ATTN_QAT_TRAIN is also always in Kandinsky5ArchConfig's
    # _supported_attention_backends, so that generic fallback never applies
    # to it here. That combination makes an exact backend match a safe,
    # false-positive-free signal *only* for ATTN_QAT_TRAIN.
    # ATTN_QAT_INFER must NOT be in this set: fastvideo.platforms.cuda
    # deliberately resolves it to FlashAttention when the sm_120 kernel
    # isn't built, and the QAD recipe README documents that fallback ("on
    # other GPUs the attention falls back to a supported dense backend
    # while the FP4 linear layers still run") -- asserting an exact match
    # would abort inference on every non-sm_120 machine relying on it.
    # Any other FASTVIDEO_ATTENTION_BACKEND value (e.g. a backend meant for
    # a different model family sharing the same process/env var) can
    # likewise legitimately resolve to something other than what was
    # requested.
    _STRICT_BACKENDS = frozenset({
        AttentionBackendEnum.ATTN_QAT_TRAIN,
    })

    def _assert_local_attention_backend_engaged(self) -> None:
        """Guard against ``LocalAttention`` silently falling back to SDPA.

        ``Kandinsky5Attention.forward`` catches the ``AssertionError``
        ``LocalAttention`` raises when no pipeline forward context is set and
        falls back to plain ``F.scaled_dot_product_attention`` -- silently
        skipping the transformer's recorded attention backend (the
        fake-quantized ``ATTN_QAT_TRAIN`` kernel).
        ``LocalAttention.backend`` is resolved once at module construction,
        independent of whether a forward context is later set, so this only
        catches a backend that failed to resolve to the recorded decision
        -- it does not prove forward context was present for a
        given forward call. Pair it with always wrapping the actual
        transformer call in ``set_forward_context`` (see the denoising loops
        below), which is what prevents the runtime fallback.

        Scoped to ``_STRICT_BACKENDS``: other backend names -- including
        ``ATTN_QAT_INFER``, whose FlashAttention fallback on non-sm_120
        GPUs is documented, intentional behavior -- can legitimately differ
        from what ``LocalAttention.backend`` resolved to (see
        ``_STRICT_BACKENDS``'s comment), so asserting on those would flag
        expected behavior as a bug.
        """
        expected = component_attention_backend(self.transformer)
        if expected not in self._STRICT_BACKENDS:
            return
        for module in self.transformer.modules():
            if isinstance(module, LocalAttention):
                assert module.backend == expected, (
                    f"Kandinsky5 local attention resolved to backend {module.backend}, expected "
                    f"{expected}, the transformer's recorded attention backend. This likely means "
                    "LocalAttention's missing-forward-context guard silently fell back to SDPA.")
                return

    def _resolve_target_dtype(self, fastvideo_args: FastVideoArgs) -> torch.dtype:
        """Resolve the transformer's actual compute dtype.

        Trust ``pipeline_config.dit_precision`` directly for a normal,
        non-FSDP load -- ``TransformerLoader.load()`` asserts every
        parameter matches it exactly (``fastvideo/models/loader/
        component_loader.py``), so this holds for standalone T2V/I2V
        inference regardless of which precision (including fp32) was
        requested.

        FSDP2-wrapped modules are the one case where that assertion doesn't
        carry forward: ``maybe_load_fsdp_model`` hardcodes its
        ``MixedPrecisionPolicy`` to ``param_dtype=torch.bfloat16`` (with
        ``cast_forward_inputs=False``) for every FSDP-wrapped load,
        independent of ``dit_precision`` -- this covers both the live
        transformer ``ValidationCallback`` reuses from training (whose
        ``pipeline_config.dit_precision`` reflects the fp32 master-weight
        load dtype, not the actual bf16 compute dtype) and multi-GPU
        ``use_fsdp_inference=True`` runs.

        For that FSDP case, read the policy itself
        (``set_mixed_precision_policy`` records it right before
        ``maybe_load_fsdp_model`` shards -- every FSDP wrap in this repo
        goes through that path), NOT the parameter storage dtypes:
        FSDP2 computes in the policy's ``param_dtype`` regardless of what
        dtype the parameters are stored/loaded in, so e.g. an fp16-loaded
        FSDP model still runs its forward in bf16 -- scanning parameters
        would resolve fp16 and, with ``cast_forward_inputs=False`` and
        autocast keyed off the wrong dtype, hand bf16-computing modules
        fp16 inputs.
        """
        declared_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
        try:
            from torch.distributed.fsdp import FSDPModule
        except Exception:  # pragma: no cover - FSDP not always available
            return declared_dtype

        if not isinstance(self.transformer, FSDPModule):
            return declared_dtype

        try:
            policy_dtype = get_mixed_precision_state().param_dtype
        except ValueError:
            policy_dtype = None
        # bf16 mirrors maybe_load_fsdp_model's hardcoded param_dtype if the
        # (thread-local) policy state is somehow unset here.
        return policy_dtype if policy_dtype is not None else torch.bfloat16

    @staticmethod
    def fast_sta_nabla(
        T: int,
        H: int,
        W: int,
        wT: int = 3,
        wH: int = 3,
        wW: int = 3,
        device: torch.device | str = "cuda",
    ) -> torch.Tensor:
        """
        Create a sparse temporal attention (STA) mask for efficient video generation.

        This method generates a mask that limits attention to nearby frames and spatial positions, reducing
        computational complexity for video generation.

        Args:
            T (int): Number of temporal frames
            H (int): Height in latent space
            W (int): Width in latent space
            wT (int): Temporal attention window size
            wH (int): Height attention window size
            wW (int): Width attention window size
            device (str): Device to create tensor on

        Returns:
            torch.Tensor: Sparse attention mask of shape (T*H*W, T*H*W)
        """
        max_extent = int(torch.tensor([T, H, W], device=device).amax().item())
        r = torch.arange(0, max_extent, 1, dtype=torch.int16, device=device)
        mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs()
        sta_t, sta_h, sta_w = (
            mat[:T, :T].flatten(),
            mat[:H, :H].flatten(),
            mat[:W, :W].flatten(),
        )
        sta_t = sta_t <= wT // 2
        sta_h = sta_h <= wH // 2
        sta_w = sta_w <= wW // 2
        sta_hw = (sta_h.unsqueeze(1) * sta_w.unsqueeze(0)).reshape(H, H, W, W).transpose(1, 2).flatten()
        sta = (sta_t.unsqueeze(1) * sta_hw.unsqueeze(0)).reshape(T, T, H * W, H * W).transpose(1, 2)
        return sta.reshape(T * H * W, T * H * W)

    def get_sparse_params(self, sample: torch.Tensor, device: torch.device) -> dict[str, Any] | None:
        """
        Generate sparse attention parameters for the transformer based on sample dimensions.

        This method computes the sparse attention configuration needed for efficient video processing in the
        transformer model.

        Args:
            sample (torch.Tensor): Input sample tensor
            device (torch.device): Device to place tensors on

        Returns:
            Dict: Dictionary containing sparse attention parameters
        """
        assert self.transformer.config.patch_size[0] == 1
        _, T, H, W, _ = sample.shape
        T, H, W = (
            T // self.transformer.config.patch_size[0],
            H // self.transformer.config.patch_size[1],
            W // self.transformer.config.patch_size[2],
        )
        if self.transformer.config.attention_type == "nabla":
            sta_mask = self.fast_sta_nabla(
                T,
                H // 8,
                W // 8,
                self.transformer.config.attention_wT,
                self.transformer.config.attention_wH,
                self.transformer.config.attention_wW,
                device=device,
            )

            sparse_params = {
                "sta_mask": sta_mask.unsqueeze_(0).unsqueeze_(0),
                "attention_type": self.transformer.config.attention_type,
                "to_fractal": True,
                "P": self.transformer.config.attention_P,
                "wT": self.transformer.config.attention_wT,
                "wW": self.transformer.config.attention_wW,
                "wH": self.transformer.config.attention_wH,
                "add_sta": self.transformer.config.attention_add_sta,
                "visual_shape": (T, H, W),
                "method": self.transformer.config.attention_method,
            }
        else:
            sparse_params = None

        return sparse_params

    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
        if batch.timesteps is None:
            raise ValueError("timesteps must be prepared before Kandinsky5 denoising.")
        if batch.latents is None:
            raise ValueError("latents must be prepared before Kandinsky5 denoising.")
        if not fastvideo_args.model_loaded["transformer"]:
            loader = TransformerLoader()
            self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
            fastvideo_args.model_loaded["transformer"] = True
        self._assert_local_attention_backend_engaged()

        device = get_local_torch_device()
        target_dtype = self._resolve_target_dtype(fastvideo_args)
        autocast_enabled = target_dtype != torch.float32 and not fastvideo_args.disable_autocast
        latents = batch.latents
        num_channels = getattr(
            self.transformer,
            "in_visual_dim",
            fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
        )

        prompt_embeds = batch.prompt_embeds[0].to(device=device, dtype=target_dtype)
        pooled = batch.prompt_embeds[1].to(device=device, dtype=target_dtype)
        if batch.prompt_attention_mask is None or not batch.prompt_attention_mask:
            raise ValueError("Kandinsky5 requires Qwen prompt attention masks.")
        text_rope_pos = self._text_rope_pos(batch.prompt_attention_mask[0].to(device), device)

        neg_prompt_embeds = None
        neg_pooled = None
        negative_text_rope_pos = None
        if batch.do_classifier_free_guidance and batch.negative_prompt_embeds:
            neg_prompt_embeds = batch.negative_prompt_embeds[0].to(device=device, dtype=target_dtype)
            neg_pooled = batch.negative_prompt_embeds[1].to(device=device, dtype=target_dtype)
            if batch.negative_attention_mask is None or not batch.negative_attention_mask:
                raise ValueError("Kandinsky5 requires Qwen negative attention masks for CFG.")
            negative_text_rope_pos = self._text_rope_pos(batch.negative_attention_mask[0].to(device), device)

        height = int(batch.height)
        width = int(batch.width)
        temporal_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
        spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
        num_latent_frames = (int(batch.num_frames) - 1) // temporal_ratio + 1
        visual_rope_pos = [
            torch.arange(num_latent_frames, device=device),
            torch.arange(height // spatial_ratio // 2, device=device),
            torch.arange(width // spatial_ratio // 2, device=device),
        ]
        scale_factor = self._scale_factor(height, width)

        sparse_params = self.get_sparse_params(latents, device)

        # I2V keeps the first (conditioning) frame fixed during denoising.
        # Key off the actual image conditioning, not transformer.visual_cond:
        # official T2V checkpoints also ship visual_cond=True, and skipping
        # frame 0 for them leaves it as undenoised noise.
        cond_frames = 1 if batch.image_latent is not None else 0

        trajectory_timesteps: list[torch.Tensor] = []
        trajectory_latents: list[torch.Tensor] = []

        with tqdm(total=batch.num_inference_steps, desc="Kandinsky5 Denoising") as progress_bar:
            for i, timestep in enumerate(batch.timesteps):
                if hasattr(self, "interrupt") and self.interrupt:
                    break

                t_expand = timestep.unsqueeze(0).repeat(latents.shape[0]).to(device=device, dtype=target_dtype)
                attn_metadata = None
                if sparse_params is not None:
                    attn_metadata = NablaAttentionMetadataBuilder().build(
                        current_timestep=i,
                        sta_mask=sparse_params["sta_mask"],
                        P=sparse_params["P"],
                        visual_shape=sparse_params["visual_shape"],
                    )
                autocast_ctx = (torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled)
                                if device.type == "cuda" else contextlib.nullcontext())
                with set_forward_context(current_timestep=i, attn_metadata=attn_metadata,
                                         forward_batch=batch), autocast_ctx:
                    pred_velocity = self.transformer(
                        hidden_states=latents.to(dtype=target_dtype),
                        encoder_hidden_states=prompt_embeds,
                        pooled_projections=pooled,
                        timestep=t_expand,
                        visual_rope_pos=visual_rope_pos,
                        text_rope_pos=text_rope_pos,
                        scale_factor=scale_factor,
                        sparse_params=sparse_params,
                        return_dict=True,
                    ).sample

                    if neg_prompt_embeds is not None and neg_pooled is not None:
                        uncond_pred_velocity = self.transformer(
                            hidden_states=latents.to(dtype=target_dtype),
                            encoder_hidden_states=neg_prompt_embeds,
                            pooled_projections=neg_pooled,
                            timestep=t_expand,
                            visual_rope_pos=visual_rope_pos,
                            text_rope_pos=negative_text_rope_pos,
                            scale_factor=scale_factor,
                            sparse_params=sparse_params,
                            return_dict=True,
                        ).sample
                        pred_velocity = uncond_pred_velocity + batch.guidance_scale * (pred_velocity -
                                                                                       uncond_pred_velocity)

                latents[:, cond_frames:, :, :, :num_channels] = self.scheduler.step(
                    pred_velocity[:, cond_frames:],
                    timestep,
                    latents[:, cond_frames:, :, :, :num_channels],
                    return_dict=False,
                )[0]

                if batch.return_trajectory_latents:
                    trajectory_timesteps.append(timestep)
                    # latents is mutated in place, so snapshot a channels-first copy.
                    trajectory_latents.append(latents[..., :num_channels].permute(0, 4, 1, 2, 3).cpu())

                if i == len(batch.timesteps) - 1 or (i + 1) % self.scheduler.order == 0:
                    progress_bar.update()

        if trajectory_latents:
            batch.trajectory_latents = torch.stack(trajectory_latents, dim=1)
            batch.trajectory_timesteps = torch.stack(trajectory_timesteps, dim=0).cpu()

        batch.latents = latents[:, :, :, :, :num_channels]
        return batch

    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
        result = VerificationResult()
        result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
        result.add_check("prompt_embeds", batch.prompt_embeds, V.min_list_length(2))
        return result

fastvideo.pipelines.stages.kandinsky5.Kandinsky5DenoisingStage.fast_sta_nabla

fast_sta_nabla(T: int, H: int, W: int, wT: int = 3, wH: int = 3, wW: int = 3, device: torch.device | str = 'cuda') -> torch.Tensor

Create a sparse temporal attention (STA) mask for efficient video generation.

This method generates a mask that limits attention to nearby frames and spatial positions, reducing computational complexity for video generation.

Parameters

T: int

Number of temporal frames

H: int

Height in latent space

W: int

Width in latent space

wT: int

Temporal attention window size

wH: int

Height attention window size

wW: int

Width attention window size

device: str

Device to create tensor on

Returns

torch.Tensor

torch.Tensor: Sparse attention mask of shape (THW, THW)

Source
@staticmethod
def fast_sta_nabla(
    T: int,
    H: int,
    W: int,
    wT: int = 3,
    wH: int = 3,
    wW: int = 3,
    device: torch.device | str = "cuda",
) -> torch.Tensor:
    """
    Create a sparse temporal attention (STA) mask for efficient video generation.

    This method generates a mask that limits attention to nearby frames and spatial positions, reducing
    computational complexity for video generation.

    Args:
        T (int): Number of temporal frames
        H (int): Height in latent space
        W (int): Width in latent space
        wT (int): Temporal attention window size
        wH (int): Height attention window size
        wW (int): Width attention window size
        device (str): Device to create tensor on

    Returns:
        torch.Tensor: Sparse attention mask of shape (T*H*W, T*H*W)
    """
    max_extent = int(torch.tensor([T, H, W], device=device).amax().item())
    r = torch.arange(0, max_extent, 1, dtype=torch.int16, device=device)
    mat = (r.unsqueeze(1) - r.unsqueeze(0)).abs()
    sta_t, sta_h, sta_w = (
        mat[:T, :T].flatten(),
        mat[:H, :H].flatten(),
        mat[:W, :W].flatten(),
    )
    sta_t = sta_t <= wT // 2
    sta_h = sta_h <= wH // 2
    sta_w = sta_w <= wW // 2
    sta_hw = (sta_h.unsqueeze(1) * sta_w.unsqueeze(0)).reshape(H, H, W, W).transpose(1, 2).flatten()
    sta = (sta_t.unsqueeze(1) * sta_hw.unsqueeze(0)).reshape(T, T, H * W, H * W).transpose(1, 2)
    return sta.reshape(T * H * W, T * H * W)

fastvideo.pipelines.stages.kandinsky5.Kandinsky5DenoisingStage.get_sparse_params

get_sparse_params(sample: torch.Tensor, device: torch.device) -> dict[str, Any] | None

Generate sparse attention parameters for the transformer based on sample dimensions.

This method computes the sparse attention configuration needed for efficient video processing in the transformer model.

Parameters

sample: torch.Tensor

Input sample tensor

device: torch.device

Device to place tensors on

Returns

Dict: dict[str, Any] | None

Dictionary containing sparse attention parameters

Source
def get_sparse_params(self, sample: torch.Tensor, device: torch.device) -> dict[str, Any] | None:
    """
    Generate sparse attention parameters for the transformer based on sample dimensions.

    This method computes the sparse attention configuration needed for efficient video processing in the
    transformer model.

    Args:
        sample (torch.Tensor): Input sample tensor
        device (torch.device): Device to place tensors on

    Returns:
        Dict: Dictionary containing sparse attention parameters
    """
    assert self.transformer.config.patch_size[0] == 1
    _, T, H, W, _ = sample.shape
    T, H, W = (
        T // self.transformer.config.patch_size[0],
        H // self.transformer.config.patch_size[1],
        W // self.transformer.config.patch_size[2],
    )
    if self.transformer.config.attention_type == "nabla":
        sta_mask = self.fast_sta_nabla(
            T,
            H // 8,
            W // 8,
            self.transformer.config.attention_wT,
            self.transformer.config.attention_wH,
            self.transformer.config.attention_wW,
            device=device,
        )

        sparse_params = {
            "sta_mask": sta_mask.unsqueeze_(0).unsqueeze_(0),
            "attention_type": self.transformer.config.attention_type,
            "to_fractal": True,
            "P": self.transformer.config.attention_P,
            "wT": self.transformer.config.attention_wT,
            "wW": self.transformer.config.attention_wW,
            "wH": self.transformer.config.attention_wH,
            "add_sta": self.transformer.config.attention_add_sta,
            "visual_shape": (T, H, W),
            "method": self.transformer.config.attention_method,
        }
    else:
        sparse_params = None

    return sparse_params

fastvideo.pipelines.stages.kandinsky5.Kandinsky5DenoisingStage.forward

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    if batch.timesteps is None:
        raise ValueError("timesteps must be prepared before Kandinsky5 denoising.")
    if batch.latents is None:
        raise ValueError("latents must be prepared before Kandinsky5 denoising.")
    if not fastvideo_args.model_loaded["transformer"]:
        loader = TransformerLoader()
        self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
        fastvideo_args.model_loaded["transformer"] = True
    self._assert_local_attention_backend_engaged()

    device = get_local_torch_device()
    target_dtype = self._resolve_target_dtype(fastvideo_args)
    autocast_enabled = target_dtype != torch.float32 and not fastvideo_args.disable_autocast
    latents = batch.latents
    num_channels = getattr(
        self.transformer,
        "in_visual_dim",
        fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
    )

    prompt_embeds = batch.prompt_embeds[0].to(device=device, dtype=target_dtype)
    pooled = batch.prompt_embeds[1].to(device=device, dtype=target_dtype)
    if batch.prompt_attention_mask is None or not batch.prompt_attention_mask:
        raise ValueError("Kandinsky5 requires Qwen prompt attention masks.")
    text_rope_pos = self._text_rope_pos(batch.prompt_attention_mask[0].to(device), device)

    neg_prompt_embeds = None
    neg_pooled = None
    negative_text_rope_pos = None
    if batch.do_classifier_free_guidance and batch.negative_prompt_embeds:
        neg_prompt_embeds = batch.negative_prompt_embeds[0].to(device=device, dtype=target_dtype)
        neg_pooled = batch.negative_prompt_embeds[1].to(device=device, dtype=target_dtype)
        if batch.negative_attention_mask is None or not batch.negative_attention_mask:
            raise ValueError("Kandinsky5 requires Qwen negative attention masks for CFG.")
        negative_text_rope_pos = self._text_rope_pos(batch.negative_attention_mask[0].to(device), device)

    height = int(batch.height)
    width = int(batch.width)
    temporal_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
    spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
    num_latent_frames = (int(batch.num_frames) - 1) // temporal_ratio + 1
    visual_rope_pos = [
        torch.arange(num_latent_frames, device=device),
        torch.arange(height // spatial_ratio // 2, device=device),
        torch.arange(width // spatial_ratio // 2, device=device),
    ]
    scale_factor = self._scale_factor(height, width)

    sparse_params = self.get_sparse_params(latents, device)

    # I2V keeps the first (conditioning) frame fixed during denoising.
    # Key off the actual image conditioning, not transformer.visual_cond:
    # official T2V checkpoints also ship visual_cond=True, and skipping
    # frame 0 for them leaves it as undenoised noise.
    cond_frames = 1 if batch.image_latent is not None else 0

    trajectory_timesteps: list[torch.Tensor] = []
    trajectory_latents: list[torch.Tensor] = []

    with tqdm(total=batch.num_inference_steps, desc="Kandinsky5 Denoising") as progress_bar:
        for i, timestep in enumerate(batch.timesteps):
            if hasattr(self, "interrupt") and self.interrupt:
                break

            t_expand = timestep.unsqueeze(0).repeat(latents.shape[0]).to(device=device, dtype=target_dtype)
            attn_metadata = None
            if sparse_params is not None:
                attn_metadata = NablaAttentionMetadataBuilder().build(
                    current_timestep=i,
                    sta_mask=sparse_params["sta_mask"],
                    P=sparse_params["P"],
                    visual_shape=sparse_params["visual_shape"],
                )
            autocast_ctx = (torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled)
                            if device.type == "cuda" else contextlib.nullcontext())
            with set_forward_context(current_timestep=i, attn_metadata=attn_metadata,
                                     forward_batch=batch), autocast_ctx:
                pred_velocity = self.transformer(
                    hidden_states=latents.to(dtype=target_dtype),
                    encoder_hidden_states=prompt_embeds,
                    pooled_projections=pooled,
                    timestep=t_expand,
                    visual_rope_pos=visual_rope_pos,
                    text_rope_pos=text_rope_pos,
                    scale_factor=scale_factor,
                    sparse_params=sparse_params,
                    return_dict=True,
                ).sample

                if neg_prompt_embeds is not None and neg_pooled is not None:
                    uncond_pred_velocity = self.transformer(
                        hidden_states=latents.to(dtype=target_dtype),
                        encoder_hidden_states=neg_prompt_embeds,
                        pooled_projections=neg_pooled,
                        timestep=t_expand,
                        visual_rope_pos=visual_rope_pos,
                        text_rope_pos=negative_text_rope_pos,
                        scale_factor=scale_factor,
                        sparse_params=sparse_params,
                        return_dict=True,
                    ).sample
                    pred_velocity = uncond_pred_velocity + batch.guidance_scale * (pred_velocity -
                                                                                   uncond_pred_velocity)

            latents[:, cond_frames:, :, :, :num_channels] = self.scheduler.step(
                pred_velocity[:, cond_frames:],
                timestep,
                latents[:, cond_frames:, :, :, :num_channels],
                return_dict=False,
            )[0]

            if batch.return_trajectory_latents:
                trajectory_timesteps.append(timestep)
                # latents is mutated in place, so snapshot a channels-first copy.
                trajectory_latents.append(latents[..., :num_channels].permute(0, 4, 1, 2, 3).cpu())

            if i == len(batch.timesteps) - 1 or (i + 1) % self.scheduler.order == 0:
                progress_bar.update()

    if trajectory_latents:
        batch.trajectory_latents = torch.stack(trajectory_latents, dim=1)
        batch.trajectory_timesteps = torch.stack(trajectory_timesteps, dim=0).cpu()

    batch.latents = latents[:, :, :, :, :num_channels]
    return batch

fastvideo.pipelines.stages.kandinsky5.Kandinsky5DenoisingStage.verify_input

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult
Source
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
    result = VerificationResult()
    result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
    result.add_check("prompt_embeds", batch.prompt_embeds, V.min_list_length(2))
    return result

fastvideo.pipelines.stages.kandinsky5.Kandinsky5DmdDenoisingStage

Kandinsky5DmdDenoisingStage(transformer, scheduler)

DMD (few fixed steps, no CFG) variant of Kandinsky5DenoisingStage.

Reuses the parent's RoPE/scale_factor/sparse-params helpers; only the denoising loop differs: a fixed short timestep schedule (pipeline_config.dmd_denoising_steps) with a single forward pass per step and no classifier-free-guidance branch, matching DMD's distilled few-step generator.

dmd_denoising_steps values (e.g. [1000, 750, 500, 250]) are literal final target timesteps -- that's how DMD2Method on the training side resolves them: nearest-sigma lookup against a scheduler that has never had set_timesteps called on it, so e.g. "750" maps to sigma 0.75. This stage therefore keeps a private FlowMatchEulerDiscreteScheduler (same shift as the pipeline scheduler) and drives it directly via predict-x0 + re-noise, exactly mirroring DMD2Method._student_rollout's "simulate" branch and Wan's own DmdDenoisingStage. It must NOT reuse the pipeline's shared scheduler object through scheduler.step(): by the time this stage runs, TimestepPreparationStage has already called scheduler.set_timesteps(timesteps=dmd_denoising_steps) on it, which re-applies the flow-match shift warp on top of values that are already final (e.g. sigma 0.75 -> 0.9375 at shift=5) -- a double shift that leaves every step far noisier than the student was trained for, so the sampled video is still mostly noise after the last step instead of converged.

Source
class Kandinsky5DmdDenoisingStage(Kandinsky5DenoisingStage):
    """DMD (few fixed steps, no CFG) variant of Kandinsky5DenoisingStage.

    Reuses the parent's RoPE/scale_factor/sparse-params helpers; only the
    denoising loop differs: a fixed short timestep schedule
    (``pipeline_config.dmd_denoising_steps``) with a single forward pass per
    step and no classifier-free-guidance branch, matching DMD's distilled
    few-step generator.

    ``dmd_denoising_steps`` values (e.g. ``[1000, 750, 500, 250]``) are
    literal *final* target timesteps -- that's how ``DMD2Method`` on the
    training side resolves them: nearest-sigma lookup against a scheduler
    that has never had ``set_timesteps`` called on it, so e.g. "750" maps to
    sigma 0.75. This stage therefore keeps a private
    ``FlowMatchEulerDiscreteScheduler`` (same ``shift`` as the pipeline
    scheduler) and drives it directly via predict-x0 + re-noise, exactly
    mirroring ``DMD2Method._student_rollout``'s "simulate" branch and Wan's
    own ``DmdDenoisingStage``. It must NOT reuse the pipeline's shared
    ``scheduler`` object through ``scheduler.step()``: by the time this
    stage runs, ``TimestepPreparationStage`` has already called
    ``scheduler.set_timesteps(timesteps=dmd_denoising_steps)`` on it, which
    re-applies the flow-match ``shift`` warp on top of values that are
    already final (e.g. sigma 0.75 -> 0.9375 at shift=5) -- a double shift
    that leaves every step far noisier than the student was trained for, so
    the sampled video is still mostly noise after the last step instead of
    converged.
    """

    def __init__(self, transformer, scheduler) -> None:
        super().__init__(transformer, scheduler)
        self._sample_scheduler = FlowMatchEulerDiscreteScheduler(shift=scheduler.shift)

    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
        if batch.latents is None:
            raise ValueError("latents must be prepared before Kandinsky5 DMD denoising.")
        if not fastvideo_args.model_loaded["transformer"]:
            loader = TransformerLoader()
            self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
            fastvideo_args.model_loaded["transformer"] = True
        self._assert_local_attention_backend_engaged()

        device = get_local_torch_device()
        target_dtype = self._resolve_target_dtype(fastvideo_args)
        autocast_enabled = target_dtype != torch.float32 and not fastvideo_args.disable_autocast
        latents = batch.latents
        num_channels = getattr(
            self.transformer,
            "in_visual_dim",
            fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
        )

        prompt_embeds = batch.prompt_embeds[0].to(device=device, dtype=target_dtype)
        pooled = batch.prompt_embeds[1].to(device=device, dtype=target_dtype)
        if batch.prompt_attention_mask is None or not batch.prompt_attention_mask:
            raise ValueError("Kandinsky5 DMD requires Qwen prompt attention masks.")
        text_rope_pos = self._text_rope_pos(batch.prompt_attention_mask[0].to(device), device)

        height = int(batch.height)
        width = int(batch.width)
        temporal_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
        spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
        num_latent_frames = (int(batch.num_frames) - 1) // temporal_ratio + 1
        visual_rope_pos = [
            torch.arange(num_latent_frames, device=device),
            torch.arange(height // spatial_ratio // 2, device=device),
            torch.arange(width // spatial_ratio // 2, device=device),
        ]
        scale_factor = self._scale_factor(height, width)
        sparse_params = self.get_sparse_params(latents, device)

        dmd_steps = fastvideo_args.pipeline_config.dmd_denoising_steps
        if not dmd_steps:
            raise ValueError("Kandinsky5 DMD denoising requires "
                             "pipeline_config.dmd_denoising_steps to be set.")
        timesteps = torch.tensor(dmd_steps, dtype=torch.long, device=device)

        # I2V keeps the first (conditioning) frame fixed during denoising.
        # Key off the actual image conditioning, not transformer.visual_cond:
        # official T2V checkpoints also ship visual_cond=True, and skipping
        # frame 0 for them leaves it as undenoised noise (see the same fix
        # in Kandinsky5DenoisingStage.forward() above).
        cond_frames = 1 if batch.image_latent is not None else 0

        with tqdm(total=len(timesteps), desc="Kandinsky5 DMD Denoising") as progress_bar:
            for i, timestep in enumerate(timesteps):
                if hasattr(self, "interrupt") and self.interrupt:
                    continue

                t_expand = timestep.unsqueeze(0).repeat(latents.shape[0]).to(device=device, dtype=target_dtype)
                attn_metadata = None
                if sparse_params is not None:
                    attn_metadata = NablaAttentionMetadataBuilder().build(
                        current_timestep=i,
                        sta_mask=sparse_params["sta_mask"],
                        P=sparse_params["P"],
                        visual_shape=sparse_params["visual_shape"],
                    )
                autocast_ctx = (torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled)
                                if device.type == "cuda" else contextlib.nullcontext())
                with set_forward_context(current_timestep=i, attn_metadata=attn_metadata,
                                         forward_batch=batch), autocast_ctx:
                    pred_velocity = self.transformer(
                        hidden_states=latents.to(dtype=target_dtype),
                        encoder_hidden_states=prompt_embeds,
                        pooled_projections=pooled,
                        timestep=t_expand,
                        visual_rope_pos=visual_rope_pos,
                        text_rope_pos=text_rope_pos,
                        scale_factor=scale_factor,
                        sparse_params=sparse_params,
                        return_dict=True,
                    ).sample

                # Channel-last [B, T', H, W, C] -> channel-first [B, T', C, H, W]
                # to match pred_noise_to_pred_video/add_noise's expected layout
                # (same convention as Kandinsky5Model.add_noise/predict_x0).
                sample_cf = latents[:, cond_frames:, :, :, :num_channels].permute(0, 1, 4, 2, 3)
                velocity_cf = pred_velocity[:, cond_frames:].permute(0, 1, 4, 2, 3)
                b, t = sample_cf.shape[:2]
                step_timestep = timestep.reshape(1).to(device=device)
                pred_x0_cf = pred_noise_to_pred_video(
                    pred_noise=velocity_cf.flatten(0, 1),
                    noise_input_latent=sample_cf.flatten(0, 1),
                    timestep=step_timestep,
                    scheduler=self._sample_scheduler,
                ).unflatten(0, (b, t))

                if i < len(timesteps) - 1:
                    next_timestep = timesteps[i + 1].reshape(1).to(device=device)
                    noise = randn_tensor(
                        sample_cf.shape,
                        generator=batch.generator,
                        device=device,
                        dtype=pred_x0_cf.dtype,
                    )
                    next_cf = self._sample_scheduler.add_noise(
                        pred_x0_cf.flatten(0, 1),
                        noise.flatten(0, 1),
                        next_timestep,
                    ).unflatten(0, (b, t))
                else:
                    next_cf = pred_x0_cf

                latents[:, cond_frames:, :, :, :num_channels] = next_cf.permute(0, 1, 3, 4, 2).to(latents.dtype)
                progress_bar.update()

        batch.latents = latents[:, :, :, :, :num_channels]
        return batch

fastvideo.pipelines.stages.kandinsky5.Kandinsky5DmdDenoisingStage.forward

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    if batch.latents is None:
        raise ValueError("latents must be prepared before Kandinsky5 DMD denoising.")
    if not fastvideo_args.model_loaded["transformer"]:
        loader = TransformerLoader()
        self.transformer = loader.load(fastvideo_args.model_paths["transformer"], fastvideo_args)
        fastvideo_args.model_loaded["transformer"] = True
    self._assert_local_attention_backend_engaged()

    device = get_local_torch_device()
    target_dtype = self._resolve_target_dtype(fastvideo_args)
    autocast_enabled = target_dtype != torch.float32 and not fastvideo_args.disable_autocast
    latents = batch.latents
    num_channels = getattr(
        self.transformer,
        "in_visual_dim",
        fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
    )

    prompt_embeds = batch.prompt_embeds[0].to(device=device, dtype=target_dtype)
    pooled = batch.prompt_embeds[1].to(device=device, dtype=target_dtype)
    if batch.prompt_attention_mask is None or not batch.prompt_attention_mask:
        raise ValueError("Kandinsky5 DMD requires Qwen prompt attention masks.")
    text_rope_pos = self._text_rope_pos(batch.prompt_attention_mask[0].to(device), device)

    height = int(batch.height)
    width = int(batch.width)
    temporal_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
    spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio
    num_latent_frames = (int(batch.num_frames) - 1) // temporal_ratio + 1
    visual_rope_pos = [
        torch.arange(num_latent_frames, device=device),
        torch.arange(height // spatial_ratio // 2, device=device),
        torch.arange(width // spatial_ratio // 2, device=device),
    ]
    scale_factor = self._scale_factor(height, width)
    sparse_params = self.get_sparse_params(latents, device)

    dmd_steps = fastvideo_args.pipeline_config.dmd_denoising_steps
    if not dmd_steps:
        raise ValueError("Kandinsky5 DMD denoising requires "
                         "pipeline_config.dmd_denoising_steps to be set.")
    timesteps = torch.tensor(dmd_steps, dtype=torch.long, device=device)

    # I2V keeps the first (conditioning) frame fixed during denoising.
    # Key off the actual image conditioning, not transformer.visual_cond:
    # official T2V checkpoints also ship visual_cond=True, and skipping
    # frame 0 for them leaves it as undenoised noise (see the same fix
    # in Kandinsky5DenoisingStage.forward() above).
    cond_frames = 1 if batch.image_latent is not None else 0

    with tqdm(total=len(timesteps), desc="Kandinsky5 DMD Denoising") as progress_bar:
        for i, timestep in enumerate(timesteps):
            if hasattr(self, "interrupt") and self.interrupt:
                continue

            t_expand = timestep.unsqueeze(0).repeat(latents.shape[0]).to(device=device, dtype=target_dtype)
            attn_metadata = None
            if sparse_params is not None:
                attn_metadata = NablaAttentionMetadataBuilder().build(
                    current_timestep=i,
                    sta_mask=sparse_params["sta_mask"],
                    P=sparse_params["P"],
                    visual_shape=sparse_params["visual_shape"],
                )
            autocast_ctx = (torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled)
                            if device.type == "cuda" else contextlib.nullcontext())
            with set_forward_context(current_timestep=i, attn_metadata=attn_metadata,
                                     forward_batch=batch), autocast_ctx:
                pred_velocity = self.transformer(
                    hidden_states=latents.to(dtype=target_dtype),
                    encoder_hidden_states=prompt_embeds,
                    pooled_projections=pooled,
                    timestep=t_expand,
                    visual_rope_pos=visual_rope_pos,
                    text_rope_pos=text_rope_pos,
                    scale_factor=scale_factor,
                    sparse_params=sparse_params,
                    return_dict=True,
                ).sample

            # Channel-last [B, T', H, W, C] -> channel-first [B, T', C, H, W]
            # to match pred_noise_to_pred_video/add_noise's expected layout
            # (same convention as Kandinsky5Model.add_noise/predict_x0).
            sample_cf = latents[:, cond_frames:, :, :, :num_channels].permute(0, 1, 4, 2, 3)
            velocity_cf = pred_velocity[:, cond_frames:].permute(0, 1, 4, 2, 3)
            b, t = sample_cf.shape[:2]
            step_timestep = timestep.reshape(1).to(device=device)
            pred_x0_cf = pred_noise_to_pred_video(
                pred_noise=velocity_cf.flatten(0, 1),
                noise_input_latent=sample_cf.flatten(0, 1),
                timestep=step_timestep,
                scheduler=self._sample_scheduler,
            ).unflatten(0, (b, t))

            if i < len(timesteps) - 1:
                next_timestep = timesteps[i + 1].reshape(1).to(device=device)
                noise = randn_tensor(
                    sample_cf.shape,
                    generator=batch.generator,
                    device=device,
                    dtype=pred_x0_cf.dtype,
                )
                next_cf = self._sample_scheduler.add_noise(
                    pred_x0_cf.flatten(0, 1),
                    noise.flatten(0, 1),
                    next_timestep,
                ).unflatten(0, (b, t))
            else:
                next_cf = pred_x0_cf

            latents[:, cond_frames:, :, :, :num_channels] = next_cf.permute(0, 1, 3, 4, 2).to(latents.dtype)
            progress_bar.update()

    batch.latents = latents[:, :, :, :, :num_channels]
    return batch

fastvideo.pipelines.stages.kandinsky5.Kandinsky5DecodingStage

Kandinsky5DecodingStage(vae: ParallelTiledVAE, pipeline = None)
Source
class Kandinsky5DecodingStage(DecodingStage):

    def __init__(self, vae: ParallelTiledVAE, pipeline=None) -> None:
        super().__init__(vae=vae, pipeline=pipeline)

    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
        if batch.latents is None:
            raise ValueError("latents must be available before Kandinsky5 decoding.")
        # Kandinsky5 latents are channels-last [B, T, H, W, C]; the base stage
        # (and the trajectory latents recorded by the denoising stage) work
        # channels-first.
        batch.latents = batch.latents.permute(0, 4, 1, 2, 3).contiguous()
        return super().forward(batch, fastvideo_args)

fastvideo.pipelines.stages.kandinsky5.Kandinsky5DecodingStage.forward

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    if batch.latents is None:
        raise ValueError("latents must be available before Kandinsky5 decoding.")
    # Kandinsky5 latents are channels-last [B, T, H, W, C]; the base stage
    # (and the trajectory latents recorded by the denoising stage) work
    # channels-first.
    batch.latents = batch.latents.permute(0, 4, 1, 2, 3).contiguous()
    return super().forward(batch, fastvideo_args)

fastvideo.pipelines.stages.kandinsky5.Kandinsky5ImageEncodingStage

Kandinsky5ImageEncodingStage(vae: ParallelTiledVAE, pipeline = None)

Encode the conditioning image into a VAE latent for I2V.

Source
class Kandinsky5ImageEncodingStage(EncodingStage):
    """Encode the conditioning image into a VAE latent for I2V."""

    def __init__(self, vae: ParallelTiledVAE, pipeline=None) -> None:
        super().__init__(vae=vae)

    @staticmethod
    def _preprocess(image, height: int, width: int) -> torch.Tensor:
        if isinstance(image, PIL.Image.Image):
            image = resize(image, height, width)
            image = numpy_to_pt(pil_to_numpy(image))  # always lands in [0, 1]
            return normalize(image)  # [0, 1] -> [-1, 1]
        # Tensor input: no reliable way to tell [0, 1] from an already
        # normalized [-1, 1] tensor whose values happen to be non-negative,
        # so mirror diffusers' heuristic and say what we assumed.
        if image.min() >= 0:
            logger.warning("Kandinsky5 conditioning image tensor has no negative values; "
                           "assuming range [0, 1] and normalizing to [-1, 1]. "
                           "Pass a [-1, 1] tensor with negative values to skip normalization.")
            image = normalize(image)  # [0, 1] -> [-1, 1]
        if image.ndim == 3:
            image = image.unsqueeze(0)
        if image.shape[-2:] != (height, width):
            image = torch.nn.functional.interpolate(image.float(),
                                                    size=(height, width),
                                                    mode="bilinear",
                                                    antialias=True)
        return image

    @torch.no_grad()
    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
        if batch.pil_image is None:
            raise ValueError("Kandinsky5 I2V requires an input image.")

        if not fastvideo_args.model_loaded["vae"]:
            vae = getattr(self, "vae", None)
            if vae is None:
                loader = VAELoader()
                vae = loader.load(fastvideo_args.model_paths["vae"], fastvideo_args)
                self.vae = vae
            fastvideo_args.model_loaded["vae"] = True

        device = get_local_torch_device()
        vae = self.vae.to(device)
        self.vae = vae
        vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
        vae_autocast_enabled = vae_dtype != torch.float32 and not fastvideo_args.disable_autocast

        # [B, C, H, W] -> [B, C, 1, H, W]
        image = self._preprocess(batch.pil_image, int(batch.height), int(batch.width))
        image = image.to(device=device, dtype=torch.float32).unsqueeze(2)

        # Encode the single conditioning frame without tiling (matches diffusers).
        # The untested causal-VAE spatial_tiled_encode path corrupts the latent.
        prev_use_tiling = vae.use_tiling
        vae.use_tiling = False
        try:
            with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
                if not vae_autocast_enabled:
                    image = image.to(vae_dtype)
                # Sample with the batch generator (diffusers parity); mode()
                # would make seed-for-seed reproduction of the reference
                # pipeline impossible.
                generator = batch.generator
                if isinstance(generator, list) and len(generator) != image.shape[0]:
                    generator = generator[0]
                image_latent = vae.encode(image).sample(generator=generator)
        finally:
            vae.use_tiling = prev_use_tiling

        image_latent = image_latent * vae.scaling_factor
        # [B, C, 1, H, W] -> [B, 1, H, W, C] to match channels-last latents
        batch.image_latent = image_latent.permute(0, 2, 3, 4, 1).contiguous()

        # Place the conditioning latent into the prepared latents: frame 0 of
        # the main channels, the visual_cond channel block, and the mask.
        # NOTE: the official kandinsky-5 repo leaves the visual_cond block
        # zeros (generation_utils.py generate()), while the diffusers port
        # copies the image latent into it. A same-seed A/B on
        # Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers showed the Diffusers
        # export requires the copy: zeroing the block produces smeared faces
        # mid-video. Keep the diffusers semantics for Diffusers-format
        # checkpoints.
        latents = batch.latents
        image_latent = batch.image_latent.to(device=latents.device, dtype=latents.dtype)
        num_channels = image_latent.shape[-1]
        latents[:, 0:1, :, :, :num_channels] = image_latent
        if latents.shape[-1] > num_channels:
            latents[:, 0:1, :, :, num_channels:2 * num_channels] = image_latent
            latents[:, 0:1, :, :, 2 * num_channels:] = 1.0
        batch.latents = latents

        if fastvideo_args.vae_cpu_offload:
            vae.to("cpu")
        return batch

    def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
        result = VerificationResult()
        result.add_check("pil_image", batch.pil_image, V.not_none)
        result.add_check("height", batch.height, V.positive_int)
        result.add_check("width", batch.width, V.positive_int)
        # This stage runs after latent preparation and writes into its output.
        result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
        return result

    def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
        result = VerificationResult()
        result.add_check("image_latent", batch.image_latent, [V.is_tensor, V.with_dims(5)])
        return result

fastvideo.pipelines.stages.kandinsky5.Kandinsky5ImageEncodingStage.forward

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    if batch.pil_image is None:
        raise ValueError("Kandinsky5 I2V requires an input image.")

    if not fastvideo_args.model_loaded["vae"]:
        vae = getattr(self, "vae", None)
        if vae is None:
            loader = VAELoader()
            vae = loader.load(fastvideo_args.model_paths["vae"], fastvideo_args)
            self.vae = vae
        fastvideo_args.model_loaded["vae"] = True

    device = get_local_torch_device()
    vae = self.vae.to(device)
    self.vae = vae
    vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
    vae_autocast_enabled = vae_dtype != torch.float32 and not fastvideo_args.disable_autocast

    # [B, C, H, W] -> [B, C, 1, H, W]
    image = self._preprocess(batch.pil_image, int(batch.height), int(batch.width))
    image = image.to(device=device, dtype=torch.float32).unsqueeze(2)

    # Encode the single conditioning frame without tiling (matches diffusers).
    # The untested causal-VAE spatial_tiled_encode path corrupts the latent.
    prev_use_tiling = vae.use_tiling
    vae.use_tiling = False
    try:
        with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
            if not vae_autocast_enabled:
                image = image.to(vae_dtype)
            # Sample with the batch generator (diffusers parity); mode()
            # would make seed-for-seed reproduction of the reference
            # pipeline impossible.
            generator = batch.generator
            if isinstance(generator, list) and len(generator) != image.shape[0]:
                generator = generator[0]
            image_latent = vae.encode(image).sample(generator=generator)
    finally:
        vae.use_tiling = prev_use_tiling

    image_latent = image_latent * vae.scaling_factor
    # [B, C, 1, H, W] -> [B, 1, H, W, C] to match channels-last latents
    batch.image_latent = image_latent.permute(0, 2, 3, 4, 1).contiguous()

    # Place the conditioning latent into the prepared latents: frame 0 of
    # the main channels, the visual_cond channel block, and the mask.
    # NOTE: the official kandinsky-5 repo leaves the visual_cond block
    # zeros (generation_utils.py generate()), while the diffusers port
    # copies the image latent into it. A same-seed A/B on
    # Kandinsky-5.0-I2V-Pro-distilled-5s-Diffusers showed the Diffusers
    # export requires the copy: zeroing the block produces smeared faces
    # mid-video. Keep the diffusers semantics for Diffusers-format
    # checkpoints.
    latents = batch.latents
    image_latent = batch.image_latent.to(device=latents.device, dtype=latents.dtype)
    num_channels = image_latent.shape[-1]
    latents[:, 0:1, :, :, :num_channels] = image_latent
    if latents.shape[-1] > num_channels:
        latents[:, 0:1, :, :, num_channels:2 * num_channels] = image_latent
        latents[:, 0:1, :, :, 2 * num_channels:] = 1.0
    batch.latents = latents

    if fastvideo_args.vae_cpu_offload:
        vae.to("cpu")
    return batch

fastvideo.pipelines.stages.kandinsky5.Kandinsky5ImageEncodingStage.verify_input

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult
Source
def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
    result = VerificationResult()
    result.add_check("pil_image", batch.pil_image, V.not_none)
    result.add_check("height", batch.height, V.positive_int)
    result.add_check("width", batch.width, V.positive_int)
    # This stage runs after latent preparation and writes into its output.
    result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)])
    return result

fastvideo.pipelines.stages.kandinsky5.Kandinsky5ImageEncodingStage.verify_output

verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult
Source
def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult:
    result = VerificationResult()
    result.add_check("image_latent", batch.image_latent, [V.is_tensor, V.with_dims(5)])
    return result

fastvideo.pipelines.stages.kandinsky5.Kandinsky5NormalizationStage

Kandinsky5NormalizationStage()

Normalize the first latent frames to reduce I2V conditioning artifacts.

Source
class Kandinsky5NormalizationStage(PipelineStage):
    """Normalize the first latent frames to reduce I2V conditioning artifacts."""

    COND_FRAMES = 4
    REFERENCE_FRAMES = 5

    @staticmethod
    def _adaptive_mean_std(source: torch.Tensor, reference: torch.Tensor) -> torch.Tensor:
        source_mean = source.mean(dim=(1, 2, 3, 4), keepdim=True)
        source_std = source.std(dim=(1, 2, 3, 4), keepdim=True)
        # Magic constants limit how far the first frames may drift.
        ref_mean = torch.clamp(reference.mean(dim=(1, 2, 3, 4), keepdim=True), source_mean - 0.05, source_mean + 0.1)
        ref_std = torch.clamp(reference.std(dim=(1, 2, 3, 4), keepdim=True), source_std - 0.1, source_std + 0.25)
        normalized = (source - source_mean) / source_std
        return normalized * ref_std + ref_mean

    def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
        latents = batch.latents
        n = self.COND_FRAMES
        if latents is None or latents.shape[1] <= n:
            return batch

        reference = latents[:, n:n + min(self.REFERENCE_FRAMES, latents.shape[1] - 1)]
        latents[:, :n] = self._adaptive_mean_std(latents[:, :n].clone(), reference)
        batch.latents = latents
        return batch

fastvideo.pipelines.stages.kandinsky5.Kandinsky5NormalizationStage.forward

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
    latents = batch.latents
    n = self.COND_FRAMES
    if latents is None or latents.shape[1] <= n:
        return batch

    reference = latents[:, n:n + min(self.REFERENCE_FRAMES, latents.shape[1] - 1)]
    latents[:, :n] = self._adaptive_mean_std(latents[:, :n].clone(), reference)
    batch.latents = latents
    return batch