imagen-pytorch

The ImagenTrainer wrapper simplifies training by automatically handling Exponential Moving Averages (EMA) for all U-Nets in the cascade. When calling the trainer, you can use max_batch_size to automatically divide a large batch into smaller sub-batches to fit in memory (gradient accumulation). Use trainer.update(unet_number = i) to update the EMA for a specific U-Net. It is highly recommended to use trainer.save() and trainer.load() instead of manual state_dict calls to ensure proper device memory management.

import torch
from imagen_pytorch import Unet, Imagen, ImagenTrainer

unet1 = Unet(
    dim = 32,
    cond_dim = 512,
    dim_mults = (1, 2, 4, 8),
    num_resnet_blocks = 3,
    layer_attns = (False, True, True, True),
)

unet2 = Unet(
    dim = 32,
    cond_dim = 512,
    dim_mults = (1, 2, 4, 8),
    num_resnet_blocks = (2, 4, 8, 8),
    layer_attns = (False, False, False, True),
    layer_cross_attns = (False, False, False, True),
)

imagen = Imagen(
    unets = (unet1, unet2),
    text_encoder_name = 't5-large',
    image_sizes = (64, 256),
    timesteps = 1000,
    cond_drop_prob = 0.1
).cuda()

trainer = ImagenTrainer(imagen)

text_embeds = torch.randn(64, 256, 1024).cuda()
images = torch.randn(64, 3, 256, 256).cuda()

# Training with gradient accumulation via max_batch_size
loss = trainer(
    images,
    text_embeds = text_embeds,
    unet_number = 1,
    max_batch_size = 4
)

trainer.update(unet_number = 1)

# Sampling via trainer
images = trainer.sample(texts = [
    'a puppy looking anxiously at a giant donut on the table',
    'the milky way galaxy in the style of monet'
], cond_scale = 3.)

# Checkpointing
trainer.save('./path/to/checkpoint.pt')
trainer.load('./path/to/checkpoint.pt')
print(trainer.steps) # (2,) step number for each of the unets

If you only want to train the super-resolution stages of the cascade, you can use NullUnet() as a placeholder for the base U-Net. When sampling, you must provide low-resolution images via start_image_or_video and specify start_at_unet_number to skip the base stage.

import torch
from imagen_pytorch import Unet, NullUnet, Imagen

unet1 = NullUnet()

unet2 = Unet(
    dim = 32,
    cond_dim = 512,
    dim_mults = (1, 2, 4, 8),
    num_resnet_blocks = (2, 4, 8, 8),
    layer_attns = (False, False, False, True),
    layer_cross_attns = (False, False, False, True)
)

imagen = Imagen(
    unets = (unet1, unet2),
    image_sizes = (64, 256),
    timesteps = 250,
    cond_drop_prob = 0.1
).cuda()

text_embeds = torch.randn(4, 256, 768).cuda()
images = torch.randn(4, 3, 256, 256).cuda()

# Train only unet 2
loss = imagen(images, text_embeds = text_embeds, unet_number = 2)
loss.backward()

# Sample starting from low-res images
lowres_images = torch.randn(3, 3, 64, 64).cuda()

images = imagen.sample(
    texts = [
        'a whale breaching from afar',
        'young girl blowing out candles on her birthday cake',
        'fireworks with blue and green sparkles'
    ],
    start_at_unet_number = 2,
    start_image_or_video = lowres_images,
    cond_scale = 3.
)

The ImagenTrainer can automatically train using DataLoader instances. To support different training modes, your DataLoader must return data in specific formats:

• Unconditional training: Return images.
• Text-guided generation: Return a tuple of ('images', 'text_embeds').

Note: The DataLoader should return inputs in the order of images, text_embeddings, and then text_masks.

from imagen_pytorch import Unet, Imagen, ImagenTrainer
from imagen_pytorch.data import Dataset

# Unconditional setup
unet = Unet(
    dim = 32,
    dim_mults = (1, 2, 4, 8),
    num_resnet_blocks = 1,
    layer_attns = (False, False, False, True),
    layer_cross_attns = False
)

imagen = Imagen(
    condition_on_text = False,
    unets = unet,
    image_sizes = 128,
    timesteps = 1000
)

trainer = ImagenTrainer(
    imagen = imagen,
    split_valid_from_train = True
).cuda()

