Update README.md

README.md

@@ -4,28 +4,57 @@ library_name: diffusers
 pipeline_tag: image-to-image
 ---
 
----
-license: mit
-library_name: diffusers
-pipeline_tag: image-to-image
----
+<h1 align="center"> REPA-E: Unlocking VAE for End-to-End Tuning of Latent Diffusion Transformers </h1>
+
+<p align="center">
+  <a href="https://scholar.google.com.au/citations?user=GQzvqS4AAAAJ" target="_blank">Xingjian&nbsp;Leng</a><sup>1*</sup> &ensp; <b>&middot;</b> &ensp;
+  <a href="https://1jsingh.github.io/" target="_blank">Jaskirat&nbsp;Singh</a><sup>1*</sup> &ensp; <b>&middot;</b> &ensp;
+  <a href="https://hou-yz.github.io/" target="_blank">Yunzhong&nbsp;Hou</a><sup>1</sup> &ensp; <b>&middot;</b> &ensp;
+  <a href="https://people.csiro.au/X/Z/Zhenchang-Xing/" target="_blank">Zhenchang&nbsp;Xing</a><sup>2</sup>&ensp; <b>&middot;</b> &ensp;
+  <a href="https://www.sainingxie.com/" target="_blank">Saining&nbsp;Xie</a><sup>3</sup>&ensp; <b>&middot;</b> &ensp;
+  <a href="https://zheng-lab-anu.github.io/" target="_blank">Liang&nbsp;Zheng</a><sup>1</sup>&ensp;
+</p>
+
+<p align="center">
+  <sup>1</sup> Australian National University &emsp; <sup>2</sup>Data61-CSIRO &emsp; <sup>3</sup>New York University &emsp; <br>
+  <sub><sup>*</sup>Project Leads&emsp;</sub>
+</p>
+
+<p align="center">
+  <a href="https://End2End-Diffusion.github.io">🌐 Project Page</a> &ensp;
+  <a href="https://huggingface.co/REPA-E">🤗 Models</a> &ensp;
+  <a href="https://arxiv.org/abs/2504.10483">📃 Paper</a> &ensp;
+  <br>
+  <a href="https://paperswithcode.com/sota/image-generation-on-imagenet-256x256?p=repa-e-unlocking-vae-for-end-to-end-tuning-of"><img src="https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/repa-e-unlocking-vae-for-end-to-end-tuning-of/image-generation-on-imagenet-256x256" alt="PWC"></a>
+</p>
+
+
+<p align="center">
+  <img src="https://github.com/End2End-Diffusion/REPA-E/raw/main/assets/vis-examples.jpg" width="100%" alt="teaser">
+</p>
 
-# REPA-E: Unlocking VAE for End-to-End Tuning of Latent Diffusion Transformers
+---
 
-## About
-This model addresses the question of whether latent diffusion models and their VAE tokenizer can be trained end-to-end.  Using a representation-alignment (REPA) loss, REPA-E enables stable and effective joint training of both components, leading to significant training acceleration and improved VAE performance. The resulting E2E-VAE serves as a drop-in replacement for existing VAEs, improving convergence and generation quality across diverse LDM architectures.
+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.
 
-This model is based on the paper [REPA-E: Unlocking VAE for End-to-End Tuning of Latent Diffusion Transformers](https://huggingface.co/papers/2504.10483) and its official implementation is available on [Github](https://github.com/REPA-E/REPA-E). The project page can be found at [https://end2end-diffusion.github.io](https://end2end-diffusion.github.io).
+<p align="center">
+  <img src="https://github.com/End2End-Diffusion/REPA-E/raw/main/assets/overview.jpg" width="100%" alt="teaser">
+</p>
 
-## Usage
+**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.26** with CFG and **1.83** without CFG.
 
-To use the REPA-E model, you can load it via the Hugging Face `DiffusionPipeline`. Below is a simplified example of how to use a pretrained REPA-E model for inference. For training examples and further details, please refer to the [Github repository](https://github.com/REPA-E/REPA-E).
+ 
+## Usage and Training
 
-```python
-from diffusers import DiffusionPipeline
+Please refer our [Github Repo](https://github.com/End2End-Diffusion/REPA-E) for detailed notes on end-to-end training and inference using REPA-E.
 
-pipeline = DiffusionPipeline.from_pretrained("REPA-E/sit-repae-sdvae", trust_remote_code=True)
-image = pipeline().images[0]
+## 📚 Citation
 
-image.save("generated_image.png")
-```
+```bibtex
+@article{leng2025repae,
+  title={REPA-E: Unlocking VAE for End-to-End Tuning with Latent Diffusion Transformers},
+  author={Xingjian Leng and Jaskirat Singh and Yunzhong Hou and Zhenchang Xing and Saining Xie and Liang Zheng},
+  year={2025},
+  journal={arXiv preprint arXiv:2504.10483},
+}
+```