Decode one latent block, mutating the VAE's persistent feature cache. Args: z_chunk: Normalized latent block ``(B, C, T_lat_block

Dilate mask to expand mask regions Args: mask: PIL Image mask kernel_size: int, dilation kernel size

DPM-Solver-1 (equivalent to DDIM) from time `s` to time `t`. Args: x: A pytorch tensor. The initial value at time `s`.

Return a list of ``(step, path)`` checkpoint candidates sorted by step descending.

Bidi chunkwise GDN forward, optionally with state-cache for autoregressive sampling (chunk 0 = full bidi with state save; chunks > 0 seed forward

Precompute the frequency tensor for complex exponentials (cis) with given dimensions. This function calculates a frequency tensor with compl

Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time. For discrete-time DPMs, we convert `t_continuou

Check if chunk_index represents uniform chunking. Returns True if all chunks are equal to chunk_size except possibly the last chunk which may

Load a yaml file.

Save debug images to disk and log to wandb during rollout.

Compute log(alpha_t) of a given continuous-time label t in [0, T].

Normalize chunk_index and detect if uniform. This function handles all the complex logic for: 1. Converting chunk_size + strategy → chunk_ind

Precompute UCPE apply functions once for a batch (shared across all blocks). Only ``camctrl_type == "UCPE"`` is supported. Accepts either precom

Lazily replace eligible refiner Linear layers with TE NVFP4 Linear modules.

Diffuse the data for a given number of diffusion steps. In other words, sample from q(x_t | x_0). :param x_start: the initial

Refine one AR block; advance internal KV state. Args: block_idx: 0-based block index in the AR schedule. Used only for

Try to load training state from *candidates* (output of :func:`find_resume_candidates`). *peft_model* should be the unwrapped PeftModel (use

Save LoRA adapters, EMA, optimizer and scaler to a checkpoint directory.

Sets the discrete timesteps used for the diffusion chain (to be run before inference). Args: num_inference_steps (`int`):

Singlestep solver DPM-Solver-2 from time `s` to time `t`. Args: x: A pytorch tensor. The initial value at time `s`.

Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned m

Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion process from the learned m

Helper function to synchronize (barrier) among all processes when using distributed training

Unproject grid with intrinsics expressed as FoV (degrees) + xi.

x: (N, T, patch_size**2 * C) imgs: (N, H, W, C)

Apply positional embedding to input tensor. Args: x: Input tensor (N, T, D) bs: Batch size start_f: Start

Fused spatial pipeline: inverted_conv -> depth_conv -> GLU -> point_conv.