# Dataset returns only images for unconditional training
dataset = Dataset('/path/to/training/images', image_size = 128)
trainer.add_train_dataset(dataset, batch_size = 16)

# Training loop
for i in range(200000):
    loss = trainer.train_step(unet_number = 1, max_batch_size = 4)
    print(f'loss: {loss}')

    if not (i % 50):
        valid_loss = trainer.valid_step(unet_number = 1, max_batch_size = 4)
        print(f'valid loss: {valid_loss}')

    if not (i % 100) and trainer.is_main:
        images = trainer.sample(batch_size = 1, return_pil_images = True)
        images[0].save(f'./sample-{i // 100}.png')

To perform unconditional training (generating images without text conditioning), set condition_on_text = False in the Imagen constructor. In this mode, the U-Nets should not expect text conditioning (e.g., set layer_cross_attns = False). You can then train using ImagenTrainer by passing only the images.

import torch
from imagen_pytorch import Unet, Imagen, SRUnet256, ImagenTrainer

unet1 = Unet(
    dim = 32,
    dim_mults = (1, 2, 4),
    num_resnet_blocks = 3,
    layer_attns = (False, True, True),
    layer_cross_attns = False,
    use_linear_attn = True
)

unet2 = SRUnet256(
    dim = 32,
    dim_mults = (1, 2, 4),
    num_resnet_blocks = (2, 4, 8),
    layer_attns = (False, False, True),
    layer_cross_attns = False
)

imagen = Imagen(
    condition_on_text = False,
    unets = (unet1, unet2),
    image_sizes = (64, 128),
    timesteps = 1000
)

trainer = ImagenTrainer(imagen).cuda()

training_images = torch.randn(4, 3, 256, 256).cuda()

# Train unet 1
loss = trainer(training_images, unet_number = 1)
trainer.update(unet_number = 1)

# Sample unconditionally
images = trainer.sample(batch_size = 16)

You can use the Imagen class to manage a cascade of U-Nets. You first define your U-Nets (e.g., using the Unet class), then initialize Imagen with these U-Nets and the target image_sizes. During training, you feed images and either precomputed text_embeds or raw texts into the imagen object, specifying which unet_number in the cascade you are currently training. To generate images, use the .sample() method with text prompts and a cond_scale for classifier-free guidance.

import torch
from imagen_pytorch import Unet, Imagen

# Define U-Nets
unet1 = Unet(
    dim = 32,
    cond_dim = 512,
    dim_mults = (1, 2, 4, 8),
    num_resnet_blocks = 3,
    layer_attns = (False, True, True, True),
    layer_cross_attns = (False, True, True, True)
)

unet2 = Unet(
    dim = 32,
    cond_dim = 512,
    dim_mults = (1, 2, 4, 8),
    num_resnet_blocks = (2, 4, 8, 8),
    layer_attns = (False, False, False, True),
    layer_cross_attns = (False, False, False, True)
)

# Initialize Imagen
imagen = Imagen(
    unets = (unet1, unet2),
    image_sizes = (64, 256),
    timesteps = 1000,
    cond_drop_prob = 0.1
).cuda()

# Mock data
text_embeds = torch.randn(4, 256, 768).cuda()
images = torch.randn(4, 3, 256, 256).cuda()

# Training step for a specific U-Net in the cascade
for i in (1, 2):
    loss = imagen(images, text_embeds = text_embeds, unet_number = i)
    loss.backward()

# Sampling
images = imagen.sample(texts = [
    'a whale breaching from afar',
    'young girl blowing out candles on her birthday cake',
    'fireworks with blue and green sparkles'
], cond_scale = 3.)

print(images.shape) # (3, 3, 256, 256)

You can perform text-guided video synthesis using the Unet3D architecture and the ElucidatedImagen class. ElucidatedImagen manages a cascade of UNets (base and super-resolution) and handles the diffusion process.

Key configuration parameters for ElucidatedImagen include:

• unets: A tuple of Unet3D instances.
• image_sizes: Tuple of target resolutions for each UNet.
• temporal_downsample_factor: Controls temporal downsampling for each UNet in the cascade.
• sigma_min / sigma_max: Noise level bounds.
• S_churn, S_tmin, S_tmax, S_noise: Parameters for stochastic sampling.

