Convert frame-valid mask to token mask shaped ``(B, N, 1)``.

Run one forward at σ=0 with capture hooks; return per-layer (K, V). ``capture_mode='pre_rope'`` uses the ``_kv_cache_capture`` hook (stored

Capture the current layer self-attention K/V without computing attention output.

Compute output size of a convolution operation.

Convert flow matching's prediction to x0 prediction. flow_pred: the prediction with shape [B, C, H, W] xt: the input noisy da

Convert x0 prediction to flow matching's prediction. x0_pred: the x0 prediction with shape [B, C, H, W] xt: the input noisy d

Dequantize 8bit quantized data to 32bit float Args: quantized_chunks: list of quantized data blocks Returns:

Emit one decoded chunk to the MP4 writer / callback (single source of truth for the in-loop flush and the final drain).

Internal encode method.

Helper function to extract non-padded token embeddings from a batch. Args: hidden_states (torch.Tensor): The output embe

Return ``arr`` fitted to ``num_frames`` along axis 0.

Shared body of ``_forward_video_only`` that takes a pre-built RoPE. Used by the AR refinement path where each block forward needs custom

Build a flex_attention BlockMask for chunk-causal attention. Token ``q_idx`` can attend to ``kv_idx`` iff ``chunk(q) >= chunk(kv)``, i.e. ``k

get model layers from state dict

Wraps `torchvision.models.inception_v3`

Initialize gates near identity to mimic Linear Attention at start.

Invert a 4x4 SE(3) matrix batch (closed-form). Mirrors the production ``_invert_SE3`` in ``sana_camctrl_blocks.py``; inlined to keep this mod

Compute ``1/RMS`` per ``(b, n)`` over full-``C`` channels. Args: raw: ``(B, N, H, D)`` raw QKV projection output (typically fp32).

Camera branch Q/K/V for softmax attention. Mirrors ``_GDNUCPEBase._prepare_cam_qkv`` but skips the ReLU kernel and GDN key scaling — standard

Convert complex RoPE ``(1, 1, N, D_half//2)`` to interleaved ``(N, D_half)`` cos/sin. Uses the interleaved-pair convention: y[2i] = x[2

Lightweight variant of ``_process_camera_conditions_ucpe`` — raymats only. Computes *only* the per-ray ``world -> ray_local`` SE(3) transforms us

Materialise the camera-to-world trajectory from --camera or --action.

Optionally simulate the generator's input from noise using backward simulation and then run the generator for one-step. Input

Run SDPA with chunk-causal masking when needed, or plain SDPA otherwise. Replicates the masking logic from ``_forward_softmax_attn`` and ``_f

Return the transformer to use for inference based on *mode*. mode: "compile_nvfp4" | "compile" | "peft"

Mirror tian's ``_inject_kv_prefix``: attach a per-layer prefix dict to each ``attn1``.

Initialize the inference pipeline for visualization on CPU, to be moved to GPU only when needed.

Re-slice WAN RoPE frequencies for a smaller rope_dim using the same (T, H, W) split.

Snap ``n`` to the nearest ``stride*k + 1`` (LTX-2 VAE constraint). Ties round up to keep the user's requested length when possible. If the ro

LTX-2 self-attention with sink/current streaming mask + AR KV-cache hooks. Two modes are layered on top of vanilla diffusers self-attention, sele

Swap pipeline.transformer to the model specified by *mode*. mode: "compile_nvfp4" | "compile" | "peft"

Swap pipeline.transformer to the model specified by *mode*. mode: "compile_nvfp4" | "compile" | "peft"

"Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the prediction of x_0 at t

Generate and save sample videos to monitor training progress.

SA-Corrector, without the "rescaling" trick in Appendix D in SA-Solver paper https://arxiv.org/pdf/2309.05019.pdf

SA-Corrector, without the "rescaling" trick in Appendix D in SA-Solver paper https://arxiv.org/pdf/2309.05019.pdf

Composite the action overlay onto each frame. Args: video_hwc: ``(T, H, W, 3)`` uint8 video. c2w: ``(T_pose, 4, 4)`` camera-to-wo

A fast approximation of the cumulative distribution function of the standard normal.

Flatten to the 8-tuple the legacy API returns.

Blend horizontally between two tiles.

Blend temporally between two tiles.

Blend vertically between two tiles.

Spatial downsample is always performed. Temporal downsample is optional.