AI绘画Imagen大力出奇迹生成图像

AI绘画Imagen大力出奇迹生成图像

介绍

Imagen是一个文本到图像的扩散模型，由Google大脑团队研究所开发。

Imagen通过创新的设计，摈弃了需要预训练视觉-语言模型的繁琐步骤，直接采用了T5等大规模语言模型作为文本编码器,与扩散模型有机结合，完成了从文本到图像的直接关联映射。

这种结合语言模型与扩散模型的端到端方式，充分利用了T5作为纯文本模型的优势，包括参数规模的可拓展性和丰富的文本预训练数据，比依赖视觉信息的CLIP等模型更加灵活和直接。

Imagen的结果表明，单纯依靠语义理解力极强的语言模型就可以完成逼真的图像合成，而不需要额外引入视觉模型作为“桥梁”。这为未来在相同框架下，继续扩大语言模型规模与语义理解能力，达到文本到图像生成更上一层楼，提供了可能性与期待。

论文作者认为有两个特点：

Imagen架构

Imagen模型是一个文本到图像的生成模型，它首先使用一个大型的固定T5-XXL编码器将输入文本编码成嵌入向量。然后，一个条件扩散模型将文本嵌入映射成一个64×64像素的图像。最后，Imagen利用文本条件的超分辨率扩散模型将图像从64×64像素上采样到256×256像素，然后进一步上采样到1024×1024像素。

https://arxiv.org/pdf/2205.11487.pdf

扩散模型的去噪目标

公式 (1) 描述了扩散模型的去噪目标：
E x , c , ϵ ; t [ w t ⋅ ( x ^ t ( x ˉ t + ϵ ˉ t ; c ) − x ) 2 ] E_{x,c,\epsilon;t} \left[ w_t \cdot \left( \hat{x}_t(\bar{x}_t + \bar{\epsilon}_t; c) - x \right)^2 \right] Ex,c,ϵ;t​[wt​⋅(x^t​(xˉt​+ϵˉt​;c)−x)2]
其中， ( x ; c ) (x;c) (x;c) 是数据-条件对， t ∼ U ( [ 0 ; 1 ] ) t \sim U([0;1]) t∼U([0;1])， ϵ ∼ N ( 0 ; I ) \epsilon \sim N(0;I) ϵ∼N(0;I)，而 x ˉ t , ϵ ˉ t , w t \bar{x}_t, \bar{\epsilon}_t, w_t xˉt​,ϵˉt​,wt​ 是 t 的函数，影响样本质量。直观地说， x ^ t \hat{x}_t x^t​被训练为使用平方误差损失去噪 z t : = x ˉ t + ϵ ˉ t ϵ z_t := \bar{x}_t + \bar{\epsilon}_t \epsilon zt​:=xˉt​+ϵˉt​ϵ 至 x 。

无分类器引导

公式 (2) 描述了无分类器引导的调整预测：
( ϵ ~ t ) = ϵ ˉ t ,  where  ϵ ~ t ( z t ; c ) = w ⋅ ϵ t ( z t ; c ) + ( 1 − w ) ⋅ ϵ t ( z t ) (\tilde{\epsilon}_t) = \bar{\epsilon}_t, \text{ where } \tilde{\epsilon}_t(zt;c) = w \cdot \epsilon_t(zt;c) + (1 - w) \cdot \epsilon_t(zt) (ϵ~t​)=ϵˉt​, where ϵ~t​(zt;c)=w⋅ϵt​(zt;c)+(1−w)⋅ϵt​(zt)
这里， ϵ t ( z t ; c ) \epsilon_t(zt;c) ϵt​(zt;c)和 ϵ t ( z t ) \epsilon_t(zt) ϵt​(zt)

