Explicitly release large KV / cross-attention caches to free GPU memory. Safe to call between independent inference calls; caches wil

Return a detached CPU state dict suitable for torch.save.

Return a high-precision tensor with the dequantized values (Triton kernel).

Distributed FlexAttention over SP-sharded heads/sequences. Inputs use the local SP layout [B, L_local, N_local, D]. The function gathers sequ

Find the latest checkpoint in the logdir.

Find the latest checkpoint in the logdir.

Dequantizes FP4 packed tensor using per-block scaling factors. Args: packed_tensor (torch.Tensor): Packed uint8 tensor of shape (M, N//2)

r""" Generates video frames from text prompt using diffusion process. Args: input_prompt (`str`): Text pr

Get the calibrated kwargs for the given model and scale rules. If this model has not yet been calibrated with these scale rules, an e

Get the kwargs for tasks that should be used to calibrate the given model for this PTQ method before running evaluation.

Return the quantization configuration for a given module.

Construct matrix used in random Hadamard transform.

Initialize sequence-parallel state for the current process.

Return True if the rule is static, False otherwise.

Return True if the backend supports the given input and quantization configuration.

Optimize a model with SpinQuant.

Prepare this layer for post-training quantization by quantizing the weight, storing the quantized weight, and deleting the original w

Register a new type of quantized module. Args: high_precision_cls: (`type[nn.Module]`): The high precision module to be

Test speed on a B200 on Modal.

Test speed on a B200 on Modal.

Sample one entry uniformly from ALL buckets (legacy SVI E_img).

calculate the number of frames for video used for model inputs. Args: ele (dict): a dict contains the configuration of video.

Start the evaluation coordinator.

Return the corresponding torch.dtype if one is available, or `None` otherwise.

Convert [B, S_total, N_local, D] to [B, S_local, N, D].

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

Initialize the FourOverSixGptOssMLP layer. Args: module (GptOssMLP): The high-precision module that this quantized layer

Save current model attention state and apply pipeline overrides.

Backward simulation for bidirectional (non-causal) generator. All frames are processed at once at each denoising step — no KV cache,

Check if from (start_video_idx, start_frame) we can sample all segments without ever wrapping to the beginning (i.e. only use this video and l

Check if a folder has enough total frames across all videos to complete all segments. This is a lenient check: we verify that the tot

Compute the KL grad (eq 7 in https://arxiv.org/abs/2311.18828). Input: - noisy_image_or_video: a tensor with shape [B, F,

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

Simulate the generator's input from noise to avoid training/inference mismatch. See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/

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

r""" Forward pass through the diffusion model Args: x (List[Tensor]): List of input video tensors, each w

Split evenly into 2 chunks to avoid int32 overflow in CUDA FP4 quantize kernel.

r""" Run the diffusion model with kv caching. See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. T

r""" Forward pass through the diffusion model Args: x (List[Tensor]): List of input video tensors, each w

Full-sequence bidirectional multi-step denoising for visualization.

Return the caption directory path for a given folder (sample).

Cached TE DotProductAttention instance keyed by attention shape + masking. Constructed lazily and reused across forward calls. TE's DPA object