### Install DALL-E 2 PyTorch

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Use pip to install the package from the Python Package Index.

```bash
$ pip install dalle2-pytorch
```

--------------------------------

### Install Open Clip

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Install the open-clip-torch library using pip. This is required for using Open Clip models.

```bash
$ pip install open-clip-torch
```

--------------------------------

### Initialize and Train Decoder Components

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Setup CLIP and multiple U-Net stages for a cascading DDPM decoder, then perform training steps for individual U-Nets.

```python
clip = CLIP(
    dim_text = 512,
    dim_image = 512,
    dim_latent = 512,
    num_text_tokens = 49408,
    text_enc_depth = 6,
    text_seq_len = 256,
    text_heads = 8,
    visual_enc_depth = 6,
    visual_image_size = 256,
    visual_patch_size = 32,
    visual_heads = 8
).cuda()

# 2 unets for the decoder (a la cascading DDPM)

unet1 = Unet(
    dim = 32,
    image_embed_dim = 512,
    cond_dim = 128,
    channels = 3,
    dim_mults = (1, 2, 4, 8)
).cuda()

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

# decoder, which contains the unet(s) and clip

decoder = Decoder(
    clip = clip,
    unet = (unet1, unet2),            # insert both unets in order of low resolution to highest resolution (you can have as many stages as you want here)
    image_sizes = (256, 512),         # resolutions, 256 for first unet, 512 for second. these must be unique and in ascending order (matches with the unets passed in)
    timesteps = 1000,
    image_cond_drop_prob = 0.1,
    text_cond_drop_prob = 0.5
).cuda()

# mock images (get a lot of this)

images = torch.randn(4, 3, 512, 512).cuda()

# feed images into decoder, specifying which unet you want to train
# each unet can be trained separately, which is one of the benefits of the cascading DDPM scheme

loss = decoder(images, unet_number = 1)
loss.backward()

loss = decoder(images, unet_number = 2)
loss.backward()
```

--------------------------------

### Initialize Decoder for Inpainting

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Initialize a Decoder with a single UNet and specified image sizes. This setup is suitable for inpainting tasks.

```python
# 2 unets for the decoder (a la cascading DDPM)

unet = Unet(
    dim = 16,
    image_embed_dim = 512,
    cond_dim = 128,
    channels = 3,
    dim_mults = (1, 1, 1, 1)
).cuda()


# decoder, which contains the unet(s) and clip

decoder = Decoder(
    clip = clip,
    unet = (unet,),
    image_sizes = (256,),
    timesteps = 1000,
    image_cond_drop_prob = 0.1,
    text_cond_drop_prob = 0.5
).cuda()
```

--------------------------------

### Initialize CLIP Model for DecoderTrainer

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Initializes the CLIP model with specified dimensions and token counts for use with DecoderTrainer. Note the different depth values compared to the first example.

```python
clip = CLIP(
    dim_text = 512,
    dim_image = 512,
    dim_latent = 512,
    num_text_tokens = 49408,
    text_enc_depth = 6,
    text_seq_len = 256,
    text_heads = 8,
    visual_enc_depth = 6,
    visual_image_size = 256,
    visual_patch_size = 32,
    visual_heads = 8
).cuda()
```

--------------------------------

### Set up Decoder for Latent Diffusion

Source: https://context7.com/lucidrains/dalle2-pytorch/llms.txt

Combine VQ-GAN VAE with Unet models for latent diffusion to generate high-resolution images. This setup involves multiple VAEs and Unets for different resolution stages. Training can be memory-efficient by loading one Unet at a time.