是条件和无条件的 ϵ \epsilon ϵ-预测，由 ϵ t : = ( z t − x ˉ t ⋅ x ^ t ) / ϵ ˉ t \epsilon_t := (zt - \bar{x}_t \cdot \hat{x}_t) / \bar{\epsilon}_t ϵt​:=(zt−xˉt​⋅x^t​)/ϵˉt​
给出，w 是引导权重。设置 w = 1 禁用无分类器引导，而增加 w > 1 加强引导的效果。

T5文本编码器

T5模型（Text-to-Text Transfer Transformer）处理不同任务的示例。T5模型是一个自然语言处理模型，设计用来处理各种文本任务，如翻译、摘要和问答。这个模型将所有任务视为文本到文本的问题，因此可以使用相同的模型架构处理翻译、分类、回归等多种不同类型的任务。

T5模型在Imagen中的作用

文本编码：T5模型接收一段文本并将其编码成一个高维的文本表征。这个表征捕捉了文本的意义和细微差别，这些都是生成与文本描述匹配的图像所必需的。

指导扩散模型：这个文本表征随后被用作条件输入，来指导图像的扩散过程。扩散模型从一个随机噪声图像开始，并逐渐将其转化为一个清晰的图像，这个过程在每一步都考虑到了T5提供的文本表征。

逐步精细化：生成的初始图像是低分辨率的，然后通过一系列超分辨率扩散模型逐步精细化。这些超分辨率模型也同样使用T5的文本表征作为条件，以确保在放大过程中细节的增加仍然符合文本描述。

Imagen与unCLIP的不同

CLIP模型：CLIP通过大量的图像和文本对进行训练，学会了理解图像内容与自然语言描述之间的关联。CLIP是一个双模态模型，它可以同时理解文本和图像，并且在两者之间建立了一个共同的表征空间。

unCLIP：unCLIP方法是将CLIP的文本和图像编码器联合使用，以生成与文本描述匹配的图像。unCLIP通过优化图像来匹配CLIP模型的文本-图像嵌入，而非使用扩散过程。

Imagen与unCLIP的区别：Imagen使用T5模型来生成文本表征，然后通过一个扩散过程生成图像，而不是优化图像以匹配预训练的文本-图像嵌入。Imagen的方法侧重于生成高质量、高分辨率的图像，而unCLIP侧重于图像和文本的紧密匹配。

对齐性（Alignment）：指的是生成的图像与文本描述的对齐程度，即图像是否准确反映了文本的内容。

保真度（Fidelity）：指的是图像的质量，即图像的清晰度、细节以及是否在视觉上令人信服。

新的扩散采样技术

Imagen的动态阈值（Dynamic Thresholding）技术能够生成更真实的图像，是因为它在采样过程中积极地防止像素饱和。

具体来说，动态阈值技术通过向内推动饱和像素（那些接近-1和1的像素值），在每一步中主动避免像素达到饱和状态​​。

这一点对于维持照片的真实感是至关重要的，因为饱和像素可能会损害图像细节，导致颜色失真。特别是在使用非常大的引导权重时，动态阈值对于保持高度的照片真实感和图像与文本的对齐有着显著的效果。

计算阈值：首先，为了决定每个采样步骤中使用的阈值 ( s )，需要计算当前像素值 x 0 t x0_t x0t​ 的绝对值的某个百分位数 ( p )。这个操作是通过 jnp.percentile() 函数进行的，其中 ( p ) 是一个超参数，代表所选择的百分位数。

调整阈值：计算得到的百分位数 ( s ) 与1.0进行比较，取两者中的最大值作为最终的阈值。这样可以确保阈值至少为1.0，从而避免任何像素值超过 [ − 1 , 1 ] [-1, 1] [−1,1]的范围。

应用阈值：然后将像素值 x 0 t x0_t x0t​ 限制在 ([-s, s]) 的范围内，并通过 ( s ) 进行归一化，即 x0_t = jnp.clip(x0_t, -s, s) / s。这样做能够防止像素饱和，并保持图像的动态范围。

