magvit2-pytorch - MagViT2视频生成和理解模型的PyTorch开源实现

MagViT2 - Pytorch

在Pytorch中实现《语言模型胜过扩散模型 - 分词器是视觉生成的关键》中的MagViT2。目前该模型在视频生成/理解方面保持着最先进的水平。

论文中提出的无查找量化器可以在单独的仓库中找到。它可能应该被用于探索所有其他模态，从音频开始。

如果你对在公开场合复现本文提出的分词器感兴趣，请加入

致谢

感谢StabilityAI和🤗 Huggingface的慷慨赞助，以及我的其他赞助商，让我能够独立地开源人工智能。
感谢Louis Serrano分享了一些早期的初步运行结果，验证了整体架构在有限标量量化下的收敛性。
你？如果你是一位才华横溢的研究工程师/科学家，欢迎为前沿开源科学做出贡献！

安装

$ pip install magvit2-pytorch

使用方法

from magvit2_pytorch import (
    VideoTokenizer,
    VideoTokenizerTrainer
)

tokenizer = VideoTokenizer(
    image_size = 128,
    init_dim = 64,
    max_dim = 512,
    codebook_size = 1024,
    layers = (
        'residual',
        'compress_space',
        ('consecutive_residual', 2),
        'compress_space',
        ('consecutive_residual', 2),
        'linear_attend_space',
        'compress_space',
        ('consecutive_residual', 2),
        'attend_space',
        'compress_time',
        ('consecutive_residual', 2),
        'compress_time',
        ('consecutive_residual', 2),
        'attend_time',
    )
)

trainer = VideoTokenizerTrainer(
    tokenizer,
    dataset_folder = '/path/to/a/lot/of/media',     # 视频或图像文件夹，取决于下面的设置
    dataset_type = 'videos',                        # 'videos' 或 'images'，先前的论文表明在图像上预训练对视频合成是有效的
    batch_size = 4,
    grad_accum_every = 8,
    learning_rate = 2e-5,
    num_train_steps = 1_000_000
)

trainer.train()

# 经过大量训练后...
# 可以使用分词器的指数移动平均（EMA）版本

ema_tokenizer = trainer.ema_tokenizer

# 模拟视频

video = torch.randn(1, 3, 17, 128, 128)

# 将视频标记为离散代码

codes = ema_tokenizer.tokenize(video) # (1, 9, 16, 16) <- 在这个例子中，时间维度下采样4倍，空间维度下采样8倍。展平token ID以进行（非）自回归训练

# 完整性检查

decoded_video = ema_tokenizer.decode_from_code_indices(codes)

assert torch.allclose(
    decoded_video,
    ema_tokenizer(video, return_recon = True)
)

要在Weights & Biases上跟踪你的实验，在VideoTokenizerTrainer上设置use_wandb_tracking = True，然后使用.trackers上下文管理器


trainer = VideoTokenizerTrainer(
    use_wandb_tracking = True,
    ...
)
使用trainer.trackers(project_name = 'magvit2', run_name = 'baseline')：
    trainer.train()

待办事项

引用

[省略引用部分的翻译]