Scalable Diffusion Models with Transformers (DiT)

Official PyTorch Implementation

Paper | Project Page | Run DiT-XL/2

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.

Scalable Diffusion Models with Transformers
William Peebles, Saining Xie

UC Berkeley, New York University

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.

This repository contains:

An implementation of DiT directly in Hugging Face diffusers can also be found here.

Setup

First, download and set up the repo:

复制代码
git clone https://github.com/facebookresearch/DiT.git
cd DiT

We provide an 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.

复制代码
conda env create -f environment.yml
conda activate DiT

Sampling

Pre-trained DiT checkpoints. You can sample from our pre-trained DiT models with 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:

复制代码
python sample.py --image-size 512 --seed 1

For convenience, our pre-trained DiT models can be downloaded directly here as well:

DiT Model Image Resolution FID-50K Inception Score Gflops
XL/2 256x256 2.27 278.24 119
XL/2 512x512 3.04 240.82 525

Custom DiT checkpoints. If you've trained a new DiT model with train.py (see below), 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:

复制代码
python sample.py --model DiT-L/4 --image-size 256 --ckpt /path/to/model.pt

Training DiT

We provide a training script for DiT in 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:

复制代码
torchrun --nnodes=1 --nproc_per_node=N train.py --model DiT-XL/2 --data-path /path/to/imagenet/train

PyTorch Training Results

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:

DiT Model Train Steps FID-50K (JAX Training) FID-50K (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

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).

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.

Enhancements

Training (and sampling) could likely be sped-up significantly by:

  • using Flash Attention in the DiT model
  • using torch.compile in PyTorch 2.0

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

🔥 Feature Update Check out this repository at GitHub - chuanyangjin/fast-DiT: Fast Diffusion Models with Transformers 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.

相关推荐
阿坡RPA5 小时前
手搓MCP客户端&服务端:从零到实战极速了解MCP是什么?
人工智能·aigc
几米哥8 小时前
从思考到行动:AutoGLM沉思如何让AI真正"动"起来
llm·aigc·chatglm (智谱)
你觉得2059 小时前
哈尔滨工业大学DeepSeek公开课:探索大模型原理、技术与应用从GPT到DeepSeek|附视频与讲义下载方法
大数据·人工智能·python·gpt·学习·机器学习·aigc
头顶秃成一缕光11 小时前
若依——基于AI+若依框架的实战项目(实战篇(下))
java·前端·vue.js·elementui·aigc
weixin_4352081612 小时前
通过 Markdown 改进 RAG 文档处理
人工智能·python·算法·自然语言处理·面试·nlp·aigc
你觉得20512 小时前
浙江大学朱霖潮研究员:《人工智能重塑科学与工程研究》以蛋白质结构预测为例|附PPT下载方法
大数据·人工智能·机器学习·ai·云计算·aigc·powerpoint
360安全应急响应中心13 小时前
基于 RAG 提升大模型安全运营效率
安全·aigc
Zhikes13 小时前
抄了AI语音克隆的家!本地部署、6G显存搞定、一键启动包免费放送
aigc
墨风如雪13 小时前
深度解析 OpenAI Academy:官方下场,AI 学习迎来新基准?
aigc
KarudoLee16 小时前
AIGC7——AIGC驱动的视听内容定制化革命:从Sora到商业化落地
人工智能·aigc