```python
import torch
from dalle2_pytorch import Decoder, Unet, VQGanVAE, OpenAIClipAdapter

clip = OpenAIClipAdapter('ViT-B/32')

# Train VAEs for each latent diffusion stage (pre-trained)
vae1 = VQGanVAE(dim=32, image_size=256, layers=3, layer_mults=(1, 2, 4)).cuda()
vae2 = VQGanVAE(dim=32, image_size=512, layers=3, layer_mults=(1, 2, 4)).cuda()

# U-Nets for latent space and pixel space
unet1 = Unet(dim=32, image_embed_dim=512, cond_dim=128, channels=3,
             sparse_attn=True, sparse_attn_window=2, dim_mults=(1, 2, 4, 8)).cuda()
unet2 = Unet(dim=32, image_embed_dim=512, channels=3, dim_mults=(1, 2, 4, 8, 16),
             cond_on_image_embeds=True, cond_on_text_encodings=False).cuda()
unet3 = Unet(dim=32, image_embed_dim=512, channels=3, dim_mults=(1, 2, 4, 8, 16),
             cond_on_image_embeds=True, attend_at_middle=False).cuda()

# Decoder with latent diffusion for first two stages
decoder = Decoder(
    clip = clip,
    vae = (vae1, vae2),              # Latent diffusion for unet1 and unet2
    unet = (unet1, unet2, unet3),
    image_sizes = (256, 512, 1024),
    timesteps = 100,
    image_cond_drop_prob = 0.1,
    text_cond_drop_prob = 0.5
).cuda()

# Training each U-Net with memory-efficient context
images = torch.randn(1, 3, 1024, 1024).cuda()

with decoder.one_unet_in_gpu(1):
    loss = decoder(images, unet_number = 1)
    loss.backward()

with decoder.one_unet_in_gpu(2):
    loss = decoder(images, unet_number = 2)
    loss.backward()

with decoder.one_unet_in_gpu(3):
    loss = decoder(images, unet_number = 3)
    loss.backward()

# Sampling high-resolution images
image_embed = torch.randn(1, 512).cuda()
with torch.no_grad():
    high_res_images = decoder.sample(image_embed)
    print(f"Output shape: {high_res_images.shape}")  # (1, 3, 1024, 1024)
```

--------------------------------

### High Resolution Image Synthesis with Cascaded Diffusion

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Example demonstrating the integration of cascaded diffusion techniques for high-resolution image synthesis within the DALL-E 2 framework. Requires CUDA.

```python
import torch
from dalle2_pytorch import Unet, Decoder, CLIP

```

--------------------------------

### Launch Training with Accelerate

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/prior.md

Use this command to launch the training script with distributed training capabilities via huggingface accelerate. Ensure you have a configuration file specified.

```bash
accelerate launch train_diffusion_prior.py --config_path <path to your config>
```

--------------------------------

### Launch Training with Python

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/prior.md

Use this command to launch the training script directly with Python, suitable for single GPU or CPU training without huggingface accelerate.

```bash
python train_diffusion_prior.py
```

--------------------------------

### Image Embedding Dataset Creation

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Example of importing ImageEmbeddingDataset and create_image_embedding_dataloader from dalle2_pytorch.dataloaders. These are used for training the decoder with image embeddings.

```python
from dalle2_pytorch.dataloaders import ImageEmbeddingDataset, create_image_embedding_dataloader
```

--------------------------------

### Initialize VQGanVAE

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Import and initialize the VQGanVAE model. This is part of the experimental section for DALL-E 2 with latent diffusion.

```python
import torch
from dalle2_pytorch import Unet, Decoder, CLIP, VQGanVAE
```

--------------------------------

### Train Unconditional DDPM Decoder

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Train an unconditional DDPM model by setting `unconditional = True` in the Decoder. This example shows setting up two UNets for cascading DDPMs and training the decoder.

```python
import torch
from dalle2_pytorch import Unet, Decoder, DecoderTrainer

# unet for the cascading ddpm

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

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

# decoder, which contains the unets

decoder = Decoder(
    unet = (unet1, unet2),
    image_sizes = (256, 512),  # first unet up to 256px, then second to 512px
    timesteps = 1000,
    unconditional = True
).cuda()

# decoder trainer

decoder_trainer = DecoderTrainer(decoder)

# images (get a lot of this)

images = torch.randn(1, 3, 512, 512).cuda()

# feed images into decoder

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

# do the above for many many many many images
# then it will learn to generate images

images = decoder_trainer.sample(batch_size = 36, max_batch_size = 4) # (36, 3, 512, 512)
```

--------------------------------

### Configure CLIP-less Diffusion Prior

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Shows how to initialize the DiffusionPrior without a CLIP model by providing the image_embed_dim directly.

```python
import torch
from dalle2_pytorch import DiffusionPriorNetwork, DiffusionPrior

# setup prior network, which contains an autoregressive transformer

prior_network = DiffusionPriorNetwork(
    dim = 512,
    depth = 6,
    dim_head = 64,
    heads = 8
).cuda()

# diffusion prior network, which contains the CLIP and network (with transformer) above

diffusion_prior = DiffusionPrior(
    net = prior_network,
    image_embed_dim = 512,               # this needs to be set
    timesteps = 100,
    cond_drop_prob = 0.2,
    condition_on_text_encodings = False  # this probably should be true, but just to get Laion started
).cuda()

# mock data

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

# precompute the text and image embeddings
# here using the diffusion prior class, but could be done with CLIP alone

clip_image_embeds = torch.randn(4, 512).cuda()
clip_text_embeds = torch.randn(4, 512).cuda()

# feed text and images into diffusion prior network

loss = diffusion_prior(
    text_embed = clip_text_embeds,
    image_embed = clip_image_embeds
)

loss.backward()

# do the above for many many many steps
# now the diffusion prior can generate image embeddings from the text embeddings
```