采样步骤：接下来，执行采样步骤 sampler_step(x0_t, z_t, t)，其中 ( z_t ) 是上一步的隐状态，( t ) 是当前的时间步。然后更新隐状态 ( z_t ) 为 z_tm1。

通过上述步骤，动态阈值技术在保持图像细节和防止过度饱和方面比静态阈值技术表现得更好，尤其是在使用大引导权重时。这是因为它能够适应不同像素值的分布和动态范围，避免了像素值的极端饱和，从而生成了更真实且与文本更加一致的图像。

Imagen 源码

https://github.com/lucidrains/imagen-pytorch

import torch
from imagen_pytorch import Unet, Imagen

# 定义Imagen的第一个U-net模型，设置维度、条件维度、维度倍数、ResNet块数量以及注意力层
unet1 = Unet(
    dim = 32,  # U-net的维度
    cond_dim = 512,  # 条件维度，通常为文本嵌入的维度
    dim_mults = (1, 2, 4, 8),  # 每个阶段维度的倍数
    num_resnet_blocks = 3,  # ResNet块的数量
    layer_attns = (False, True, True, True),  # 每个阶段是否使用自注意力层
    layer_cross_attns = (False, True, True, True)  # 每个阶段是否使用交叉注意力层
)

# 定义第二个U-net模型，与第一个有不同的配置
unet2 = Unet(
    dim = 32,
    cond_dim = 512,
    dim_mults = (1, 2, 4, 8),
    num_resnet_blocks = (2, 4, 8, 8),
    layer_attns = (False, False, False, True),
    layer_cross_attns = (False, False, False, True)
)

# 创建Imagen模型，包含上面定义的两个U-net模型
imagen = Imagen(
    unets = (unet1, unet2),  # 包含的U-net模型
    image_sizes = (64, 256),  # 生成图像的尺寸，这里定义了两个尺寸
    timesteps = 1000,  # 时间步长
    cond_drop_prob = 0.1  # 条件丢弃概率
).cuda()

# 准备模拟数据：文本嵌入和图像
text_embeds = torch.randn(4, 256, 768).cuda()  # 文本嵌入
images = torch.randn(4, 3, 256, 256).cuda()  # 图像

# 使用Imagen模型训练每个U-net模型
for i in (1, 2):  # 针对每个U-net
    loss = imagen(images, text_embeds = text_embeds, unet_number = i)  # 计算损失
    loss.backward()  # 反向传播

# 通过文本嵌入生成图像
images = imagen.sample(texts = [
    'a whale breaching from afar',
    'young girl blowing out candles on her birthday cake',
    'fireworks with blue and green sparkles'
], cond_scale = 3.)  # 条件缩放因子

images.shape # 输出图像的形状 (3, 3, 256, 256)

在imagen_pytorch库中，ImagenTrainer包装类简化了Imagen模型的训练过程。这个类的一个关键功能是自动处理级联去噪扩散概率模型（DDPM）中所有U-net的指数移动平均（EMA）。

当在ImagenTrainer实例上调用update方法时，它会更新每个U-net的EMA，确保模型权重随时间进行平均。这个过程对于稳定训练和提高最终图像生成质量至关重要。

import torch
from imagen_pytorch import Unet, Imagen, ImagenTrainer

# 初始化Imagen的两个U-net模型
unet1 = Unet(
    dim = 32,  # U-net的基础维度
    cond_dim = 512,  # 条件编码的维度
    dim_mults = (1, 2, 4, 8),  # 定义每个U-net块的维度倍增
    num_resnet_blocks = 3,  # ResNet块的数量
    layer_attns = (False, True, True, True),  # 定义哪些层使用自注意力
)

unet2 = Unet(
    dim = 32,
    cond_dim = 512,
    dim_mults = (1, 2, 4, 8),
    num_resnet_blocks = (2, 4, 8, 8),  # 不同层的ResNet块数量
    layer_attns = (False, False, False, True),  # 自注意力层配置
    layer_cross_attns = (False, False, False, True)  # 交叉注意力层配置
)

