Initialize the RMSNorm normalization layer. Args: dim (int): The dimension of the input tensor. eps (float,

Create zero tensor(s) matching x, which can be Tensor or list of Tensors.

Return sorted, deduplicated starting positions that cover *length*.

Return sorted, deduplicated starting positions that cover *length*.

Same as _prepare_cfg_conditions but uses args.weak_cond_strength_aelq (scalar) for distill losses.

x[indices] = src[indices] for Tensor or list of Tensors.

Compute teacher/self-distill target v with CFG.

img_gt: GPU Tensor of shape (B, C, H, W) in [0, 1].

embed_dim: output dimension for each position pos: a list of positions to be encoded: size (M,) out: (M, D)

torch.repeat_interleave(x, dim=2, repeats=n_rep)

Set requires_grad flag for all parameters in a model.

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

2-D Gaussian blend mask (1, C, tile_h, tile_w) peaked at the centre.

2-D Gaussian blend mask (1, C, tile_h, tile_w) peaked at the centre.

Dynamically generate RoPE for variable-size inputs. Args: hw_seq_len: patch grid side length (H // patch_size)

Apply the RMSNorm normalization to the input tensor. Args: x (torch.Tensor): The input tensor. Returns:

Prepare lq_weak/lq_noised/lq_mixed and z_weak/z_noised/z_mixed for CFG training. Fixed: use_aelq=True, use_venc=True, weak_cond_strength_venc=

Compute RCGM consistency loss.

Given a weight file like .../checkpoint-00010000/clean_weights/ema_model.safetensors, walk up to find the checkpoint-XXXXXXXX directory and p

Given a weight file like .../checkpoint-00010000/clean_weights/ema_model.safetensors, walk up to find the checkpoint-XXXXXXXX directory and p

Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings to the given query 'xq' a

Create a logger that writes to a log file and stdout.

Create a logger that writes to a log file and stdout.

Resolve checkpoint path and parse global step. - If args.resume_ckpt is a .safetensors / .pth file, trace up to find the checkpoint-XXX

Resolve checkpoint path and parse global step. - If args.resume_ckpt is a .safetensors / .pth file, trace up to find the checkpoint-XXX

Support OrderedDict for yaml. Returns: yaml Loader and Dumper.

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

Set requires_grad flag for all parameters in a model.

Reshape frequency tensor for broadcasting it with another tensor. This function reshapes the frequency tensor to have the same shape as the

Latent-space tiled inference for VOSR DiT (multi-step FM with CFG). 1. VAE-encode the full image once. 2. For each tile, crop the corres

Latent-space tiled inference for VOSR DiT. 1. VAE-encode the full image once (no tile seams from encoder). 2. For each tile, crop the co

Create sinusoidal timestep embeddings. Args: t: A 1-D Tensor of N indices, one per batch element. These may be fractional

Step the EMA model towards the current model. Handles DDP (module.) and torch.compile (_orig_mod.) name prefixes.

Step the EMA model towards the current model. Strips DDP (module.) and torch.compile (_orig_mod.) name prefixes.