--------------------------------

### Initialize VQGanVAE Models

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Initializes two VQGanVAE models with different image sizes. These are used for latent diffusion in the decoder.

```python
vae1 = VQGanVAE(
    dim = 32,
    image_size = 256,
    layers = 3,
    layer_mults = (1, 2, 4)
)

vae2 = VQGanVAE(
    dim = 32,
    image_size = 512,
    layers = 3,
    layer_mults = (1, 2, 4)
)
```

--------------------------------

### Initialize and Train a Custom CLIP Model

Source: https://context7.com/lucidrains/dalle2-pytorch/llms.txt

Configure a CLIP model with advanced training objectives like FILIP and DCL, then perform a training step and extract embeddings.

```python
import torch
from dalle2_pytorch import CLIP

# Initialize CLIP model with custom configuration
clip = CLIP(
    dim_text = 512,
    dim_image = 512,
    dim_latent = 512,
    num_text_tokens = 49408,
    text_enc_depth = 6,
    text_seq_len = 256,
    text_heads = 8,
    visual_enc_depth = 6,
    visual_image_size = 256,
    visual_patch_size = 32,
    visual_heads = 8,
    use_all_token_embeds = True,            # Enable fine-grained contrastive learning (FILIP)
    decoupled_contrastive_learning = True,  # Use DCL objective function
    extra_latent_projection = True,         # Separate projections for comparisons
    use_visual_ssl = True,                  # Self-supervised learning on images
    visual_ssl_type = 'simclr',             # 'simclr' or 'simsiam'
    use_mlm = False,                        # Masked language modeling
    text_ssl_loss_weight = 0.05,
    image_ssl_loss_weight = 0.05
).cuda()

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

# Training step - compute contrastive loss
loss = clip(
    text,
    images,
    return_loss = True
)
loss.backward()

# After training, embed text and images separately
with torch.no_grad():
    text_embeds = clip.embed_text(text)      # Returns text embeddings
    image_embeds = clip.embed_image(images)  # Returns image embeddings
```

--------------------------------

### Create Image Embedding Dataloader and Dataset

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/dalle2_pytorch/dataloaders/README.md

Demonstrates how to initialize a dataloader for image and embedding pairs using webdataset, or configure the dataset manually for custom loading.

```python
from dalle2_pytorch.dataloaders import ImageEmbeddingDataset, create_image_embedding_dataloader

# Create a dataloader directly.
dataloader = create_image_embedding_dataloader(
    tar_url="/path/or/url/to/webdataset/{0000..9999}.tar", # Uses bracket expanding notation. This specifies to read all tars from 0000.tar to 9999.tar
    embeddings_url="path/or/url/to/embeddings/folder",     # Included if .npy files are not in webdataset. Left out or set to None otherwise
    num_workers=4,
    batch_size=32,
    shard_width=4,                                         # If a file in the webdataset shard 3 is named 0003039.jpg, we know the shard width is 4 and the last three digits are the index
    shuffle_num=200,                                       # Does a shuffle of the data with a buffer size of 200
    shuffle_shards=True,                                   # Shuffle the order the shards are read in
    resample_shards=False,                                 # Sample shards with replacement. If true, an epoch will be infinite unless stopped manually
)
for img, emb in dataloader:
    print(img.shape)  # torch.Size([32, 3, 256, 256])
    print(emb.shape)  # torch.Size([32, 512])
    # Train decoder only as shown above

# Or create a dataset without a loader so you can configure it manually
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
)
```

--------------------------------

### Initialize and Train DiffusionPriorTrainer

Source: https://context7.com/lucidrains/dalle2-pytorch/llms.txt

Initializes a DiffusionPriorTrainer with EMA and scheduling, then demonstrates a training loop and sampling from the EMA model. Load checkpoints using the .load() method.

