Guide · Open source
How to use DALLE2-pytorch: install, train and run it
Updated · written with AI from the project's own documentation and checked against our catalogue data
DALLE2-pytorch is an independent Python library that reimplements the components of OpenAI's DALL-E 2 text-to-image model in PyTorch. It is built for researchers and developers who want to train or study the architecture, not for typing a prompt and getting images. This guide shows you how to install the package, what the three-step training process looks like in code, and where the project's limits are.
DALL-E 2lucidrains/DALLE2-pytorch · ★ 11.3k · MIT Install options & downloads Quick start
- 1Check prerequisites. Confirm you have Python, PyTorch and a CUDA-capable GPU, since the examples call .cuda().
- 2Install the package. Run pip install dalle2-pytorch from the command line.
- 3Train CLIP. Build the CLIP model with dalle2_pytorch and train it on your text and image data using return_loss = True.
- 4Train the decoder. Wrap a Unet and your trained CLIP in a Decoder, then train it on images for many steps.
- 5Train the diffusion prior. Build a DiffusionPriorNetwork and wrap it with CLIP in a DiffusionPrior to learn image embeddings from text embeddings.
What DALLE2-pytorch is and who it is for
The README describes this repository as an implementation of OpenAI's DALL-E 2 in PyTorch. It focuses on the diffusion prior network, which takes CLIP text embeddings and predicts CLIP image embeddings. It also includes a decoder that turns those image embeddings into images.
It suits machine learning researchers, PyTorch developers, teams with GPUs, and students who want to see how CLIP-based diffusion text-to-image systems fit together.
- Best for: people replicating or studying the DALL-E 2 architecture
- Best for: developers who want to train their own text-to-image components
- Not for: anyone who wants ready-made image generation without training
What you need before you start
The README lists Python, PyTorch and a CUDA-capable GPU as prerequisites. The examples call .cuda() on models and tensors, so they will not run as written on a CPU-only machine.
The project is an independent implementation. The README notes that pre-trained checkpoints come from LAION and are still in progress, so do not expect official OpenAI DALL-E 2 weights.
- Python
- PyTorch
- A CUDA-capable GPU
Install dalle2-pytorch from PyPI
The package is published on PyPI. Install it with pip using the command from the README.
$ pip install dalle2-pytorchStep 1: train CLIP
The README says training has three steps, and CLIP is the most important. You can train CLIP with the x-clip package, or use the open_clip work underway in the LAION community. The repository demonstrates integration with x-clip to start.
The example imports CLIP from dalle2_pytorch, builds it with options such as dim_text, dim_image and visual_patch_size, and computes a contrastive loss on mock data. Replace the mock tensors with real text and image batches for actual training.
- Import: from dalle2_pytorch import CLIP
- Pass return_loss = True to get the contrastive loss, then call loss.backward()
import torch
from dalle2_pytorch import CLIP
clip = CLIP(
dim_text = 512,
dim_image = 512,
dim_latent = 512,
num_text_tokens = 49408,
text_enc_depth = 1,
text_seq_len = 256,
text_heads = 8,
visual_enc_depth = 1,
visual_image_size = 256,
visual_patch_size = 32,
visual_heads = 8,
use_all_token_embeds = True,
decoupled_contrastive_learning = True,
extra_latent_projection = True,
use_visual_ssl = True,
visual_ssl_type = 'simclr',
use_mlm = False,
text_ssl_loss_weight = 0.05,
image_ssl_loss_weight = 0.05
).cuda()
text = torch.randint(0, 49408, (4, 256)).cuda()
images = torch.randn(4, 3, 256, 256).cuda()
loss = clip(
text,
images,
return_loss = True
)
loss.backward()Step 2: train the decoder
The decoder learns to generate images from the image embeddings produced by the trained CLIP. The example builds a Unet, wraps it in a Decoder together with the CLIP model, and feeds images through it. The README says to repeat this for many, many steps.
The Decoder takes arguments such as timesteps, image_cond_drop_prob and text_cond_drop_prob. The example uses mock images, so you will need your own image dataset.
- Build a Unet for the decoder
- Wrap the Unet and CLIP in a Decoder
- Call decoder(images) and then loss.backward()
import torch
from dalle2_pytorch import Unet, Decoder, CLIP
unet = Unet(
dim = 128,
image_embed_dim = 512,
cond_dim = 128,
channels = 3,
dim_mults=(1, 2, 4, 8)
).cuda()
decoder = Decoder(
unet = unet,
clip = clip,
timesteps = 100,
image_cond_drop_prob = 0.1,
text_cond_drop_prob = 0.5
).cuda()
images = torch.randn(4, 3, 256, 256).cuda()
loss = decoder(images)
loss.backward()Step 3: train the diffusion prior
The diffusion prior is the main contribution of the paper, according to the README. It takes CLIP text embeddings and generates CLIP image embeddings. You need the trained CLIP from step one.
The example builds a DiffusionPriorNetwork, wraps it with the CLIP model in a DiffusionPrior, and uses mock text and images as input. The README notes that the prior network contains an autoregressive transformer in this example.
- Build a DiffusionPriorNetwork with dim, depth, dim_head and heads
- Wrap it with CLIP in a DiffusionPrior using timesteps and cond_drop_prob
import torch
from dalle2_pytorch import DiffusionPriorNetwork, DiffusionPrior, CLIP
prior_network = DiffusionPriorNetwork(
dim = 512,
depth = 6,
dim_head = 64,
heads = 8
).cuda()
diffusion_prior = DiffusionPrior(
net = prior_network,
clip = clip,
timesteps = 100,
cond_drop_prob = 0.2
).cuda()Check the project status and checkpoints before relying on it
The README reports that a research group used the code to train a functional diffusion prior, and that the decoder has been verified for unconditional generation on one experimental setup. It also says LAION is training prior models, with checkpoints on Hugging Face and training statistics on Weights & Biases.
The README also says the repository's model is no longer state of the art, and that the project will shift toward the Imagen architecture for Jax and text-to-video work. Read those notes before choosing this project for production.
- Prior checkpoints: LAION, hosted on Hugging Face
- Decoder: test runs are listed on Weights & Biases, marked in progress
What people use it for
- Replicating the DALL-E 2 prior and decoder architecture for a research paper
- Studying how CLIP embeddings feed a diffusion model
- Training a custom CLIP-to-image pipeline on your own dataset with GPUs
- Teaching a course on CLIP-based text-to-image systems
Frequently asked questions
Is DALLE2-pytorch free to use?
The repository is released under the MIT licence, according to the project listing. You can use, modify and redistribute it under that licence, but you still need your own compute and training data.
What do I need to run dalle2-pytorch?
The README lists Python, PyTorch and a CUDA-capable GPU as prerequisites. The examples call .cuda(), so a GPU is needed for them to run as written.
How do I install dalle2-pytorch?
Install it from PyPI with pip install dalle2-pytorch. The package name uses a hyphen, while the Python import uses dalle2_pytorch.
Can I generate images from a text prompt with dalle2-pytorch?
Not out of the box. The README describes a multi-step training process, and the project does not present a finished model for typing a prompt. You would need to train CLIP, the decoder and the diffusion prior first.
Does dalle2-pytorch include official OpenAI DALL-E 2 weights?
No. It is an independent implementation. The README says pre-trained checkpoints come from LAION and are still in progress.
How does dalle2-pytorch compare with imagen-pytorch or stable-diffusion?
This repository focuses on the DALL-E 2 diffusion prior and decoder. The README points to imagen-pytorch as the project it plans to move toward, since the Imagen architecture is simpler. Stable Diffusion is a separate, widely used model family that is not covered by this README.
New agent skills, repos and AI tools — every week
The best new skills and open-source projects, the fastest-growing repos, and one free prompt. Five minutes to read. Unsubscribe anytime.
Our verdict
DALLE2-pytorch suits researchers, PyTorch developers and students who want to study or train the DALL-E 2 prior and decoder on their own GPUs. It is a poor fit for anyone who wants to type a prompt and get images without training, or who needs official OpenAI weights. If you want a ready-to-use text-to-image tool, pick a finished model instead.
Install options & downloads