EDiT

EDiT is a diffusion transformer (DiT) model architecture. For now, E means editted as it supports text prompts instead of class labels.

EDiT adopts SDXL-VAE and CLIP to encode images and text.

• SDXL-VAE
• CLIP-ViT-L/14

This repository contains:

• 🪐 A simple PyTorch implementation of EDiT
• 🛸 Training script on ImageNet with text prompt.

Setup

Please refer to DiT and PixArt-α.

Training

Training EDiT

accelerate launch --mixed_precision fp16 train.py --data_path /path/to/ImageNet/train

TODO

• Support text prompt for DiT
• Training script using accelerate
• Gradio for inference

Acknowledgments

EDiT has been greatly inspired by the following amazing works and teams:

• DiT: Scalable Diffusion Models with Transformers
• fast-DiT: Scalable Diffusion Models with Transformers
• PixArt-α: PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis