$ show bytetriper/RAE
This repository implements a two-stage generative image model: Representation Autoencoders (RAE) followed by Diffusion Transformers (DiT^DH). Stage 1 trains a ViT decoder to reconstruct images from frozen pretrained encoders (DINOv2, MAE, SigLIP), with optional GAN and LPIPS losses. Stage 2 trains a class-conditional diffusion transformer on the resulting latent space, supporting both classifier-free guidance and autoguidance. The codebase is config-driven (OmegaConf), PyTorch DDP, and provides training, distributed sampling, FID evaluation, and latent statistics computation. It is a research codebase with strong engineering quality, modular design (transport, discriminator, encoders, decoders), and reproducibility through pinned dependencies and checkpointing.
The LLM identified 2 files worth highlighting for reuse:
Defines Transport class that abstracts the forward diffusion/flow process with configurable path types (Linear, GVP, VP), model outputs (velocity, score, noise), and loss weighting. Provides training_losses() and get_drift()/get_score() methods. The companion integrators (ode, sde classes in integrators.py) and path plans (ICPlan, VPCPlan, GVPCPlan in path.py) are self-contained and reusable for any diffusion/flow matching project. This module is cleanly separated from model architectures and can be extracted with minimal adaptation.
Provides a factory function build_discriminator() that creates a DinoDiscriminator (frozen DINO backbone + trainable heads) and a DiffAug augmentation policy. The DinoDiscriminator and DiffAug are independently reusable components for GAN-based representation learning. The module encapsulates all discriminator logic without coupling to the RAE architecture.