```python
trainer = DiffusionPriorTrainer(
    diffusion_prior = diffusion_prior,
    lr = 1.1e-4,
    wd = 6.02e-2,
    eps = 1e-6,
    max_grad_norm = 0.5,
    use_ema = True,
    ema_beta = 0.99,
    ema_update_after_step = 1000,
    ema_update_every = 10,
    warmup_steps = 1000,
    cosine_decay_max_steps = 100000
)

# Training loop
text = torch.randint(0, 49408, (32, 256)).cuda()
images = torch.randn(32, 3, 256, 256).cuda()

for step in range(10000):
    # Forward pass with gradient accumulation
    loss = trainer(text, images, max_batch_size = 8)

    # Update weights and EMA
    trainer.update()

    if step % 100 == 0:
        print(f"Step {step}: loss = {loss:.4f}")

    # Save checkpoint periodically
    if step % 1000 == 0:
        trainer.save(f"./checkpoints/prior_step_{step}.pt")

# Sampling from EMA model
with torch.no_grad():
    image_embeds = trainer.sample(text[:4], max_batch_size = 4)

# Load checkpoint
trainer.load("./checkpoints/prior_step_9000.pt")
```

--------------------------------

### Initialize and Train Unconditional Decoder

Source: https://context7.com/lucidrains/dalle2-pytorch/llms.txt

Configures a multi-UNet decoder for unconditional image generation and demonstrates the training loop.

```python
decoder = Decoder(
    unet = (unet1, unet2),
    image_sizes = (256, 512),
    timesteps = 1000,
    unconditional = True  # No CLIP conditioning
).cuda()

# Create trainer
trainer = DecoderTrainer(decoder)

# Training loop
images = torch.randn(4, 3, 512, 512).cuda()

for step in range(1000):
    loss1 = trainer(images, unet_number = 1)
    trainer.update(unet_number = 1)

    loss2 = trainer(images, unet_number = 2)
    trainer.update(unet_number = 2)

# Sample unconditionally
with torch.no_grad():
    generated = trainer.sample(
        batch_size = 16,
        max_batch_size = 4
    )
    print(f"Generated shape: {generated.shape}")  # (16, 3, 512, 512)
```

--------------------------------

### Tokenize Text for Prior Sampling

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/prior.md

Shows how to tokenize a user-provided prompt using CLIP's tokenizer, preparing it as input for the diffusion prior's sampling function.

```python
# tokenize the text
tokenized_text = clip.tokenize("<your amazing prompt>")
```

--------------------------------

### Configure Prior Embedding Dataset for Distributed Training

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/dalle2_pytorch/dataloaders/README.md

Shows how to initialize an embedding reader and generate training, evaluation, and test splits for distributed training environments.

```python
from dalle2_pytorch.dataloaders import get_reader, make_splits

# grab embeddings from some specified location
IMG_URL = "data/img_emb/"
META_URL = "data/meta/"

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

# some config for training
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,
}

# specifying a rank will handle allocation internally
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)
```

--------------------------------

### Initialize Decoder with Cascading U-Nets

Source: https://context7.com/lucidrains/dalle2-pytorch/llms.txt

Sets up a cascading Decoder using two U-Net models for different resolutions (128x128 and 256x256). It utilizes CLIP for image embeddings and supports learned variance and classifier-free guidance.

```python
import torch
from dalle2_pytorch import Decoder, Unet, OpenAIClipAdapter

clip = OpenAIClipAdapter('ViT-B/32')

# Setup cascading U-Nets
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 = 64,
    image_embed_dim = 512,
    cond_dim = 128,
    channels = 3,
    dim_mults = (1, 2, 4, 8, 16)
).cuda()

# Create cascading decoder
decoder = Decoder(
    unet = (unet1, unet2),
    image_sizes = (128, 256),        # Resolutions for each U-Net
    clip = clip,
    timesteps = 1000,
    sample_timesteps = (250, 27),    # DDIM steps per U-Net
    image_cond_drop_prob = 0.1,
    text_cond_drop_prob = 0.5,
    learned_variance = True,         # Learn variance (improved DDPM)
    use_blur_for_lowres_cond = True, # Blur conditioning for cascade
    beta_schedule = ('cosine', 'linear')
).cuda()
```

--------------------------------

### Initialize DiffusionPriorTrainer

Source: https://context7.com/lucidrains/dalle2-pytorch/llms.txt

Sets up the training wrapper for the diffusion prior model. This utility supports advanced training features like EMA and distributed training.

