dalle2-pytorch

For training the diffusion prior, it is most efficient to use pre-computed embeddings via PriorEmbeddingDataset. This abstraction allows you to use the same training logic for both embedding-only and text-conditioned prior training.

Workflow:

1. Call get_reader() to create an EmbeddingReader object.
2. Use make_splits() to generate DataLoader objects for training, evaluation, and testing.

Distributed Training: make_splits() supports distributed training via the rank and world_size arguments. This ensures data is correctly partitioned across different processes. For single-process training, you can omit these or use the defaults (rank=0, world_size=1).

from dalle2_pytorch.dataloaders import get_reader, make_splits

IMG_URL = "data/img_emb/"
META_URL = "data/meta/"

# Create the reader
reader = get_reader(text_conditioned=True, img_url=IMG_URL, meta_url=META_URL)

# Configuration for splits
TRAIN_ARGS = {
    "world_size": 3,
    "text_conditioned": True,
    "start": 0,
    "num_data_points": 10000,
    "batch_size": 2,
    "train_split": 0.5,
    "eval_split": 0.25,
    "image_reader": reader,
}

# Create distributed splits
rank0_train, rank0_eval, rank0_test = make_splits(rank=0, **TRAIN_ARGS)
rank1_train, rank1_eval, rank1_test = make_splits(rank=1, **TRAIN_ARGS)
rank2_train, rank2_eval, rank2_test = make_splits(rank=2, **TRAIN_ARGS)

When training the decoder (or up samplers) in isolation, use ImageEmbeddingDataset to load images and their corresponding image embeddings.

This dataset supports two formats:

1. Webdataset format: A .tar file containing both .jpg and .npy files.
2. External embeddings: If embeddings are not in the webdataset, you can provide an embedding_folder_url. The folder must contain .npy files that match the shard numbers of the webdataset. For example, if a webdataset shard is 0001.tar and contains image 00010509.jpg, there must be a corresponding img_emb_0001.npy where the embedding for that image is located at index 509.

To generate compatible data, it is recommended to use img2dataset for the webdataset, clip-retrieval for embeddings, and embedding-dataset-reordering to ensure the correct format.

from dalle2_pytorch.dataloaders import ImageEmbeddingDataset

dataset = ImageEmbeddingDataset(
    urls="/path/or/url/to/webdataset/{0000..9999}.tar",
    embedding_folder_url="/path/or/url/to/embeddings/folder",
    shard_width=4,
    shuffle_shards=True,
    resample=False
)

The DiffusionPriorTrainer automates the training of a DiffusionPrior and manages its exponential moving average (EMA).

To train:

1. Instantiate CLIP, a DiffusionPriorNetwork (typically a transformer), and a DiffusionPrior.
2. Wrap the prior in a DiffusionPriorTrainer.
3. In the training loop, call diffusion_prior_trainer(text, images, max_batch_size=N).
4. Call diffusion_prior_trainer.update() to update the optimizer and the EMA.
5. Use diffusion_prior_trainer.sample(text, max_batch_size=N) to generate image embeddings from the EMA weights.

import torch
from dalle2_pytorch import DiffusionPriorNetwork, DiffusionPrior, DiffusionPriorTrainer, CLIP

clip = CLIP(...).cuda()

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

diffusion_prior_trainer = DiffusionPriorTrainer(
    diffusion_prior,
    lr = 3e-4,
    wd = 1e-2,
    ema_beta = 0.99,
    ema_update_after_step = 1000,
    ema_update_every = 10,
)

# Training loop
loss = diffusion_prior_trainer(text, images, max_batch_size = 4)
diffusion_prior_trainer.update()

# Sampling
image_embeds = diffusion_prior_trainer.sample(text, max_batch_size = 4)

The first step in training DALL-E 2 is training the CLIP model. This model provides the shared embedding space for text and images. You can initialize CLIP with various parameters for dimensionality, depth, and self-supervised learning (SSL) objectives like simclr or simsiam. To compute the contrastive loss during training, you must set return_loss = True when calling the model.

