Generate image/videos from noise and train the critic with generated samples. The noisy input to the generator is backward simulated.

ID for the rule in the CUDA implementation.

Decode latent to pixel video. Args: vae: VAE wrapper exposing `decode_to_pixel(latent)` latent: shape (batch, T, C, H, W) or

Decode latent to pixel video; same logic as pipeline/causal_diffusion_inference.py. Args: vae: Wan2.2-TI2V-5B VAE wrapper la

Dequantize multiple AR KV-cache chunks with one CUDA launch.

Performs distributed attention based on DeepSpeed Ulysses attention mechanism. please refer to https://arxiv.org/pdf/2309.14509 Args:

One step for the first-order DPMSolver (equivalent to DDIM). Args: model_output (`torch.Tensor`): The dir

Evaluate a model with a given PTQ method.

Evaluate a quantized model with lm-eval.

r""" Args: x(Tensor): Shape [B, L1, C] e(Tensor): Shape [B, L1, C]

Perform a matrix multiplication (`a @ b.T`) between two quantized tensors using the backend.

Get all checkpoints in the logdir sorted by step number.

Get all checkpoints in the logdir sorted by step number.

Get the quantized module for a given high-precision module.

Add -gencode flags based on nvcc capabilities. Uses the following rules: - sm_100/120 on CUDA >= 12.8 - Use 100f on CUDA >= 12.9

Construct a 16x16 Hadamard matrix.

Non-random sign vector for Hadamard transform.

Get the remote URL of the submodule.

Return the Wan2.2-TI2V-5B VAE wrapper.

Get the weight conversion for a given model type determined by the model config type.

Hard-coded random signs for Hadamard transform. https://xkcd.com/221/

Check if the submodule has untracked or unstaged changes.

r""" Generates video frames from input image and text prompt using diffusion process. Args: input_prompt (`str`):

Initialize distributed process group if launched via torchrun.

Initialize distributed process group if launched via torchrun. Returns (rank, world_size).

r""" Initialize model parameters using Xavier initialization.

r""" Initialize model parameters using Xavier initialization.

Return True if the backend is available on the current machine.

Map a latent range to the raw-frame window needed by Wan VAE encode.

Load a LongLive generator checkpoint into ``generator``.

One step for the second-order multistep DPMSolver. Args: model_output_list (`List[torch.Tensor]`): The di

One step for the third-order multistep DPMSolver. Args: model_output_list (`List[torch.Tensor]`): The dir

One step for the UniC (B(h) version). Args: this_model_output (`torch.Tensor`): The model outputs at `x_

One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified. Args: model_output (`torch

Move ``pipeline.vae`` to ``config.vae_device`` for streaming-pipeline decode. Only acts when both ``streaming_vae`` and ``vae_device`` are set; o

A few custom initialization steps that should be called after the object is created. Currently, the only one we have is to bind a few

Quantize a model.

Replace eligible modules with TransformerEngine NVFP4 wrappers.

Quantize a model with FourOverSix by default, or replace `nn.Linear` with TransformerEngine wrappers when `use_transformer_engine=True`.

Quantize a tensor to FP4 using the backend. Args: x (torch.Tensor): The input tensor to quantize. config (Qu

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

Return quantized parameters as QuantizedTensor.

Return quantized parameters as QuantizedTensor.

Apply rotary embedding via Triton kernel. iter-46: when `seq_len < x.size(0)`, the kernel rotates the first `seq_len` rows and copies through

Run any tasks that should be used to calibrate models for a given PTQ method and set of parameters before running evaluation.

Run any tasks that should be used to calibrate models for a given PTQ method and set of parameters before running evaluation.

For 2D buffers: sample at the given position, with random timestep. This is the natural choice for context (E_img) injection — the clean