Dequantize list of QuantizedTensor to a contiguous bf16 tensor. kv_list[block_idx] -> QuantizedTensor(block_token_size * num_heads, 128)

Return the quantization configuration for the gradient tensors.

Get the path of the submodule relative to the root directory.

Update the current scheduler with the interface's static method

Get the tasks that should be evaluated. If a group name is set, tasks will only be evaluated if they have not yet been evaluated for

Perform inference on the given noise and text prompts. Inputs: noise (torch.Tensor): The input noise tensor of shape

Return True when a state dict contains materialized FourOverSix NVFP4 buffers.

Returns the closest integer to 'number' that is divisible by 'factor'.

Save the results of a PTQ experiment to the SQLite database.

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

Rescales the image so that the following conditions are met: 1. Both dimensions (height and width) are divisible by 'factor'. 2. The to

Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep DPMSolver.

Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with the multistep UniPC.

Return a hash of the configuration.

Load caption for a video file.

Materialize NVFP4 weights before FSDP wraps quantized modules.

Materialize quantized weights and drop master weights. Optionally cache an additional transposed quantized layout for training paths that

Slice a block-wise conditional dict to keep only the last `new_num_segments` segments.

Return a high-precision tensor with the dequantized values (PyTorch impl).

Drop high-precision master weights after loading materialized NVFP4 buffers.

Returns the largest integer less than or equal to 'number' that is divisible by 'factor'.

Generate image/videos from noise and compute the DMD loss. The noisy input to the generator is backward simulated. This remov

Return True for checkpoints saved with TransformerEngine module state.

Run the command.

Select the offset slice that applies to one sample. ``temporal_offset`` accepts a scalar, ``[B]`` per-sample constants, ``[F]`` shared per-fr

Frames that are sharded by training sequence parallelism.

r""" Reconstruct video tensors from patch embeddings. Args: x (List[Tensor]): List of patchified features

Applies cache updates collected from multiple blocks. Args: kv_cache: List of cache dictionaries for each block

Pure local: drop ``items`` into ``buffer``. No collective, no cross-rank coordination — each rank may invoke this independently (or sk

Configure LoRA for a WanDiffusionWrapper model

Configures a model object. This includes setting evaluation modes, applying distributed parallel strategy, and handling device placem

Configures a model object. This includes setting evaluation modes, applying distributed parallel strategy, and handling device placem

All-gather errors/timesteps across the appropriate group and return a list of ready-to-add ``(err_block, t_idx, pos_or_None)`` items.

On rank-0, gather FULL_STATE_DICT, then filter only LoRA weights

Get sorted video files in a folder, keeping only those with a per-video caption (cached).

Lazy initialize the inference pipeline during the first backward simulation run. Here we encapsulate the inference code with a model-

Return True when *chunk_index* starts a new shot (prompt-based detection). Pure prompt check — no dependency on sink config so that Narrative

Mark the current chunk's buffer position as pinned for multi-shot sink. The pinned region REPLACES the original sink on the next rolling even

Block-wise causal mask. The mask is defined only by the AR chunk size: a token can attend to all tokens before the end of its current

we will divide the token sequence into the following format [1 latent frame] [N latent frame] ... [N latent frame] The first

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

Sample frames from a video preserving original duration. Args: video_path: Path to video file num_frames: Num

Toggle SP self-attention between training and inference. This only applies to 5B runs with SP enabled.

Bidirectionally swap model weights with EMA shadow weights. Calling this twice restores both the model and EMA to their original stat

"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

"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

Try to get item from a specific folder. Returns (True, result) on success and (False, reason) on failure. The reason is used by __get

Reset KV cache for clean recache, preserving global sink.

Fused (LayerNorm + (1+e_scale)*x + e_shift) over [B, F*frame_seqlen, C]. Replaces the WanAttentionBlock norm1/norm2 + modulate pattern. Output dt

Build an ErrorBuffer from an OmegaConf/dict config node. When ``num_blocks > 0`` the buffer becomes 2D (position × timestep), enabling teache