Lumiere PyTorch

The MPLumiere class wraps a pre-existing Karras U-Net (from denoising-diffusion-pytorch) to extend it with temporal capabilities for video generation.

To use it, you must provide a KarrasUnet instance and specify the names of the modules that will be used for temporal inflation (convolutions, attention, upsampling, and downsampling).

import torch
from lumiere_pytorch import MPLumiere
from denoising_diffusion_pytorch import KarrasUnet

# 1. Initialize the base Karras U-Net
karras_unet = KarrasUnet(
    image_size = 256,
    dim = 8,
    channels = 3,
    dim_max = 768,
)

# 2. Wrap it with MPLumiere, specifying module names for temporal layers
lumiere = MPLumiere(
    karras_unet,
    image_size = 256,
    unet_time_kwarg = 'time',
    conv_module_names = [
        'downs.1',
        'ups.1',
        'downs.2',
        'ups.2',
    ],
    attn_module_names = [
        'mids.0'
    ],
    upsample_module_names = [
        'ups.2',
        'ups.1',
    ],
    downsample_module_names = [
        'downs.1',
        'downs.2'
    ]
)

# 3. Prepare input: (batch, channels, time, height, width)
noised_video = torch.randn(2, 3, 8, 256, 256)
time = torch.ones(2,)

# 4. Perform denoising
denoised_video = lumiere(noised_video, time = time)

assert noised_video.shape == denoised_video.shape