q: [B, Lq, Nq, C1]. k: [B, Lk, Nk, C1]. v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. q_lens

config: class_name: str (module.class) other_params, etc. will be passed as a whole config to the class

modified from https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py#L49 Major change: G is [b, d, d] rather than [d, d]

freqs: list of 3 tensors, freq_f, freq_h, freq_w of shape [1024, 22], [1024, 21], [1024, 21] for head_dim=64. Does not support variable

grid_sizes: list of 3 ints, F, H, W where F is the number of latent frames for each video in the batch. freqs: list of 3 tensors, fr

freqs: list of 3 tensors, freq_f, freq_h, freq_w of shape [1024, 22], [1024, 21], [1024, 21] for head_dim=64. Does not support variable

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

x of shape [B, f, c, h, w]

x: [b, l, d]

This is an updated version of the zeropower_via_newtonschulz5 function in here: https://github.com/KellerJordan/modded-nanogpt/blob/master/tr

Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'.

Args: W0_W2: [B, 2 * Hidden, D] W1: [B, D, Hidden] X: [B, num_Tokens, D] Outs: Hidden: [B, num_token

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

Args: dy: [b, d, l], gradient of the outer loss wrt the y x: [b, d, l], input of the silu activation outs: dx: [b, d,

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

This is an updated version of the zeropower_via_newtonschulz5 function in here: https://github.com/KellerJordan/modded-nanogpt/blob/master/tr

Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.

Args: grad_out: [B, num_tokens, D] Outs: grad_W0_W2: [B, 2 * Hidden, D] grad_W1: [B, D, Hidden]

Transform target before computing loss Args: fxfycxcy (torch.tensor): [b, v, 4] c2w (torch.tensor): [b, v, 4, 4] Returns:

Input: - videos_rgb: a tensor with shape [B, F+1, C, H, W] in RGB format, [0-1] Output: - video_rgb: a tensor

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

hf api is not reliable, retry when failed with max tries :param repo_path: The path of the repo to download :param odir: output path

x: [B, T, D] x_add: [B, T, D] tgt_scale: [B, T] eps: float tgt_dtype: if None, will use the same dtype as x

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

Adjusts the state of each module in the object. Parameters: - module_grad (dict): A dictionary where each key is the name of

Sets value in nested dictionary

This is an updated version of the zeropower_via_newtonschulz5 function in here: https://github.com/KellerJordan/modded-nanogpt/blob/master/tr

This is an updated version of the zeropower_via_newtonschulz5 function in here: https://github.com/KellerJordan/modded-nanogpt/blob/master/tr

modified from https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py#L49 Major change: G is [b, d, d] rather than [d, d]

"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

Remove older checkpoint directories, keeping only the most recent ones within the specified iteration range. Args: keep_last

latents: [batch_size, num_frames, num_channels, height, width] output: [batch_size, num_channels, num_frames, height, width]

Exports a video from a numpy array of frames. Args: video_frames (np.ndarray): Video frames in numpy format with shape [num_fram

Args: W0_W2: [B, 2 * Hidden, D] W1: [B, D, Hidden] K, V: [B, num_Tokens, D] lr0, lr1, lr2: [B, N]

W0, W1: [B, K, M] (bf16/fp16/fp32) X0T, X1T: [B, N, K] Returns O: [B, M, N] with O = W0^T @ X0T.T + W1^T @ X1T.T

W0, W1: [B, K, M] X0, X1: [B, K, N] Returns O: [B, M, N] with O = W0^T @ X0 + W1^T @ X1

Wraps the Triton kernel. Shapes: W0, W2: [B, M, K] (bf16) X : [B, N, K] (bf16) returns O: [B, M, N] (bf16) where O = SiLU

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

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

x: [b, l, d]

x: [b, l, d]

x: [b, l, d]

When launched with torchrun --nproc_per_node=<> --nnodes=<> --rdzv-endpoint=${master_addr}:${master_port} The following environment variable

Compute the integral of the logit normal distribution over unit intervals for given time steps. This is a legacy version that works directly

freqs: list of 3 tensors, freq_f, freq_h, freq_w of shape [1024, 22], [1024, 21], [1024, 21] for head_dim=64.

Args: x: [B, D1, D2], must be fp32 dx: [B, D1, D2], bf16 dx_momentum: [B, D1, 1], fp32 recommended momentum_mu: f