```python
import torch
from dalle2_pytorch import DiffusionPrior, DiffusionPriorNetwork, OpenAIClipAdapter
from dalle2_pytorch.trainer import DiffusionPriorTrainer

# Setup model
prior_network = DiffusionPriorNetwork(
    dim = 768,
    depth = 24,
    dim_head = 64,
    heads = 32
).cuda()

diffusion_prior = DiffusionPrior(
    net = prior_network,
    clip = OpenAIClipAdapter('ViT-L/14'),
    timesteps = 1000,
    cond_drop_prob = 0.1,
    condition_on_text_encodings = True
).cuda()
```

--------------------------------

### Initialize DALLE2 Model

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Combine a diffusion prior and a decoder to create the main DALLE2 model.

```python
dalle2 = DALLE2(
    prior = diffusion_prior,
    decoder = decoder
)
```

--------------------------------

### Initialize Super-resolution Unet for 512x512 Upsampling

Source: https://context7.com/lucidrains/dalle2-pytorch/llms.txt

Initializes a super-resolution U-Net for upsampling images to 512x512. It supports low-resolution conditioning and image embedding conditioning but not text conditioning.

```python
# Super-resolution U-Net for 512x512 upsampling
unet_sr = Unet(
    dim = 64,
    image_embed_dim = 512,
    cond_dim = 128,
    channels = 3,
    dim_mults = (1, 2, 4, 8, 16),
    num_resnet_blocks = 2,
    lowres_cond = True,             # Condition on low-res output
    cond_on_image_embeds = True,
    cond_on_text_encodings = False, # SR doesn't need text
    attend_at_middle = False        # Speed up high-res training
).cuda()
```

--------------------------------

### Initialize Open Clip Adapter

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Initialize the Open Clip adapter with a specified model variant, such as 'ViT-H/14'. This adapter is used for CLIP integration.

```python
from dalle2_pytorch import OpenClipAdapter

clip = OpenClipAdapter('ViT-H/14')
```

--------------------------------

### Initialize Unet Models for Decoder

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Initializes three Unet models with varying configurations for different stages of the decoder. Note the conditional inputs and attention mechanisms.

```python
unet1 = Unet(
    dim = 32,
    image_embed_dim = 512,
    cond_dim = 128,
    channels = 3,
    sparse_attn = True,
    sparse_attn_window = 2,
    dim_mults = (1, 2, 4, 8)
)

unet2 = Unet(
    dim = 32,
    image_embed_dim = 512,
    channels = 3,
    dim_mults = (1, 2, 4, 8, 16),
    cond_on_image_embeds = True,
    cond_on_text_encodings = False
)

unet3 = Unet(
    dim = 32,
    image_embed_dim = 512,
    channels = 3,
    dim_mults = (1, 2, 4, 8, 16),
    cond_on_image_embeds = True,
    cond_on_text_encodings = False,
    attend_at_middle = False
)
```

--------------------------------

### Initialize DALLE2 Pipeline

Source: https://context7.com/lucidrains/dalle2-pytorch/llms.txt

Combines a DiffusionPrior and a Decoder into a single pipeline for end-to-end text-to-image generation.

```python
import torch
from dalle2_pytorch import DALLE2, DiffusionPrior, DiffusionPriorNetwork, Decoder, Unet, OpenAIClipAdapter

# Setup and train prior (abbreviated)
clip = OpenAIClipAdapter('ViT-B/32')

prior_network = DiffusionPriorNetwork(dim=512, depth=6, dim_head=64, heads=8).cuda()
diffusion_prior = DiffusionPrior(
    net=prior_network, clip=clip, timesteps=1000, cond_drop_prob=0.2
).cuda()

# Setup and train decoder (abbreviated)
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=64, 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()

# Combine into DALLE2 pipeline
dalle2 = DALLE2(
    prior = diffusion_prior,
    decoder = decoder,
    prior_num_samples = 2  # Generate 2 embeddings, pick best
)
```

--------------------------------

### Initialize and Train VQGanVAE

Source: https://context7.com/lucidrains/dalle2-pytorch/llms.txt

Create and train a VQ-GAN VAE for latent diffusion. This class supports various configurations for reconstruction loss, GAN loss, and codebook initialization. Ensure CUDA is available if using .cuda().