If training on text-image pairs instead of video, Unet3D will automatically treat them as single-frame videos by setting ignore_time = True internally.

import torch
from imagen_pytorch import Unet3D, ElucidatedImagen, ImagenTrainer

unet1 = Unet3D(dim = 64, dim_mults = (1, 2, 4, 8)).cuda()
unet2 = Unet3D(dim = 64, dim_mults = (1, 2, 4, 8)).cuda()

# elucidated imagen, which contains the unets above (base unet and super resoluting ones)
imagen = ElucidatedImagen(
    unets = (unet1, unet2),
    image_sizes = (16, 32),
    random_crop_sizes = (None, 16),
    temporal_downsample_factor = (2, 1),
    num_sample_steps = 10,
    cond_drop_prob = 0.1,
    sigma_min = 0.002,
    sigma_max = (80, 160),
    sigma_data = 0.5,
    rho = 7,
    P_mean = -1.2,
    P_std = 1.2,
    S_churn = 80,
    S_tmin = 0.05,
    S_tmax = 50,
    S_noise = 1.003,
).cuda()

# mock videos and text encodings
texts = ['a whale breaching from afar', 'young girl blowing out candles']
videos = torch.randn(4, 3, 10, 32, 32).cuda() # (batch, channels, time, height, width)

# training
trainer = ImagenTrainer(imagen)
trainer(videos, texts = texts, unet_number = 1, ignore_time = False)
trainer.update(unet_number = 1)

# sampling
videos = trainer.sample(texts = texts, video_frames = 20)

ElucidatedImagen is an experimental implementation using the elucidated DDPM for text-guided cascading generation. It requires specific hyperparameters for the sampling schedule and noise levels.

Key Parameters:

• unets: A tuple of UNet instances.
• image_sizes: Tuple of image sizes for the cascade.
• cond_drop_prob: Conditioning dropout probability.
• num_sample_steps: Tuple of sample steps for each UNet.
• sigma_min: Minimum noise level.
• sigma_max: Tuple of maximum noise levels for each UNet.
• sigma_data: Standard deviation of data distribution.
• rho: Controls the sampling schedule.
• P_mean / P_std: Parameters for the log-normal noise distribution.
• S_churn, S_tmin, S_tmax, S_noise: Parameters for stochastic sampling.

from imagen_pytorch import ElucidatedImagen

# instantiate your unets ...

imagen = ElucidatedImagen(
    unets = (unet1, unet2),
    image_sizes = (64, 128),
    cond_drop_prob = 0.1,
    num_sample_steps = (64, 32),
    sigma_min = 0.002,
    sigma_max = (80, 160),
    sigma_data = 0.5,
    rho = 7,
    P_mean = -1.2,
    P_std = 1.2,
    S_churn = 80,
    S_tmin = 0.05,
    S_tmax = 50,
    S_noise = 1.003,
).cuda()

Inpainting is implemented following the Repaint formulation. You can use the .sample() method on either Imagen or ElucidatedImagen by providing inpaint_images and inpaint_masks (for images) or inpaint_videos and inpaint_masks (for video).

Image Inpainting Requirements:

• inpaint_images: Tensor of shape (batch, channels, height, width).
• inpaint_masks: Boolean tensor of shape (batch, height, width).

Video Inpainting Requirements:

• inpaint_videos: Tensor of shape (batch, channels, frames, height, width).
• inpaint_masks: Boolean tensor of shape (batch, height, width) (same mask for all frames) or (batch, frames, height, width) (per-frame masks).

import torch

# Image Inpainting
inpaint_images = torch.randn(4, 3, 512, 512).cuda()
inpaint_masks = torch.ones((4, 512, 512)).bool().cuda()

inpainted_images = trainer.sample(
    texts = ['a whale breaching from afar'], 
    inpaint_images = inpaint_images, 
    inpaint_masks = inpaint_masks, 
    cond_scale = 5.
)

# Video Inpainting
inpaint_videos = torch.randn(4, 3, 8, 512, 512).cuda()
inpaint_masks = torch.ones((4, 8, 512, 512)).bool().cuda()

inpainted_videos = trainer.sample(
    texts = ['a whale breaching from afar'], 
    inpaint_videos = inpaint_videos, 
    inpaint_masks = inpaint_masks, 
    cond_scale = 5.
)