muse-maskgit-pytorch

Muse - Pytorch

使用Pytorch实现<a href="https://muse-model.github.io/">Muse</a>：基于掩码生成式Transformer的文本到图像生成

如果您有兴趣参与<a href="https://laion.ai/">LAION</a>社区的复现工作，请加入<a href="https://discord.gg/xBPBXfcFHd"><img alt="加入我们的Discord" src="https://img.shields.io/discord/823813159592001537?color=5865F2&logo=discord&logoColor=white"></a>

安装

$ pip install muse-maskgit-pytorch

使用方法

首先训练您的VAE - VQGanVAE

import torch
from muse_maskgit_pytorch import VQGanVAE, VQGanVAETrainer

vae = VQGanVAE(
    dim = 256,
    codebook_size = 65536
)

# 在包含尽可能多图像的文件夹上进行训练

trainer = VQGanVAETrainer(
    vae = vae,
    image_size = 128,             # 您可能想从小图像开始，然后逐步学习到更大的图像，但由于vae全是卷积层，它应该可以泛化到512（如论文中所述），无需在512上训练
    folder = '/path/to/images',
    batch_size = 4,
    grad_accum_every = 8,
    num_train_steps = 50000
).cuda()

trainer.train()

然后将训练好的VQGanVAE和一个Transformer传递给MaskGit

import torch
from muse_maskgit_pytorch import VQGanVAE, MaskGit, MaskGitTransformer

# 首先实例化您的vae

vae = VQGanVAE(
    dim = 256,
    codebook_size = 65536
).cuda()

vae.load('/path/to/vae.pt') # 您需要加载指数移动平均的VAE

# 然后将vae和transformer插入到MaskGit中，如下所示

# (1) 创建您的transformer / 注意力网络

transformer = MaskGitTransformer(
    num_tokens = 65536,       # 必须与上面的codebook_size相同
    seq_len = 256,            # 必须等于vae中的fmap_size ** 2
    dim = 512,                # 模型维度
    depth = 8,                # 深度
    dim_head = 64,            # 注意力头维度
    heads = 8,                # 注意力头数量
    ff_mult = 4,              # 前馈扩展因子
    t5_name = 't5-small',     # 您的T5名称
)

# (2) 将训练好的VAE和基础transformer传递给MaskGit

base_maskgit = MaskGit(
    vae = vae,                 # vqgan vae
    transformer = transformer, # transformer
    image_size = 256,          # 图像大小
    cond_drop_prob = 0.25,     # 条件丢弃概率，用于无分类器引导
).cuda()

# 准备您的训练文本和图像

texts = [
    '一个孩子在咬了一半的苹果中发现虫子时尖叫',
    '蜥蜴在沙漠中用两只脚奔跑',
    '醒来发现身处迷幻景观',
    '浅水中闪闪发光的贝壳'
]

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

# 将其输入到您的maskgit实例中，并将return_loss设置为True

loss = base_maskgit(
    images,
    texts = texts
)

loss.backward()

# 在大量数据上长时间进行这个过程
# 然后...

images = base_maskgit.generate(texts = [
    '远处一头鲸鱼跃出水面',
    '小女孩在生日蛋糕上吹蜡烛',
    '蓝色和绿色火花的烟花'
], cond_scale = 3.) # 无分类器引导的条件缩放

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

要训练超分辨率maskgit，您需要在MaskGit实例化时更改1个字段（现在您需要传入cond_image_size，作为之前被条件化的图像大小）

可选地，您可以为条件低分辨率图像传入一个不同的VAE作为cond_vae。默认情况下，它将使用vae来同时对超分辨率和低分辨率图像进行标记化。

import torch
import torch.nn.functional as F
from muse_maskgit_pytorch import VQGanVAE, MaskGit, MaskGitTransformer

# 首先实例化您的ViT VQGan VAE
# 一个由transformer组成的VQGan VAE

vae = VQGanVAE(
    dim = 256,
    codebook_size = 65536
).cuda()

vae.load('./path/to/vae.pt') # 您需要加载指数移动平均的VAE

# 然后将VqGan VAE插入到MaskGit中，如下所示

# (1) 创建您的transformer / 注意力网络

transformer = MaskGitTransformer(
    num_tokens = 65536,       # 必须与上面的codebook_size相同
    seq_len = 1024,           # 必须等于vae中的fmap_size ** 2
    dim = 512,                # 模型维度
    depth = 2,                # 深度
    dim_head = 64,            # 注意力头维度
    heads = 8,                # 注意力头数量
    ff_mult = 4,              # 前馈扩展因子
    t5_name = 't5-small',     # 您的T5名称
)

# (2) 将训练好的VAE和基础transformer传递给MaskGit

superres_maskgit = MaskGit(
    vae = vae,
    transformer = transformer,
    cond_drop_prob = 0.25,
    image_size = 512,                     # 更大的图像尺寸
    cond_image_size = 256,                # 条件图像尺寸 <- 必须设置这个
).cuda()

# 准备您的训练文本和图像

texts = [
    '一个孩子在咬了一半的苹果中发现虫子时尖叫',
    '蜥蜴在沙漠中用两只脚奔跑',
    '醒来发现身处迷幻景观',
    '浅水中闪闪发光的贝壳'
]

images = torch.randn(4, 3, 512, 512).cuda()

# 将其输入到您的maskgit实例中，并将return_loss设置为True

loss = superres_maskgit(
    images,
    texts = texts
)

loss.backward()

# 在大量数据上长时间进行这个过程
# 然后...

images = superres_maskgit.generate(
    texts = [
        '远处一头鲸鱼跃出水面',
        '小女孩在生日蛋糕上吹蜡烛',
        '蓝色和绿色火花的烟花',
        '醒来发现身处迷幻景观'
    ],
    cond_images = F.interpolate(images, 256),  # 生成超分辨率图像时必须传入条件图像
    cond_scale = 3.
)

images.shape # (4, 3, 512, 512)

现在把它们组合在一起

from muse_maskgit_pytorch import Muse

base_maskgit.load('./path/to/base.pt')

superres_maskgit.load('./path/to/superres.pt')

# 传入上面训练好的base_maskgit和superres_maskgit

muse = Muse(
    base = base_maskgit,
    superres = superres_maskgit
)
images = muse([
    '远处一头鲸鱼跃出水面',
    '小女孩吹灭生日蛋糕上的蜡烛',
    '蓝色和绿色闪烁的烟花',
    '醒来后看到迷幻景象'
])

images # List[PIL.Image.Image]

致谢

感谢<a href="https://stability.ai/">StabilityAI</a>的赞助，以及我的其他赞助商，让我能够独立地开源人工智能。
感谢<a href="https://huggingface.co/">🤗 Huggingface</a>提供的transformers和accelerate库，它们都非常出色。

待办事项

测试端到端流程
分离cond_images_or_ids，目前实现不正确
添加VAE的训练代码
添加嵌入的可选自条件
结合token critic论文，已在<a href="https://github.com/lucidrains/phenaki-pytorch">Phenaki</a>中实现
为maskgit接入accelerate训练代码