# 创建包含上述U-net模型的Imagen模型
imagen = Imagen(
    unets = (unet1, unet2),  # 包含的U-net模型列表
    text_encoder_name = 't5-large',  # 文本编码器的名称
    image_sizes = (64, 256),  # 生成图像的尺寸
    timesteps = 1000,  # 时间步数
    cond_drop_prob = 0.1  # 条件丢失概率
).cuda()  # 将模型移动到GPU

# 使用ImagenTrainer类封装Imagen模型
trainer = ImagenTrainer(imagen)

# 生成模拟的文本嵌入和图像数据
text_embeds = torch.randn(64, 256, 1024).cuda()  # 随机生成的文本嵌入
images = torch.randn(64, 3, 256, 256).cuda()  # 随机生成的图像

# 将图像和文本嵌入数据输入到Imagen模型中，训练每个U-net模型
loss = trainer(
    images,
    text_embeds = text_embeds,
    unet_number = 1,  # 在这个例子中训练第一个U-net
    max_batch_size = 4  # 将64个样本的批次分成16个批次进行训练
)

trainer.update(unet_number = 1)  # 更新模型参数

# 经过多次迭代后，可以根据文本嵌入从级联的DDPM中采样图像
images = trainer.sample(texts = [
 'a puppy looking anxiously at a giant donut on the table',
    'the milky way galaxy in the style of monet'
], cond_scale = 3.)  # 条件缩放因子，控制图像生成的一致性

images.shape # 输出图像的形状 (2, 3, 256, 256)

imagen_pytorch 库进行无条件图像生成的训练，即不依赖于文本输入。

首先定义了两个U-net网络，一个是标准的Unet，另一个是专门为256像素图像设计的SRUnet256。

然后创建了一个Imagen模型，明确设置condition_on_text = False，表示不依赖于文本条件。接着使用ImagenTrainer类对模型进行训练。

代码中先对一组随机生成的图像数据进行训练，然后更新模型。最后，从训练好的模型中无条件地采样生成图像。这个过程演示了如何在不依赖文本输入的情况下使用Imagen模型进行图像生成。

import torch
from imagen_pytorch import Unet, Imagen, SRUnet256, ImagenTrainer

# 定义无条件图像生成所需的两个U-net模型
unet1 = Unet(
    dim = 32,  # 基础维度
    dim_mults = (1, 2, 4),  # 维度倍增配置
    num_resnet_blocks = 3,  # ResNet块的数量
    layer_attns = (False, True, True),  # 自注意力层配置
    layer_cross_attns = False,  # 交叉注意力层配置
    use_linear_attn = True  # 是否使用线性注意力
)

unet2 = SRUnet256(
    dim = 32,
    dim_mults = (1, 2, 4),
    num_resnet_blocks = (2, 4, 8),  # ResNet块数量
    layer_attns = (False, False, True),  # 自注意力层配置
    layer_cross_attns = False  # 交叉注意力层配置
)

# 创建Imagen模型，包含上述U-net模型，用于无条件图像生成
imagen = Imagen(
    condition_on_text = False,  # 设置为无条件生成
    unets = (unet1, unet2),  # 包含的U-net模型
    image_sizes = (64, 128),  # 生成图像的尺寸
    timesteps = 1000  # 时间步长
).cuda()  # 移至GPU

# 初始化Imagen训练器
trainer = ImagenTrainer(imagen).cuda()

# 准备训练图像数据
training_images = torch.randn(4, 3, 256, 256).cuda()  # 随机生成的训练图像

# 训练过程：分别训练每个U-net模型
loss = trainer(training_images, unet_number = 1)  # 训练第一个U-net
trainer.update(unet_number = 1)  # 更新第一个U-net的参数

# 无条件地从级联U-net模型中采样生成图像
images = trainer.sample(batch_size = 16)  # 采样生成的图像批次 (16, 3, 128, 128)