```python
import torch
from dalle2_pytorch import VQGanVAE
from dalle2_pytorch.vqgan_vae_trainer import VQGanVAETrainer

# Create VQ-GAN VAE
vae = VQGanVAE(
    dim = 64,
    image_size = 256,
    channels = 3,
    layers = 4,
    l2_recon_loss = False,           # Use L1 loss
    use_hinge_loss = True,           # Hinge GAN loss
    vq_codebook_dim = 256,
    vq_codebook_size = 8192,         # Number of codebook entries
    vq_decay = 0.8,
    vq_commitment_weight = 1.0,
    vq_kmeans_init = True,           # K-means codebook initialization
    vq_use_cosine_sim = True,        # Cosine similarity for quantization
    use_vgg_and_gan = True,          # Enable perceptual + GAN loss
    vae_type = 'resnet',             # 'resnet' or 'vit'
    discr_layers = 4
).cuda()

# Training the VAE
images = torch.randn(8, 3, 256, 256).cuda()

# Forward pass returns reconstruction
recon = vae(images)
print(f"Reconstruction shape: {recon.shape}")  # (8, 3, 256, 256)

# Training with autoencoder loss
ae_loss = vae(images, return_loss = True)
ae_loss.backward()

# Training discriminator
discr_loss = vae(images, return_discr_loss = True)
discr_loss.backward()

# Encode/decode separately
encoded = vae.encode(images)  # Latent representation
print(f"Encoded shape: {encoded.shape}")  # (8, 256, 16, 16) for 4 layers

decoded = vae.decode(encoded)  # Reconstruct from latents
print(f"Decoded shape: {decoded.shape}")  # (8, 3, 256, 256)
```

--------------------------------

### Generate Image from Text Embedding (Naive Approach)

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/prior.md

Demonstrates an initial attempt to generate an image from a text embedding using CLIP and a decoder. This approach highlights the issue of embeddings not being in the same space.

```python
# Load Models
clip_model = clip.load("ViT-L/14")
decoder = Decoder(checkpoint="best.pth") # A decoder trained on CLIP Image embeddings

# Retrieve prompt from user and encode with CLIP
prompt = "A corgi wearing sunglasses"
tokenized_text = tokenize(prompt)
text_embedding = clip_model.encode_text(tokenized_text)

# Now, pass the text embedding to the decoder
predicted_image = decoder.sample(text_embedding)
```

--------------------------------

### Integrate Open CLIP Models

Source: https://context7.com/lucidrains/dalle2-pytorch/llms.txt

Load SOTA Open CLIP models using the OpenClipAdapter to generate text and image embeddings.

```python
from dalle2_pytorch import OpenClipAdapter

# Load SOTA Open CLIP model (ViT-H/14 trained on LAION)
clip = OpenClipAdapter('ViT-H/14')

# Alternative configurations
# clip = OpenClipAdapter('ViT-B/32', pretrained='laion400m_e32')
# clip = OpenClipAdapter('ViT-L/14', pretrained='laion2b_s32b_b82k')

# Embed text
text_tokens = torch.randint(0, 49408, (2, 77)).cuda()
text_embed, text_encodings = clip.embed_text(text_tokens)
print(f"Text embedding shape: {text_embed.shape}")  # (2, 768) for ViT-H

# Embed images
images = torch.randn(2, 3, 224, 224).cuda()
image_embed, image_encodings = clip.embed_image(images)
print(f"Image embedding shape: {image_embed.shape}")  # (2, 1024) for ViT-H
```

--------------------------------

### Initialize Base Unet for 256x256 Generation

Source: https://context7.com/lucidrains/dalle2-pytorch/llms.txt

Initializes a base U-Net model for 256x256 image generation, supporting various conditioning methods like text and image embeddings, self-attention, and cross-attention.

```python
import torch
from dalle2_pytorch import Unet

# Base U-Net for 256x256 generation
unet_base = Unet(
    dim = 128,
    image_embed_dim = 512,          # CLIP embedding dimension
    text_embed_dim = 512,           # Text encoding dimension
    cond_dim = 128,                 # Conditioning dimension
    channels = 3,
    dim_mults = (1, 2, 4, 8),       # Channel multipliers per resolution
    num_resnet_blocks = 2,
    cond_on_text_encodings = True,  # Enable text conditioning
    cond_on_image_embeds = True,    # Enable image embedding conditioning
    attend_at_middle = True,        # Attention at bottleneck
    attn_heads = 16,
    attn_dim_head = 32,
    init_cross_embed = True,        # Cross-scale embedding
    memory_efficient = False,
    self_attn = (False, False, False, True),  # Self-attention per resolution
    sparse_attn = True,             # Efficient linear attention
    cosine_sim_cross_attn = True
).cuda()
```

--------------------------------

