magvit2-pytorch

MagViT2 - Pytorch

在Pytorch中实现<a href="https://arxiv.org/abs/2310.05737">《语言模型胜过扩散模型 - 分词器是视觉生成的关键》</a>中的MagViT2。目前该模型在视频生成/理解方面保持着最先进的水平。

论文中提出的无查找量化器可以在<a href="https://github.com/lucidrains/vector-quantize-pytorch/blob/master/vector_quantize_pytorch/lookup_free_quantization.py">单独的仓库</a>中找到。它可能应该被用于探索所有其他模态，从<a href="https://github.com/lucidrains/audiolm-pytorch/commit/c748fcdb565964bc562277bd73fbeb2e5df0ffca">音频</a>开始。

如果你对在公开场合复现本文提出的分词器感兴趣，请加入<a href="https://discord.gg/xBPBXfcFHd"><img alt="加入我们的Discord" src="https://img.shields.io/discord/823813159592001537?color=5865F2&logo=discord&logoColor=white"></a>

致谢

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

安装

$ 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)
)

要在<a href="https://wandb.ai">Weights & Biases</a>上跟踪你的实验，在VideoTokenizerTrainer上设置use_wandb_tracking = True，然后使用.trackers上下文管理器


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