Key parameters for CLIP include:

• use_all_token_embeds: Enables fine-grained contrastive learning (FILIP).
• decoupled_contrastive_learning: Uses the DCL objective function (CLOOB + DCL).
• extra_latent_projection: Uses separate projections for text-to-image vs image-to-text (CLOOB).
• use_visual_ssl: Enables self-supervised learning on images.
• visual_ssl_type: Set to 'simclr' or 'simsiam'.

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

# mock data
text = torch.randint(0, 49408, (4, 256)).cuda()
images = torch.randn(4, 3, 256, 256).cuda()

# train
loss = clip(
    text,
    images,
    return_loss = True
)

loss.backward()

The Decoder can consist of multiple Unets (e.g., for cascading diffusion). Training them manually is complex because each Unet requires its own optimizer and exponential moving average (EMA). DecoderTrainer simplifies this by managing multiple optimizers and EMAs automatically.

To train, iterate through the unet_numbers (the index of the Unet in the Decoder.unet tuple) and call the trainer. Use max_batch_size to implement gradient accumulation.

Key steps:

1. Instantiate CLIP, Unets, and the Decoder.
2. Wrap the Decoder in a DecoderTrainer.
3. In a training loop, call decoder_trainer(images, text=text, unet_number=i, max_batch_size=N).
4. Call decoder_trainer.update(unet_number=i) to update the specific Unet and its EMA.
5. Use decoder_trainer.sample(...) to generate images from the EMA weights.

import torch
from dalle2_pytorch import Unet, Decoder, CLIP, DecoderTrainer

clip = CLIP(...).cuda()

unet1 = Unet(dim=128, image_embed_dim=512, text_embed_dim=512, cond_dim=128, channels=3, dim_mults=(1, 2, 4, 8), cond_on_text_encodings=True).cuda()
unet2 = Unet(dim=16, image_embed_dim=512, cond_dim=128, channels=3, dim_mults=(1, 2, 4, 8, 16)).cuda()

decoder = Decoder(
    unet = (unet1, unet2),
    image_sizes = (128, 256),
    clip = clip,
    timesteps = 1000
).cuda()

decoder_trainer = DecoderTrainer(
    decoder,
    lr = 3e-4,
    wd = 1e-2,
    ema_beta = 0.99,
    ema_update_after_step = 1000,
    ema_update_every = 10,
)

# Training loop
for unet_number in (1, 2):
    loss = decoder_trainer(
        images, 
        text = text, 
        unet_number = unet_number, 
        max_batch_size = 4
    )
    decoder_trainer.update(unet_number)

# Sampling
images = decoder_trainer.sample(image_embed = mock_image_embed, text = text)

You can train an unconditional DDPM or cascading DDPMs by setting unconditional = True in the Decoder configuration. This removes the requirement for text or image embeddings as conditioning during training.

To train:

1. Define Unets and a Decoder with unconditional = True.
2. Use DecoderTrainer to manage training.
3. In the loop, call decoder_trainer(images, unet_number=i) and decoder_trainer.update(unet_number=i).
4. Sample using decoder_trainer.sample(batch_size=N, max_batch_size=M).

import torch
from dalle2_pytorch import Unet, Decoder, DecoderTrainer

unet1 = Unet(dim=128, dim_mults=(1, 2, 4, 8)).cuda()
unet2 = Unet(dim=32, dim_mults=(1, 2, 4, 8, 16)).cuda()

decoder = Decoder(
    unet = (unet1, unet2),
    image_sizes = (256, 512),
    timesteps = 1000,
    unconditional = True
).cuda()

decoder_trainer = DecoderTrainer(decoder)

# Training loop
for i in (1, 2):
    loss = decoder_trainer(images, unet_number = i)
    decoder_trainer.update(unet_number = i)

# Sampling
images = decoder_trainer.sample(batch_size = 36, max_batch_size = 4)