### Initialize Decoder Model for DecoderTrainer

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Initializes the Decoder model with a tuple of Unets and image sizes for use with DecoderTrainer. The CLIP model is also passed.

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

--------------------------------

### Initialize DecoderTrainer

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Initializes the DecoderTrainer with the Decoder model and training parameters such as learning rate, weight decay, and EMA settings. Gradient accumulation is configured via max_batch_size.

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

--------------------------------

### Full DALL-E 2 Training Script

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

A complete workflow including CLIP training, Diffusion Prior training, and Decoder training.

```python
import torch
from dalle2_pytorch import DALLE2, DiffusionPriorNetwork, DiffusionPrior, Unet, Decoder, CLIP

clip = CLIP(
    dim_text = 512,
    dim_image = 512,
    dim_latent = 512,
    num_text_tokens = 49408,
    text_enc_depth = 6,
    text_seq_len = 256,
    text_heads = 8,
    visual_enc_depth = 6,
    visual_image_size = 256,
    visual_patch_size = 32,
    visual_heads = 8
).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()

# do above for many steps ...

# prior networks (with transformer)

prior_network = DiffusionPriorNetwork(
    dim = 512,
    depth = 6,
    dim_head = 64,
    heads = 8
).cuda()

diffusion_prior = DiffusionPrior(
    net = prior_network,
    clip = clip,
    timesteps = 1000,
    sample_timesteps = 64,
    cond_drop_prob = 0.2
).cuda()

loss = diffusion_prior(text, images)
loss.backward()

# do above for many steps ...

# decoder (with unet)

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    # set to True for any unets that need to be conditioned on text encodings
).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 = 100,
    image_cond_drop_prob = 0.1,
    text_cond_drop_prob = 0.5
).cuda()

for unet_number in (1, 2):
    loss = decoder(images, text = text, unet_number = unet_number) # this can optionally be decoder(images, text) if you wish to condition on the text encodings as well, though it was hinted in the paper it didn't do much
    loss.backward()

# do above for many steps

dalle2 = DALLE2(
    prior = diffusion_prior,
    decoder = decoder
)

images = dalle2(
    ['cute puppy chasing after a squirrel'],
    cond_scale = 2. # classifier free guidance strength (> 1 would strengthen the condition)
)

# save your image (in this example, of size 256x256)
```

--------------------------------

### Sample Images using DecoderTrainer

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Samples images using the exponentially moving averaged Unets from the DecoderTrainer. Requires image embeddings and optionally text embeddings.

```python
# after much training
# you can sample from the exponentially moving averaged unets as so

mock_image_embed = torch.randn(32, 512).cuda()
images = decoder_trainer.sample(image_embed = mock_image_embed, text = text) # (4, 3, 256, 256)
```

--------------------------------

### Initialize CLIP Model

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Initialize a CLIP model with specified dimensions and parameters for text and visual encoders. This model is used for conditioning.

```python
from dalle2_pytorch import Unet, Decoder, CLIP

clip = CLIP(
    dim_text = 512,
    dim_image = 512,
    dim_latent = 512,
    num_text_tokens = 49408,
    text_enc_depth = 6,
    text_seq_len = 256,
    text_heads = 8,
    visual_enc_depth = 6,
    visual_image_size = 256,
    visual_patch_size = 32,
    visual_heads = 8
).cuda()
```

--------------------------------

### Initialize Decoder Model

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Initializes the Decoder model, combining CLIP, VAEs, and Unets. Specify image sizes and timesteps for the diffusion process.

```python
decoder = Decoder(
    clip = clip,
    vae = (vae1, vae2),                # latent diffusion for unet1 (vae1) and unet2 (vae2), but not for the last unet3
    unet = (unet1, unet2, unet3),      # insert unets in order of low resolution to highest resolution (you can have as many stages as you want here)
    image_sizes = (256, 512, 1024),    # resolutions, 256 for first unet, 512 for second, 1024 for third
    timesteps = 100,
    image_cond_drop_prob = 0.1,
    text_cond_drop_prob = 0.5
).cuda()
```

--------------------------------

### Initialize DALL-E 2 Decoder with UNets

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Initialize two UNet models and a Decoder for DALL-E 2. The decoder takes UNets, image sizes, CLIP model, and timestep information.

```python
unet1 = Unet(
    dim = 128,
    image_embed_dim = 512,
    cond_dim = 128,
    channels = 3,
    dim_mults=(1, 2, 4, 8),
    text_embed_dim = 512,
    cond_on_text_encodings = True  # set to True for any unets that need to be conditioned on text encodings (ex. first unet in cascade)
).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,
    sample_timesteps = (250, 27),
    image_cond_drop_prob = 0.1,
    text_cond_drop_prob = 0.5
).cuda()
```

