Diffusion model with a Transformer backbone.

Lightweight VAE Decoder for SD2 with halved channels. Accepts latent z of shape [B, 4, H/8, W/8] and outputs RGB image [B, 3, H, W].

Embeds scalar timesteps into vector representations. Same as DiT.

Outputs HQ images as Tensors in [0,1]. Degradation is applied in the main training loop as a batched GPU operation.

Attention module of LightningDiT.

Multi-head attention module.

Reads images, applies geometric transforms, and outputs [0,1] Tensors. Degradation is handled in the GPU training loop.

The final layer of LightningDiT.

Lightning DiT Block. We add features including: - ROPE - QKNorm - RMSNorm - SwiGLU - No shift AdaLN. Not all of them ar

FileHandler that silently ignores flush errors on filesystems that don't support it.

2D Qwen/Wan VAE for image-only use cases. API: encode(x) -> latent_dist, decode(z) -> sample, forward(sample) -> reconstruct.