phenaki-pytorch

You can improve generation quality by using a TokenCritic. This involves training an extra critic to decide which tokens to iteratively mask during sampling.

Alternatively, you can use MaskGit itself as a self-critic by setting self_token_critic = True when initializing Phenaki.

from phenaki_pytorch import CViViT, MaskGit, TokenCritic, Phenaki

# ... setup cvivit and maskgit ...

# 1. Define the critic
critic = TokenCritic(
    num_tokens = 65536,
    max_seq_len = 1024,
    dim = 512,
    dim_context = 768,
    depth = 6,
    has_cross_attn = True
)

# 2. Pass critic into Phenaki
trainer = Phenaki(
    maskgit = maskgit,
    cvivit = cvivit,
    critic = critic
).cuda()

# OR: Use MaskGit as a self-critic
# phenaki = Phenaki(cvivit=cvivit, maskgit=maskgit, self_token_critic=True)

Install the package using pip:

$ pip install phenaki-pytorch

To train the CViViT component, use the CViViTTrainer. You can train on images first for sample efficiency before fine-tuning on video by setting train_on_images = True.

import torch
from phenaki_pytorch import CViViT, CViViTTrainer

cvivit = CViViT(
    dim = 512,
    codebook_size = 65536,
    image_size = 256,
    patch_size = 32,
    temporal_patch_size = 2,
    spatial_depth = 4,
    temporal_depth = 4,
    dim_head = 64,
    heads = 8
).cuda()

trainer = CViViTTrainer(
    cvivit,
    folder = '/path/to/images/or/videos',
    batch_size = 4,
    grad_accum_every = 4,
    train_on_images = False,  # set to True to train on images first
    use_ema = False,          # recommended to be True to keep exponential moving averaged cvivit
    num_train_steps = 10000
)

trainer.train()               # reconstructions and checkpoints saved to ./results

Use PhenakiTrainer to train the full model. You can provide a custom Dataset that returns (video, caption) tuples, or point to a folder of images/videos for unconditional training.

import torch
from torch.utils.data import Dataset
from phenaki_pytorch import CViViT, MaskGit, Phenaki, PhenakiTrainer

# ... setup phenaki ...

class MyDataset(Dataset):
    def __len__(self): return 100
    def __getitem__(self, idx):
        # Return (video_tensor, caption_string)
        return torch.randn(3, 17, 256, 256), 'a video caption'

dataset = MyDataset()

trainer = PhenakiTrainer(
    phenaki = phenaki,
    batch_size = 4,
    grad_accum_every = 4,
    train_on_images = False, 
    dataset = dataset,       
    sample_texts_file_path = '/path/to/captions.txt' # used for sampling during training
)

trainer.train()

To achieve long-form video generation (beyond the single scene limit), use the make_video function. This function handles the process of conditioning new scenes on the previous $K$ frames to maintain coherence across multiple text prompts.

from phenaki_pytorch import make_video

# entire_video: (1, 3, total_frames, height, width)
# scenes: List[Tensor[3]] - video segment of each scene
entire_video, scenes = make_video(
    phenaki, 
    texts = [
        'a squirrel examines an acorn buried in the snow',
        'a cat watches the squirrel from a frosted window sill',
        'zoom out to show the entire living room'
    ],
    num_frames = (17, 14, 14), # frames per scene
    prime_lengths = (5, 5)     # number of previous frames to use as conditioning
)