--------------------------------

### Initialize Unet Models for DecoderTrainer

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Initializes two Unet models with different configurations for use with DecoderTrainer. One Unet is conditioned on text encodings.

```python
# decoder (with unet)

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

--------------------------------

### Train CLIP Model

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Initialize and train the CLIP model for text and image embeddings. Ensure CUDA is available for GPU acceleration. This step is crucial for subsequent training stages.

```python
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,            # whether to use fine-grained contrastive learning (FILIP)
    decoupled_contrastive_learning = True,  # use decoupled contrastive learning (DCL) objective function, removing positive pairs from the denominator of the InfoNCE loss (CLOOB + DCL)
    extra_latent_projection = True,         # whether to use separate projections for text-to-image vs image-to-text comparisons (CLOOB)
    use_visual_ssl = True,                  # whether to do self supervised learning on images
    visual_ssl_type = 'simclr',             # can be either 'simclr' or 'simsiam', depending on using DeCLIP or SLIP
    use_mlm = False,                        # use masked language learning (MLM) on text (DeCLIP)
    text_ssl_loss_weight = 0.05,            # weight for text MLM loss
    image_ssl_loss_weight = 0.05            # weight for image self-supervised learning loss
).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              # needs to be set to True to return contrastive loss
)

loss.backward()

# do the above with as many texts and images as possible in a loop

```

--------------------------------

### Initialize CLIP Model

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Initializes the CLIP model with specified dimensions and token counts. Ensure dimensions match downstream models.

```python
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
)
```

--------------------------------

### Create Image Embedding Dataset Manually

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/README.md

Instantiate the ImageEmbeddingDataset directly to configure it manually. Specify the webdataset URLs, embedding folder URL, shard width, and shuffling options.

```python
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
)
```

--------------------------------

### Generate Image from Text Embedding (Diffusion Prior Approach)

Source: https://github.com/lucidrains/dalle2-pytorch/blob/main/prior.md

Illustrates the improved method of generating an image by first using a diffusion prior to translate text embeddings into CLIP's image embedding space, then feeding this to the decoder. This ensures embeddings are compatible.

```python
# Load Models
prior= Prior(checkpoint="prior.pth") # A decoder trained to go from: text-> clip text emb -> clip img emb
decoder = Decoder(checkpoint="decoder.pth") # A decoder trained on CLIP Image embeddings

# Retrieve prompt from user and encode with a prior
prompt = "A corgi wearing sunglasses"
tokenized_text = tokenize(prompt)
text_embedding = prior.sample(tokenized_text) # <-- now we get an embedding in the same space as images!

# Now, pass the predicted image embedding to the decoder
predicted_image = decoder.sample(text_embedding)
```

--------------------------------

### Integrate OpenAI CLIP with DiffusionPrior

Source: https://context7.com/lucidrains/dalle2-pytorch/llms.txt

Use the OpenAIClipAdapter to wrap pretrained models for use within the DALL-E 2 diffusion prior pipeline.

```python
import torch
from dalle2_pytorch import OpenAIClipAdapter, DiffusionPrior, DiffusionPriorNetwork

# Load pretrained OpenAI CLIP (defaults to ViT-B/32)
clip = OpenAIClipAdapter('ViT-B/32')

# Alternative: Load larger model
# clip = OpenAIClipAdapter('ViT-L/14')

# Use with DiffusionPrior for text-to-image embedding translation
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()

# The adapter handles text and image embedding automatically
text = torch.randint(0, 49408, (4, 256)).cuda()
images = torch.randn(4, 3, 256, 256).cuda()

loss = diffusion_prior(text, images)
loss.backward()
```

--------------------------------

### Train Decoder U-Nets

Source: https://context7.com/lucidrains/dalle2-pytorch/llms.txt

Demonstrates training individual U-Nets within the Decoder by feeding images and optional text, then performing backpropagation for each U-Net.

```python
# Training: feed images, optionally with text
images = torch.randn(4, 3, 256, 256).cuda()
text = torch.randint(0, 49408, (4, 256)).cuda()

# Train each U-Net separately
loss1 = decoder(images, text = text, unet_number = 1)
loss1.backward()

loss2 = decoder(images, text = text, unet_number = 2)
loss2.backward()
```