Update README.md

README.md

@@ -1,198 +1,166 @@
 ---
 license: mit
 ---
+## Scalable Diffusion Models with Transformers (DiT)<br><sub>Official PyTorch Implementation</sub>
 
-# Model Card for Model ID
+### [Paper](http://arxiv.org/abs/2212.09748) | [Project Page](https://www.wpeebles.com/DiT) | Run DiT-XL/2 [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/wpeebles/DiT) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](http://colab.research.google.com/github/facebookresearch/DiT/blob/main/run_DiT.ipynb) <a href="https://replicate.com/arielreplicate/scalable_diffusion_with_transformers"><img src="https://replicate.com/arielreplicate/scalable_diffusion_with_transformers/badge"></a>
 
-<!-- Provide a quick summary of what the model is/does. -->
+![DiT samples](visuals/sample_grid_0.png)
 
-This modelcard aims to be a base template for new models. It has been generated using [this raw template](https://github.com/huggingface/huggingface_hub/blob/main/src/huggingface_hub/templates/modelcard_template.md?plain=1).
+This repo contains PyTorch model definitions, pre-trained weights and training/sampling code for our paper exploring 
+diffusion models with transformers (DiTs). You can find more visualizations on our [project page](https://www.wpeebles.com/DiT).
 
-## Model Details
+> [**Scalable Diffusion Models with Transformers**](https://www.wpeebles.com/DiT)<br>
+> [William Peebles](https://www.wpeebles.com), [Saining Xie](https://www.sainingxie.com)
+> <br>UC Berkeley, New York University<br>
 
-### Model Description
+We train latent diffusion models, replacing the commonly-used U-Net backbone with a transformer that operates on 
+latent patches. We analyze the scalability of our Diffusion Transformers (DiTs) through the lens of forward pass 
+complexity as measured by Gflops. We find that DiTs with higher Gflops---through increased transformer depth/width or
+increased number of input tokens---consistently have lower FID. In addition to good scalability properties, our 
+DiT-XL/2 models outperform all prior diffusion models on the class-conditional ImageNet 512×512 and 256×256 benchmarks, 
+achieving a state-of-the-art FID of 2.27 on the latter.
 
-<!-- Provide a longer summary of what this model is. -->
+This repository contains:
 
+* 🪐 A simple PyTorch [implementation](models.py) of DiT
+* ⚡️ Pre-trained class-conditional DiT models trained on ImageNet (512x512 and 256x256)
+* 💥 A self-contained [Hugging Face Space](https://huggingface.co/spaces/wpeebles/DiT) and [Colab notebook](http://colab.research.google.com/github/facebookresearch/DiT/blob/main/run_DiT.ipynb) for running pre-trained DiT-XL/2 models
+* 🛸 A DiT [training script](train.py) using PyTorch DDP
 
+An implementation of DiT directly in Hugging Face `diffusers` can also be found [here](https://github.com/huggingface/diffusers/blob/main/docs/source/en/api/pipelines/dit.mdx).
 
-- **Developed by:** [More Information Needed]
-- **Funded by [optional]:** [More Information Needed]
-- **Shared by [optional]:** [More Information Needed]
-- **Model type:** [More Information Needed]
-- **Language(s) (NLP):** [More Information Needed]
-- **License:** [More Information Needed]
-- **Finetuned from model [optional]:** [More Information Needed]
 
-### Model Sources [optional]
+## Setup
 
-<!-- Provide the basic links for the model. -->
+First, download and set up the repo:
 
-- **Repository:** [More Information Needed]
-- **Paper [optional]:** [More Information Needed]
-- **Demo [optional]:** [More Information Needed]
+```bash
+git clone https://github.com/facebookresearch/DiT.git
+cd DiT
+```
 
-## Uses
+We provide an [`environment.yml`](environment.yml) file that can be used to create a Conda environment. If you only want 
+to run pre-trained models locally on CPU, you can remove the `cudatoolkit` and `pytorch-cuda` requirements from the file.
 
-<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+```bash
+conda env create -f environment.yml
+conda activate DiT
+```
 
-### Direct Use
 
-<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+## Sampling [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/wpeebles/DiT) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](http://colab.research.google.com/github/facebookresearch/DiT/blob/main/run_DiT.ipynb)
+![More DiT samples](visuals/sample_grid_1.png)
 
-[More Information Needed]
+**Pre-trained DiT checkpoints.** You can sample from our pre-trained DiT models with [`sample.py`](sample.py). Weights for our pre-trained DiT model will be 
+automatically downloaded depending on the model you use. The script has various arguments to switch between the 256x256
+and 512x512 models, adjust sampling steps, change the classifier-free guidance scale, etc. For example, to sample from
+our 512x512 DiT-XL/2 model, you can use:
 
-### Downstream Use [optional]
+```bash
+python sample.py --image-size 512 --seed 1
+```
 
-<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+For convenience, our pre-trained DiT models can be downloaded directly here as well:
 
-[More Information Needed]
+| DiT Model     | Image Resolution | FID-50K | Inception Score | Gflops | 
+|---------------|------------------|---------|-----------------|--------|
+| [XL/2](https://dl.fbaipublicfiles.com/DiT/models/DiT-XL-2-256x256.pt) | 256x256          | 2.27    | 278.24          | 119    |
+| [XL/2](https://dl.fbaipublicfiles.com/DiT/models/DiT-XL-2-512x512.pt) | 512x512          | 3.04    | 240.82          | 525    |
 
-### Out-of-Scope Use
 
-<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+**Custom DiT checkpoints.** If you've trained a new DiT model with [`train.py`](train.py) (see [below](#training-dit)), you can add the `--ckpt`
+argument to use your own checkpoint instead. For example, to sample from the EMA weights of a custom 
+256x256 DiT-L/4 model, run:
 
-[More Information Needed]
+```bash
+python sample.py --model DiT-L/4 --image-size 256 --ckpt /path/to/model.pt
+```
 
-## Bias, Risks, and Limitations
 
-<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+## Training DiT
 
-[More Information Needed]
+We provide a training script for DiT in [`train.py`](train.py). This script can be used to train class-conditional 
+DiT models, but it can be easily modified to support other types of conditioning. To launch DiT-XL/2 (256x256) training with `N` GPUs on 
+one node:
 
-### Recommendations
+```bash
+torchrun --nnodes=1 --nproc_per_node=N train.py --model DiT-XL/2 --data-path /path/to/imagenet/train
+```
 
-<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+### PyTorch Training Results
 
-Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+We've trained DiT-XL/2 and DiT-B/4 models from scratch with the PyTorch training script
+to verify that it reproduces the original JAX results up to several hundred thousand training iterations. Across our experiments, the PyTorch-trained models give 
+similar (and sometimes slightly better) results compared to the JAX-trained models up to reasonable random variation. Some data points:
 
-## How to Get Started with the Model
+| DiT Model  | Train Steps | FID-50K<br> (JAX Training) | FID-50K<br> (PyTorch Training) | PyTorch Global Training Seed |
+|------------|-------------|----------------------------|--------------------------------|------------------------------|
+| XL/2       | 400K        | 19.5                       | **18.1**                       | 42                           |
+| B/4        | 400K        | **68.4**                   | 68.9                           | 42                           |
+| B/4        | 400K        | 68.4                       | **68.3**                       | 100                          |
 
-Use the code below to get started with the model.
+These models were trained at 256x256 resolution; we used 8x A100s to train XL/2 and 4x A100s to train B/4. Note that FID 
+here is computed with 250 DDPM sampling steps, with the `mse` VAE decoder and without guidance (`cfg-scale=1`). 
 
-[More Information Needed]
+**TF32 Note (important for A100 users).** When we ran the above tests, TF32 matmuls were disabled per PyTorch's defaults. 
+We've enabled them at the top of `train.py` and `sample.py` because it makes training and sampling way way way faster on 
+A100s (and should for other Ampere GPUs too), but note that the use of TF32 may lead to some differences compared to 
+the above results.
 
-## Training Details
+### Enhancements
+Training (and sampling) could likely be sped-up significantly by:
+- [ ] using [Flash Attention](https://github.com/HazyResearch/flash-attention) in the DiT model
+- [ ] using `torch.compile` in PyTorch 2.0
 
-### Training Data
+Basic features that would be nice to add:
+- [ ] Monitor FID and other metrics
+- [ ] Generate and save samples from the EMA model periodically
+- [ ] Resume training from a checkpoint
+- [ ] AMP/bfloat16 support
 
-<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+**🔥 Feature Update** Check out this repository at https://github.com/chuanyangjin/fast-DiT to preview a selection of training speed acceleration and memory saving features including gradient checkpointing, mixed precision training and pre-extrated VAE features. With these advancements, we have achieved a training speed of 0.84 steps/sec for DiT-XL/2 using just a single A100 GPU.
 
-[More Information Needed]
+## Evaluation (FID, Inception Score, etc.)
 
-### Training Procedure
+We include a [`sample_ddp.py`](sample_ddp.py) script which samples a large number of images from a DiT model in parallel. This script 
+generates a folder of samples as well as a `.npz` file which can be directly used with [ADM's TensorFlow
+evaluation suite](https://github.com/openai/guided-diffusion/tree/main/evaluations) to compute FID, Inception Score and
+other metrics. For example, to sample 50K images from our pre-trained DiT-XL/2 model over `N` GPUs, run:
 
-<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+```bash
+torchrun --nnodes=1 --nproc_per_node=N sample_ddp.py --model DiT-XL/2 --num-fid-samples 50000
+```
 
-#### Preprocessing [optional]
+There are several additional options; see [`sample_ddp.py`](sample_ddp.py) for details. 
 
-[More Information Needed]
 
+## Differences from JAX
 
-#### Training Hyperparameters
+Our models were originally trained in JAX on TPUs. The weights in this repo are ported directly from the JAX models. 
+There may be minor differences in results stemming from sampling with different floating point precisions. We re-evaluated 
+our ported PyTorch weights at FP32, and they actually perform marginally better than sampling in JAX (2.21 FID 
+versus 2.27 in the paper).
 
-- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
 
-#### Speeds, Sizes, Times [optional]
+## BibTeX
 
-<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+```bibtex
+@article{Peebles2022DiT,
+  title={Scalable Diffusion Models with Transformers},
+  author={William Peebles and Saining Xie},
+  year={2022},
+  journal={arXiv preprint arXiv:2212.09748},
+}
+```
 
-[More Information Needed]
 
-## Evaluation
+## Acknowledgments
+We thank Kaiming He, Ronghang Hu, Alexander Berg, Shoubhik Debnath, Tim Brooks, Ilija Radosavovic and Tete Xiao for helpful discussions. 
+William Peebles is supported by the NSF Graduate Research Fellowship.
 
-<!-- This section describes the evaluation protocols and provides the results. -->
+This codebase borrows from OpenAI's diffusion repos, most notably [ADM](https://github.com/openai/guided-diffusion).
 
-### Testing Data, Factors & Metrics
 
-#### Testing Data
-
-<!-- This should link to a Dataset Card if possible. -->
-
-[More Information Needed]
-
-#### Factors
-
-<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
-
-[More Information Needed]
-
-#### Metrics
-
-<!-- These are the evaluation metrics being used, ideally with a description of why. -->
-
-[More Information Needed]
-
-### Results
-
-[More Information Needed]
-
-#### Summary
-
-
-
-## Model Examination [optional]
-
-<!-- Relevant interpretability work for the model goes here -->
-
-[More Information Needed]
-
-## Environmental Impact
-
-<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
-
-Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
-
-- **Hardware Type:** [More Information Needed]
-- **Hours used:** [More Information Needed]
-- **Cloud Provider:** [More Information Needed]
-- **Compute Region:** [More Information Needed]
-- **Carbon Emitted:** [More Information Needed]
-
-## Technical Specifications [optional]
-
-### Model Architecture and Objective
-
-[More Information Needed]
-
-### Compute Infrastructure
-
-[More Information Needed]
-
-#### Hardware
-
-[More Information Needed]
-
-#### Software
-
-[More Information Needed]
-
-## Citation [optional]
-
-<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
-
-**BibTeX:**
-
-[More Information Needed]
-
-**APA:**
-
-[More Information Needed]
-
-## Glossary [optional]
-
-<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
-
-[More Information Needed]
-
-## More Information [optional]
-
-[More Information Needed]
-
-## Model Card Authors [optional]
-
-[More Information Needed]
-
-## Model Card Contact
-
-[More Information Needed]
+## License
+The code and model weights are licensed under CC-BY-NC. See [`LICENSE.txt`](LICENSE.txt) for details.