Xingjian Leng^1*   ·   Jaskirat Singh^1*   ·   Yunzhong Hou¹   ·   Zhenchang Xing²  ·   Saining Xie³  ·   Liang Zheng¹ 

¹ Australian National University   ²Data61-CSIRO   ³New York University  
_{^*Project Leads}

🌐 Project Page   🤗 Models   📃 Paper  

We address a fundamental question: Can latent diffusion models and their VAE tokenizer be trained end-to-end? While training both components jointly with standard diffusion loss is observed to be ineffective — often degrading final performance — we show that this limitation can be overcome using a simple representation-alignment (REPA) loss. Our proposed method, REPA-E, enables stable and effective joint training of both the VAE and the diffusion model.

REPA-E significantly accelerates training — achieving over 17× speedup compared to REPA and 45× over the vanilla training recipe. Interestingly, end-to-end tuning also improves the VAE itself: the resulting E2E-VAE provides better latent structure and serves as a drop-in replacement for existing VAEs (e.g., SD-VAE), improving convergence and generation quality across diverse LDM architectures. Our method achieves state-of-the-art FID scores on ImageNet 256×256: 1.12 with CFG and 1.69 without CFG.

We previously released the REPA-E VAE checkpoint, which required loading through the model class in our REPA-E repository.
This new version provides a Hugging Face–compatible AutoencoderKL checkpoint that can be loaded directly via the diffusers API — no extra code or custom wrapper needed.

It offers plug-and-play compatibility with diffusion pipelines and can be seamlessly used to build or train new diffusion models.

The following packages are required to load and run the REPA-E VAEs with the diffusers library:

Below is a minimal example showing how to load and use the REPA-E end-to-end trained SD-VAE with diffusers: