Initialize the Attention module. Args: args (ModelArgs): Model configuration parameters. Attributes:

Initialize the FeedForward module. Args: dim (int): Input dimension. hidden_dim (int): Hidden dimension of t

Initialize a TransformerBlock. Args: layer_id (int): Identifier for the layer. args (ModelArgs): Model confi

Initialize a Transformer model. Args: params (ModelArgs): Model configuration parameters. Attributes:

Crop DINOv2 feature list from latent-space coordinates.

Collate that keeps only Tensor-like values and drops strings. Prevents metadata keys like 'base_name' from breaking accelerate.

Collate that keeps only Tensor (and numeric) fields and drops strings. Prevents metadata like 'base_name' from breaking accelerate.

Args: z: Latent tensor of shape [B, in_channels, H', W'] Returns: Decoded image of shape [B, out_channels, H,

Forward pass of LightningDiT. x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) t: (N,) t

Forward pass through the RMSNorm layer. Args: x (torch.Tensor): The input tensor. Returns: torch.Te

Forward pass of the attention module. Args: x (torch.Tensor): Input tensor. start_pos (int): Starting positi

Perform a forward pass through the TransformerBlock. Args: x (torch.Tensor): Input tensor. start_pos (int):

Perform a forward pass through the Transformer model. Args: tokens (torch.Tensor): Input token indices. star

Forward pass that supports variable input sizes. Unlike the standard forward, this dynamically generates RoPE to accommodate

grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/

Interpolate 2D positional embeddings via bicubic resize. pos_embed: [1, old_H*old_W, D] new_size: (new_H, new_W) old_size: (old_H, ol

# random color jitter if np.random.uniform() < self.opt['color_jitter_prob']: jitter_val = np.random.uniform(-shift, shi