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</code></pre></details>
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&#91;1&#93;
required_divisor_w = spatial_ratio * patch_size&#91;2&#93;
if height % required_divisor_h != 0 or width % required_divisor_w != 0:
    raise ValueError(f&quot;Kandinsky5 height must be divisible by &#123;required_divisor_h&#125; and width by &quot;
                     f&quot;&#123;required_divisor_w&#125;; &quot;
                     f&quot;got height=&#123;height&#125;, width=&#123;width&#125;.&quot;)

# 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, &quot;config&quot;, None) or fastvideo_args.pipeline_config.dit_config.arch_config
if getattr(arch_cfg, &quot;attention_type&quot;, &quot;regular&quot;) == &quot;nabla&quot;:
    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&quot;Kandinsky5 NABLA checkpoints require height divisible by &#123;nabla_divisor_h&#125; and &quot;
                         f&quot;width divisible by &#123;nabla_divisor_w&#125;; &quot;
                         f&quot;got height=&#123;height&#125;, width=&#123;width&#125;.&quot;)

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&#91;0&#93;.shape&#91;0&#93;
batch_size *= batch.num_videos_per_prompt

dtype = PRECISION_TO_TYPE&#91;fastvideo_args.pipeline_config.dit_precision&#93;
device = get_local_torch_device()
num_latent_frames = (num_frames - 1) // temporal_ratio + 1
num_channels = getattr(
    self.transformer,
    &quot;in_visual_dim&quot;,
    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&quot;generator list length &#123;len(batch.generator)&#125; does not match batch size &#123;batch_size&#125;.&quot;)

visual_cond = getattr(self.transformer, &quot;visual_cond&quot;, False)

if batch.latents is None:
    latents = randn_tensor(shape, generator=batch.generator, device=device, dtype=dtype)
    if hasattr(self.scheduler, &quot;init_noise_sigma&quot;):
        latents = latents * self.scheduler.init_noise_sigma
else:
    valid_shapes = &#91;shape&#93;
    if visual_cond:
        valid_shapes.append((*shape&#91;:-1&#93;, num_channels * 2 + 1))
    if tuple(batch.latents.shape) not in valid_shapes:
        raise ValueError(f&quot;Provided latents shape &#123;list(batch.latents.shape)&#125; does not match expected &quot;
                         f&quot;Kandinsky5 latent shape(s): &#123;&#91;list(s) for s in valid_shapes&#93;&#125;.&quot;)
    latents = batch.latents.to(device=device, dtype=dtype)

if visual_cond and latents.shape&#91;-1&#93; == num_channels:
    cond = torch.zeros_like(latents)
    cond_mask = torch.zeros(
        (*latents.shape&#91;:-1&#93;, 1),
        device=latents.device,
        dtype=latents.dtype,
    )
    latents = torch.cat(&#91;latents, cond, cond_mask&#93;, dim=-1)

# I2V image conditioning is placed by Kandinsky5ImageEncodingStage,
# which runs AFTER this stage so the initial noise is the generator&#x27;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</code></pre></details>
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) -&gt; None:
    super().__init__()
    self.transformer = transformer
    self.scheduler = scheduler

@staticmethod
def _scale_factor(height: int, width: int) -&gt; tuple&#91;float, float, float&#93;:
    if 480 &lt;= height &lt;= 854 and 480 &lt;= width &lt;= 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) -&gt; 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&#x27;t
# available (training must never silently de-quantize), rather than the
# generic &quot;requested backend unsupported by this layer, falling back to
# automatic selection&quot; path _cached_get_attn_backend takes for most
# other backend names (see fastvideo/attention/selector.py).
# ATTN_QAT_TRAIN is also always in Kandinsky5ArchConfig&#x27;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&#x27;t built, and the QAD recipe README documents that fallback (&quot;on
# other GPUs the attention falls back to a supported dense backend
# while the FP4 linear layers still run&quot;) -- 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(&#123;
    AttentionBackendEnum.ATTN_QAT_TRAIN,
&#125;)

def _assert_local_attention_backend_engaged(self) -&gt; None:
    &quot;&quot;&quot;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&#x27;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``&#x27;s comment), so asserting on those would flag
    expected behavior as a bug.
    &quot;&quot;&quot;
    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&quot;Kandinsky5 local attention resolved to backend &#123;module.backend&#125;, expected &quot;
                f&quot;&#123;expected&#125;, the transformer&#x27;s recorded attention backend. This likely means &quot;
                &quot;LocalAttention&#x27;s missing-forward-context guard silently fell back to SDPA.&quot;)
            return

def _resolve_target_dtype(self, fastvideo_args: FastVideoArgs) -&gt; torch.dtype:
    &quot;&quot;&quot;Resolve the transformer&#x27;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&#x27;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&#x27;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.
    &quot;&quot;&quot;
    declared_dtype = PRECISION_TO_TYPE&#91;fastvideo_args.pipeline_config.dit_precision&#93;
    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&#x27;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 = &quot;cuda&quot;,
) -&gt; torch.Tensor:
    &quot;&quot;&quot;
    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)
    &quot;&quot;&quot;
    max_extent = int(torch.tensor(&#91;T, H, W&#93;, 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&#91;:T, :T&#93;.flatten(),
        mat&#91;:H, :H&#93;.flatten(),
        mat&#91;:W, :W&#93;.flatten(),
    )
    sta_t = sta_t &lt;= wT // 2
    sta_h = sta_h &lt;= wH // 2
    sta_w = sta_w &lt;= 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) -&gt; dict&#91;str, Any&#93; | None:
    &quot;&quot;&quot;
    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
    &quot;&quot;&quot;
    assert self.transformer.config.patch_size&#91;0&#93; == 1
    _, T, H, W, _ = sample.shape
    T, H, W = (
        T // self.transformer.config.patch_size&#91;0&#93;,
        H // self.transformer.config.patch_size&#91;1&#93;,
        W // self.transformer.config.patch_size&#91;2&#93;,
    )
    if self.transformer.config.attention_type == &quot;nabla&quot;:
        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 = &#123;
            &quot;sta_mask&quot;: sta_mask.unsqueeze_(0).unsqueeze_(0),
            &quot;attention_type&quot;: self.transformer.config.attention_type,
            &quot;to_fractal&quot;: True,
            &quot;P&quot;: self.transformer.config.attention_P,
            &quot;wT&quot;: self.transformer.config.attention_wT,
            &quot;wW&quot;: self.transformer.config.attention_wW,
            &quot;wH&quot;: self.transformer.config.attention_wH,
            &quot;add_sta&quot;: self.transformer.config.attention_add_sta,
            &quot;visual_shape&quot;: (T, H, W),
            &quot;method&quot;: self.transformer.config.attention_method,
        &#125;
    else:
        sparse_params = None

    return sparse_params

def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -&gt; ForwardBatch:
    if batch.timesteps is None:
        raise ValueError(&quot;timesteps must be prepared before Kandinsky5 denoising.&quot;)
    if batch.latents is None:
        raise ValueError(&quot;latents must be prepared before Kandinsky5 denoising.&quot;)
    if not fastvideo_args.model_loaded&#91;&quot;transformer&quot;&#93;:
        loader = TransformerLoader()
        self.transformer = loader.load(fastvideo_args.model_paths&#91;&quot;transformer&quot;&#93;, fastvideo_args)
        fastvideo_args.model_loaded&#91;&quot;transformer&quot;&#93; = 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,
        &quot;in_visual_dim&quot;,
        fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
    )

    prompt_embeds = batch.prompt_embeds&#91;0&#93;.to(device=device, dtype=target_dtype)
    pooled = batch.prompt_embeds&#91;1&#93;.to(device=device, dtype=target_dtype)
    if batch.prompt_attention_mask is None or not batch.prompt_attention_mask:
        raise ValueError(&quot;Kandinsky5 requires Qwen prompt attention masks.&quot;)
    text_rope_pos = self._text_rope_pos(batch.prompt_attention_mask&#91;0&#93;.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&#91;0&#93;.to(device=device, dtype=target_dtype)
        neg_pooled = batch.negative_prompt_embeds&#91;1&#93;.to(device=device, dtype=target_dtype)
        if batch.negative_attention_mask is None or not batch.negative_attention_mask:
            raise ValueError(&quot;Kandinsky5 requires Qwen negative attention masks for CFG.&quot;)
        negative_text_rope_pos = self._text_rope_pos(batch.negative_attention_mask&#91;0&#93;.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 = &#91;
        torch.arange(num_latent_frames, device=device),
        torch.arange(height // spatial_ratio // 2, device=device),
        torch.arange(width // spatial_ratio // 2, device=device),
    &#93;
    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&#91;torch.Tensor&#93; = &#91;&#93;
    trajectory_latents: list&#91;torch.Tensor&#93; = &#91;&#93;

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

            t_expand = timestep.unsqueeze(0).repeat(latents.shape&#91;0&#93;).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&#91;&quot;sta_mask&quot;&#93;,
                    P=sparse_params&#91;&quot;P&quot;&#93;,
                    visual_shape=sparse_params&#91;&quot;visual_shape&quot;&#93;,
                )
            autocast_ctx = (torch.autocast(device_type=&quot;cuda&quot;, dtype=target_dtype, enabled=autocast_enabled)
                            if device.type == &quot;cuda&quot; 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&#91;:, cond_frames:, :, :, :num_channels&#93; = self.scheduler.step(
                pred_velocity&#91;:, cond_frames:&#93;,
                timestep,
                latents&#91;:, cond_frames:, :, :, :num_channels&#93;,
                return_dict=False,
            )&#91;0&#93;

            if batch.return_trajectory_latents:
                trajectory_timesteps.append(timestep)
                # latents is mutated in place, so snapshot a channels-first copy.
                trajectory_latents.append(latents&#91;..., :num_channels&#93;.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&#91;:, :, :, :, :num_channels&#93;
    return batch

def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -&gt; VerificationResult:
    result = VerificationResult()
    result.add_check(&quot;latents&quot;, batch.latents, &#91;V.is_tensor, V.with_dims(5)&#93;)
    result.add_check(&quot;prompt_embeds&quot;, batch.prompt_embeds, V.min_list_length(2))
    return result</code></pre></details>
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)
&quot;&quot;&quot;
max_extent = int(torch.tensor(&#91;T, H, W&#93;, 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&#91;:T, :T&#93;.flatten(),
    mat&#91;:H, :H&#93;.flatten(),
    mat&#91;:W, :W&#93;.flatten(),
)
sta_t = sta_t &lt;= wT // 2
sta_h = sta_h &lt;= wH // 2
sta_w = sta_w &lt;= 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)</code></pre></details>
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
&quot;&quot;&quot;
assert self.transformer.config.patch_size&#91;0&#93; == 1
_, T, H, W, _ = sample.shape
T, H, W = (
    T // self.transformer.config.patch_size&#91;0&#93;,
    H // self.transformer.config.patch_size&#91;1&#93;,
    W // self.transformer.config.patch_size&#91;2&#93;,
)
if self.transformer.config.attention_type == &quot;nabla&quot;:
    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 = &#123;
        &quot;sta_mask&quot;: sta_mask.unsqueeze_(0).unsqueeze_(0),
        &quot;attention_type&quot;: self.transformer.config.attention_type,
        &quot;to_fractal&quot;: True,
        &quot;P&quot;: self.transformer.config.attention_P,
        &quot;wT&quot;: self.transformer.config.attention_wT,
        &quot;wW&quot;: self.transformer.config.attention_wW,
        &quot;wH&quot;: self.transformer.config.attention_wH,
        &quot;add_sta&quot;: self.transformer.config.attention_add_sta,
        &quot;visual_shape&quot;: (T, H, W),
        &quot;method&quot;: self.transformer.config.attention_method,
    &#125;
else:
    sparse_params = None

return sparse_params</code></pre></details>
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,
    &quot;in_visual_dim&quot;,
    fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
)

prompt_embeds = batch.prompt_embeds&#91;0&#93;.to(device=device, dtype=target_dtype)
pooled = batch.prompt_embeds&#91;1&#93;.to(device=device, dtype=target_dtype)
if batch.prompt_attention_mask is None or not batch.prompt_attention_mask:
    raise ValueError(&quot;Kandinsky5 requires Qwen prompt attention masks.&quot;)
text_rope_pos = self._text_rope_pos(batch.prompt_attention_mask&#91;0&#93;.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&#91;0&#93;.to(device=device, dtype=target_dtype)
    neg_pooled = batch.negative_prompt_embeds&#91;1&#93;.to(device=device, dtype=target_dtype)
    if batch.negative_attention_mask is None or not batch.negative_attention_mask:
        raise ValueError(&quot;Kandinsky5 requires Qwen negative attention masks for CFG.&quot;)
    negative_text_rope_pos = self._text_rope_pos(batch.negative_attention_mask&#91;0&#93;.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 = &#91;
    torch.arange(num_latent_frames, device=device),
    torch.arange(height // spatial_ratio // 2, device=device),
    torch.arange(width // spatial_ratio // 2, device=device),
&#93;
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&#91;torch.Tensor&#93; = &#91;&#93;
trajectory_latents: list&#91;torch.Tensor&#93; = &#91;&#93;

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

        t_expand = timestep.unsqueeze(0).repeat(latents.shape&#91;0&#93;).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&#91;&quot;sta_mask&quot;&#93;,
                P=sparse_params&#91;&quot;P&quot;&#93;,
                visual_shape=sparse_params&#91;&quot;visual_shape&quot;&#93;,
            )
        autocast_ctx = (torch.autocast(device_type=&quot;cuda&quot;, dtype=target_dtype, enabled=autocast_enabled)
                        if device.type == &quot;cuda&quot; 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&#91;:, cond_frames:, :, :, :num_channels&#93; = self.scheduler.step(
            pred_velocity&#91;:, cond_frames:&#93;,
            timestep,
            latents&#91;:, cond_frames:, :, :, :num_channels&#93;,
            return_dict=False,
        )&#91;0&#93;

        if batch.return_trajectory_latents:
            trajectory_timesteps.append(timestep)
            # latents is mutated in place, so snapshot a channels-first copy.
            trajectory_latents.append(latents&#91;..., :num_channels&#93;.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&#91;:, :, :, :, :num_channels&#93;
return batch</code></pre></details>
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&#x27;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&#x27;s distilled
few-step generator.

``dmd_denoising_steps`` values (e.g. ``&#91;1000, 750, 500, 250&#93;``) are
literal *final* target timesteps -- that&#x27;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. &quot;750&quot; 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``&#x27;s &quot;simulate&quot; branch and Wan&#x27;s
own ``DmdDenoisingStage``. It must NOT reuse the pipeline&#x27;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 -&gt; 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.
&quot;&quot;&quot;

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

def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -&gt; ForwardBatch:
    if batch.latents is None:
        raise ValueError(&quot;latents must be prepared before Kandinsky5 DMD denoising.&quot;)
    if not fastvideo_args.model_loaded&#91;&quot;transformer&quot;&#93;:
        loader = TransformerLoader()
        self.transformer = loader.load(fastvideo_args.model_paths&#91;&quot;transformer&quot;&#93;, fastvideo_args)
        fastvideo_args.model_loaded&#91;&quot;transformer&quot;&#93; = 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,
        &quot;in_visual_dim&quot;,
        fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
    )

    prompt_embeds = batch.prompt_embeds&#91;0&#93;.to(device=device, dtype=target_dtype)
    pooled = batch.prompt_embeds&#91;1&#93;.to(device=device, dtype=target_dtype)
    if batch.prompt_attention_mask is None or not batch.prompt_attention_mask:
        raise ValueError(&quot;Kandinsky5 DMD requires Qwen prompt attention masks.&quot;)
    text_rope_pos = self._text_rope_pos(batch.prompt_attention_mask&#91;0&#93;.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 = &#91;
        torch.arange(num_latent_frames, device=device),
        torch.arange(height // spatial_ratio // 2, device=device),
        torch.arange(width // spatial_ratio // 2, device=device),
    &#93;
    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(&quot;Kandinsky5 DMD denoising requires &quot;
                         &quot;pipeline_config.dmd_denoising_steps to be set.&quot;)
    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=&quot;Kandinsky5 DMD Denoising&quot;) as progress_bar:
        for i, timestep in enumerate(timesteps):
            if hasattr(self, &quot;interrupt&quot;) and self.interrupt:
                continue

            t_expand = timestep.unsqueeze(0).repeat(latents.shape&#91;0&#93;).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&#91;&quot;sta_mask&quot;&#93;,
                    P=sparse_params&#91;&quot;P&quot;&#93;,
                    visual_shape=sparse_params&#91;&quot;visual_shape&quot;&#93;,
                )
            autocast_ctx = (torch.autocast(device_type=&quot;cuda&quot;, dtype=target_dtype, enabled=autocast_enabled)
                            if device.type == &quot;cuda&quot; 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 &#91;B, T&#x27;, H, W, C&#93; -&gt; channel-first &#91;B, T&#x27;, C, H, W&#93;
            # to match pred_noise_to_pred_video/add_noise&#x27;s expected layout
            # (same convention as Kandinsky5Model.add_noise/predict_x0).
            sample_cf = latents&#91;:, cond_frames:, :, :, :num_channels&#93;.permute(0, 1, 4, 2, 3)
            velocity_cf = pred_velocity&#91;:, cond_frames:&#93;.permute(0, 1, 4, 2, 3)
            b, t = sample_cf.shape&#91;:2&#93;
            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 &lt; len(timesteps) - 1:
                next_timestep = timesteps&#91;i + 1&#93;.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&#91;:, cond_frames:, :, :, :num_channels&#93; = next_cf.permute(0, 1, 3, 4, 2).to(latents.dtype)
            progress_bar.update()

    batch.latents = latents&#91;:, :, :, :, :num_channels&#93;
    return batch</code></pre></details>
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,
    &quot;in_visual_dim&quot;,
    fastvideo_args.pipeline_config.dit_config.arch_config.in_visual_dim,
)

prompt_embeds = batch.prompt_embeds&#91;0&#93;.to(device=device, dtype=target_dtype)
pooled = batch.prompt_embeds&#91;1&#93;.to(device=device, dtype=target_dtype)
if batch.prompt_attention_mask is None or not batch.prompt_attention_mask:
    raise ValueError(&quot;Kandinsky5 DMD requires Qwen prompt attention masks.&quot;)
text_rope_pos = self._text_rope_pos(batch.prompt_attention_mask&#91;0&#93;.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 = &#91;
    torch.arange(num_latent_frames, device=device),
    torch.arange(height // spatial_ratio // 2, device=device),
    torch.arange(width // spatial_ratio // 2, device=device),
&#93;
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(&quot;Kandinsky5 DMD denoising requires &quot;
                     &quot;pipeline_config.dmd_denoising_steps to be set.&quot;)
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=&quot;Kandinsky5 DMD Denoising&quot;) as progress_bar:
    for i, timestep in enumerate(timesteps):
        if hasattr(self, &quot;interrupt&quot;) and self.interrupt:
            continue

        t_expand = timestep.unsqueeze(0).repeat(latents.shape&#91;0&#93;).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&#91;&quot;sta_mask&quot;&#93;,
                P=sparse_params&#91;&quot;P&quot;&#93;,
                visual_shape=sparse_params&#91;&quot;visual_shape&quot;&#93;,
            )
        autocast_ctx = (torch.autocast(device_type=&quot;cuda&quot;, dtype=target_dtype, enabled=autocast_enabled)
                        if device.type == &quot;cuda&quot; 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 &#91;B, T&#x27;, H, W, C&#93; -&gt; channel-first &#91;B, T&#x27;, C, H, W&#93;
        # to match pred_noise_to_pred_video/add_noise&#x27;s expected layout
        # (same convention as Kandinsky5Model.add_noise/predict_x0).
        sample_cf = latents&#91;:, cond_frames:, :, :, :num_channels&#93;.permute(0, 1, 4, 2, 3)
        velocity_cf = pred_velocity&#91;:, cond_frames:&#93;.permute(0, 1, 4, 2, 3)
        b, t = sample_cf.shape&#91;:2&#93;
        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 &lt; len(timesteps) - 1:
            next_timestep = timesteps&#91;i + 1&#93;.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&#91;:, cond_frames:, :, :, :num_channels&#93; = next_cf.permute(0, 1, 3, 4, 2).to(latents.dtype)
        progress_bar.update()

batch.latents = latents&#91;:, :, :, :, :num_channels&#93;
return batch</code></pre></details>
fastvideo.pipelines.stages.kandinsky5.Kandinsky5DecodingStage
Kandinsky5DecodingStage(vae: ParallelTiledVAE, pipeline = None)
Source
class Kandinsky5DecodingStage(DecodingStage):
def __init__(self, vae: ParallelTiledVAE, pipeline=None) -&gt; None:
    super().__init__(vae=vae, pipeline=pipeline)

def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -&gt; ForwardBatch:
    if batch.latents is None:
        raise ValueError(&quot;latents must be available before Kandinsky5 decoding.&quot;)
    # Kandinsky5 latents are channels-last &#91;B, T, H, W, C&#93;; 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)</code></pre></details>
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) -&gt; None:
    super().__init__(vae=vae)

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

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

    if not fastvideo_args.model_loaded&#91;&quot;vae&quot;&#93;:
        vae = getattr(self, &quot;vae&quot;, None)
        if vae is None:
            loader = VAELoader()
            vae = loader.load(fastvideo_args.model_paths&#91;&quot;vae&quot;&#93;, fastvideo_args)
            self.vae = vae
        fastvideo_args.model_loaded&#91;&quot;vae&quot;&#93; = True

    device = get_local_torch_device()
    vae = self.vae.to(device)
    self.vae = vae
    vae_dtype = PRECISION_TO_TYPE&#91;fastvideo_args.pipeline_config.vae_precision&#93;
    vae_autocast_enabled = vae_dtype != torch.float32 and not fastvideo_args.disable_autocast

    # &#91;B, C, H, W&#93; -&gt; &#91;B, C, 1, H, W&#93;
    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=&quot;cuda&quot;, 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&#91;0&#93;:
                generator = generator&#91;0&#93;
            image_latent = vae.encode(image).sample(generator=generator)
    finally:
        vae.use_tiling = prev_use_tiling

    image_latent = image_latent * vae.scaling_factor
    # &#91;B, C, 1, H, W&#93; -&gt; &#91;B, 1, H, W, C&#93; 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&#91;-1&#93;
    latents&#91;:, 0:1, :, :, :num_channels&#93; = image_latent
    if latents.shape&#91;-1&#93; &gt; num_channels:
        latents&#91;:, 0:1, :, :, num_channels:2 * num_channels&#93; = image_latent
        latents&#91;:, 0:1, :, :, 2 * num_channels:&#93; = 1.0
    batch.latents = latents

    if fastvideo_args.vae_cpu_offload:
        vae.to(&quot;cpu&quot;)
    return batch

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

def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -&gt; VerificationResult:
    result = VerificationResult()
    result.add_check(&quot;image_latent&quot;, batch.image_latent, &#91;V.is_tensor, V.with_dims(5)&#93;)
    return result</code></pre></details>
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&#91;&quot;vae&quot;&#93;:
    vae = getattr(self, &quot;vae&quot;, None)
    if vae is None:
        loader = VAELoader()
        vae = loader.load(fastvideo_args.model_paths&#91;&quot;vae&quot;&#93;, fastvideo_args)
        self.vae = vae
    fastvideo_args.model_loaded&#91;&quot;vae&quot;&#93; = True

device = get_local_torch_device()
vae = self.vae.to(device)
self.vae = vae
vae_dtype = PRECISION_TO_TYPE&#91;fastvideo_args.pipeline_config.vae_precision&#93;
vae_autocast_enabled = vae_dtype != torch.float32 and not fastvideo_args.disable_autocast

# &#91;B, C, H, W&#93; -&gt; &#91;B, C, 1, H, W&#93;
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=&quot;cuda&quot;, 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&#91;0&#93;:
            generator = generator&#91;0&#93;
        image_latent = vae.encode(image).sample(generator=generator)
finally:
    vae.use_tiling = prev_use_tiling

image_latent = image_latent * vae.scaling_factor
# &#91;B, C, 1, H, W&#93; -&gt; &#91;B, 1, H, W, C&#93; 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&#91;-1&#93;
latents&#91;:, 0:1, :, :, :num_channels&#93; = image_latent
if latents.shape&#91;-1&#93; &gt; num_channels:
    latents&#91;:, 0:1, :, :, num_channels:2 * num_channels&#93; = image_latent
    latents&#91;:, 0:1, :, :, 2 * num_channels:&#93; = 1.0
batch.latents = latents

if fastvideo_args.vae_cpu_offload:
    vae.to(&quot;cpu&quot;)
return batch</code></pre></details>
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) -&gt; 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) -&gt; ForwardBatch:
    latents = batch.latents
    n = self.COND_FRAMES
    if latents is None or latents.shape&#91;1&#93; &lt;= n:
        return batch

    reference = latents&#91;:, n:n + min(self.REFERENCE_FRAMES, latents.shape&#91;1&#93; - 1)&#93;
    latents&#91;:, :n&#93; = self._adaptive_mean_std(latents&#91;:, :n&#93;.clone(), reference)
    batch.latents = latents
    return batch</code></pre></details>
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&#91;:, n:n + min(self.REFERENCE_FRAMES, latents.shape&#91;1&#93; - 1)&#93;
latents&#91;:, :n&#93; = self._adaptive_mean_std(latents&#91;:, :n&#93;.clone(), reference)
batch.latents = latents
return batch</code></pre></details>