Initialize per-chunk KV cache containers for SANA cached modules.

Optionally zero-out the camera branch during training (drop-path style).

Convert frame-valid mask to token/beta/decay masks used by GDN blocks.

Build ``(apply_q, apply_kv, apply_o)`` block-diagonal callables for UCPE.

Toggle ``_kv_cache_capture`` (pre-RoPE) or ``_tf_capture_kv`` (post-RoPE) per block.

Get shape of tensor if it is a torch.Tensor.

Average latest n values or all values.

Drop-in chunkwise replacement for `cam_scan_func`. Args mirror `cam_scan_func` exactly: q, k, v: ``(B, H, D, N)`` fp32 contiguous (cam-prep

Clear all cache state.

Inverse of :func:`compute_fx_from_fov_xi`.

Recover focal length ``fx`` from horizontal FoV (degrees) + UCM xi.

Args: clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W)

Enable tiled VAE decoding for memory-efficient processing of large videos.

Forward pass of Sana. x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) t: (N,) tensor of

Calculate the coefficient of gradients.

Initialize and get a logger by name. If the logger has not been initialized, this method will initialize the logger by adding one or two hand

Compute a hash of the input string and return the first 16 characters of the hash.

Initialize random seed. If the seed is not set, the seed will be automatically randomized, and then broadcast to all processes to prevent som

Args: strategy: init strategy name **kwargs: strategy specific parameters Returns: initialized

Multistep DPM-Solver with the order `order` from time `t_prev_list[-1]` to time `t`. Args: x: A pytorch tensor. The init

Keep LTX-2 audio-only branches off GPU for this wrapper's video-only forward.

Apply the model to get p(x_{t-1} | x_t), as well as a prediction of the initial x, x_0. :param model: the model, which takes

Retrieve a specific row from a specific array in the LMDB.

Return the directory part of a url.

Apply causal ShortConvolution along T, with S merged into batch. Under CP, a causal conv of kernel size K needs K-1 left-context fram

Route short conv: chunk-causal when chunk boundaries exist, else bidirectional. For single-chunk (``chunk_size >= T``) or no chunk_size, use

Cached identity RMS/RoPE tables used by ``cam_scan_chunkwise``.

Internal decode method.

Returns (a_warps, a_BLOCK_S, b_warps, b_stages, b_use_acc_fusion, c_warps, c_BLOCK_S, c_stages). dot_precision: 0=bf16 TC, 1=TF32

Get the conditional_dict to use for the current chunk

Randomly generate a timestep tensor based on the generator's task type. It uniformly samples a timestep from the range [min_timestep,

initialize training pipeline for DMDSana, using SanaTrainingPipeline

Vectorized categorical sampling from a single column of log-weights. logw_col: [T] starts, ends: [B], slice is [start, end)

Convert ``(B, F, 20)`` camera conditions (C2W flat + fx,fy,cx,cy) into ``(raymats, absmap)``. ``raymats`` is ``(B, F, H, W, 4, 4)`` ``ray<-wo

Reshape (B*S, T, C) back to (B, T*S, C).

Reshape (B, T*S, C) to (B*S, T, C) for temporal conv. Returns: Reshaped tensor and (B, S, T) for later restoration.

Slice rotary embedding freqs to the trailing ``current_n`` token positions. When ``sink_token=true``, upstream rope is built for sink + current c

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

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

Build train/test datasets with distributed samplers and dataloaders. Returns ``(train_dataset, train_dataloader, train_sampler, test_dataset, tes