$ show LTH14/mar
This repository provides the official PyTorch implementation of MAR (Masked Autoregressive with Diffusion Loss), a research method for class-conditional image generation on ImageNet 256x256. The core contribution is an autoregressive model that operates on continuous-valued latent tokens (via a pre-trained VAE) and uses a diffusion loss instead of vector quantization. The codebase includes model definitions (MAR encoder-decoder transformer, DiffLoss MLP), training and evaluation scripts with PyTorch DDP, a VAE latent caching pipeline, and pre-trained checkpoints for base/large/huge variants. While primarily a research prototype, it contains reusable components such as the DiffLoss module, diffusion utilities, and a clean data caching pattern. The Gradio app and Colab notebook make the model accessible for interactive use.
The LLM identified 5 files worth highlighting for reuse:
Implements the diffusion loss module used to predict continuous latent tokens. Contains a SimpleMLPAdaLN network with timestep embedding, conditional embedding, and residual blocks with adaptive layer norm. The forward method computes training loss using a cosine noise schedule, and the sample method performs ancestral sampling (with optional CFG). This is a self-contained, well-structured module that can be extracted and reused in other projects requiring a diffusion loss on continuous features.
Contains the CachedFolder class, a PyTorch Dataset that loads pre-computed VAE latent moments from .npz files and applies random horizontal flipping at load time (by selecting either moments or moments_flip). This pattern is highly reusable for any project that pre-computes features and wants to apply simple augmentations without re-encoding. The ImageFolderWithFilename class is also included, which returns filenames alongside images and labels, useful for tracking.
Provides a single function, center_crop_arr, which implements the center cropping routine used in ADM (guided diffusion). It downscales the image if it is too large, then resizes so the smaller dimension matches the target size, and finally crops the center. This utility is standalone and can be reused in any computer vision pipeline requiring center cropping.
A distributed script that pre-computes VAE latents for the entire ImageNet training set and saves them as .npz files (including both original and horizontally flipped moments). This is an excellent example of a reusable ETL pattern: it uses a custom dataset class with filename tracking, processes images in parallel across GPUs, and caches results to disk. The resulting cache can then be used by the training script to avoid repeated VAE encoding, significantly speeding up training.
Factory function for creating SpacedDiffusion instances with configurable noise schedules (linear or cosine) and timestep respacing. This module wraps the full Gaussian diffusion utilities from the diffusion folder and provides a clean interface for both training and generation. While adapted from DiT, the organization is clean and the create_diffusion function is a reusable pattern for any project needing a diffusion process.