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

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)

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)

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)