14个文本转图像AI API
扩散模型实战（九）：使用CLIP模型引导和控制扩散模型
2024-12-09
上篇文章中介绍了如何微调扩散模型，有时候微调的效果仍然不能满足需求，比如图片编辑，3D模型输出等都需要对生成的内容进行控制，本文将初步探索一下如何控制扩散模型的输出。
我们将使用在LSUM bedrooms数据集上训练并在WikiArt数据集上微调的模型，首先加载模型来查看一下模型的生成效果：
!pip install -qq diffusers datasets accelerate wandb open-clip-torchimport numpy as np
import torch
import torch.nn.functional as F
import torchvision
from datasets import load_dataset
from diffusers import DDIMScheduler, DDPMPipeline
from matplotlib import pyplot as plt
from PIL import Image
from torchvision import transforms
from tqdm.auto import tqdmdevice = (
"mps"
if torch.backends.mps.is_available()
else "cuda"
if torch.cuda.is_available()
else "cpu"
)# 载入一个预训练过的管线
pipeline_name = "johnowhitaker/sd-class-wikiart-from-bedrooms"
image_pipe = DDPMPipeline.from_pretrained(pipeline_name).to(device)
# 使用DDIM调度器，仅用40步生成一些图片
scheduler = DDIMScheduler.from_pretrained(pipeline_name)
scheduler.set_timesteps(num_inference_steps=40)
# 将随机噪声作为出发点
x = torch.randn(8, 3, 256, 256).to(device)
# 使用一个最简单的采样循环
for i, t in tqdm(enumerate(scheduler.timesteps)):
model_input = scheduler.scale_model_input(x, t)
with torch.no_grad():
noise_pred = image_pipe.unet(model_input, t)["sample"]
x = scheduler.step(noise_pred, t, x).prev_sample
# 查看生成结果，如图5-10所示
grid = torchvision.utils.make_grid(x, nrow=4)
plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5)
正如上图所示，模型可以生成一些图片，那么如何进行控制输出呢？下面我们以控制图片生成绿色风格为例介绍AIGC模型控制：
思路是：定义一个均方误差损失函数，让生成的图片像素值尽量接近目标颜色；
def color_loss(images, target_color=(0.1, 0.9, 0.5)):
"""给定一个RGB值，返回一个损失值，用于衡量图片的像素值与目标颜色相差多少；
这里的目标颜色是一种浅蓝绿色，对应的RGB值为(0.1, 0.9, 0.5)"""
target = (
torch.tensor(target_color).to(images.device) * 2 - 1
) # 首先对target_color进行归一化，使它的取值区间为(-1, 1)
target = target[
None, :, None, None
] # 将所生成目标张量的形状改为(b, c, h, w)，以适配输入图像images的
# 张量形状
error = torch.abs(
images - target
).mean() # 计算图片的像素值以及目标颜色的均方误差
return error接下来，需要修改采样循环操作，具体操作步骤如下：
- 创建输入图像X，并设置requires_grad设置为True；
- 计算“去噪”后的图像X0；
- 将“去噪”后的图像X0传递给损失函数；
- 计算损失函数对输入图像X的梯度；
- 在使用调度器之前，先用计算出来的梯度修改输入图像X，使输入图像X朝着减少损失值的方向改进
实现上述步骤有两种方法：
方法一：从UNet中获取噪声预测，并将输入图像X的requires_grad属性设置为True，这样可以充分利用内存（因为不需要通过扩散模型追踪梯度），但是这会导致梯度的精度降低；
方法二：先将输入图像X的requires_grad属性设置为True，然后传递给UNet并计算“去噪”后的图像X0；
下面分别看一下这两种方法的效果：
# 第一种方法
# guidance_loss_scale用于决定引导的强度有多大
guidance_loss_scale = 40 # 可设定为5~100的任意数字
x = torch.randn(8, 3, 256, 256).to(device)
for i, t in tqdm(enumerate(scheduler.timesteps)):
# 准备模型输入
model_input = scheduler.scale_model_input(x, t)
# 预测噪声
with torch.no_grad():
noise_pred = image_pipe.unet(model_input, t)["sample"]
# 设置x.requires_grad为True
x = x.detach().requires_grad_()
# 得到“去噪”后的图像
x0 = scheduler.step(noise_pred, t, x).pred_original_sample
# 计算损失值
loss = color_loss(x0) * guidance_loss_scale
if i % 10 == 0:
print(i, "loss:", loss.item())
# 获取梯度
cond_grad = -torch.autograd.grad(loss, x)[0]
# 使用梯度更新x
x = x.detach() + cond_grad
# 使用调度器更新x
x = scheduler.step(noise_pred, t, x).prev_sample
# 查看结果
grid = torchvision.utils.make_grid(x, nrow=4)
im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
Image.fromarray(np.array(im * 255).astype(np.uint8))# 输出
0 loss: 29.37018394470215
10 loss: 12.116650581359863
20 loss: 11.641704559326172
30 loss: 11.78276252746582
# 第二种方法：在模型预测前设置好x.requires_grad
guidance_loss_scale = 40
x = torch.randn(4, 3, 256, 256).to(device)
for i, t in tqdm(enumerate(scheduler.timesteps)):
# 首先设置好requires_grad
x = x.detach().requires_grad_()
model_input = scheduler.scale_model_input(x, t)
# 预测
noise_pred = image_pipe.unet(model_input, t)["sample"]
# 得到“去噪”后的图像
x0 = scheduler.step(noise_pred, t, x).pred_original_sample
# 计算损失值
loss = color_loss(x0) * guidance_loss_scale
if i % 10 == 0:
print(i, "loss:", loss.item())
# 获取梯度
cond_grad = -torch.autograd.grad(loss, x)[0]
# 根据梯度修改x
x = x.detach() + cond_grad
# 使用调度器更新x
x = scheduler.step(noise_pred, t, x).prev_sample
grid = torchvision.utils.make_grid(x, nrow=4)
im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
Image.fromarray(np.array(im * 255).astype(np.uint8))# 输出
0 loss: 27.62268829345703
10 loss: 16.842506408691406
20 loss: 15.54642105102539
30 loss: 15.545379638671875
从上图看出，第二种方法效果略差，但是第二种方法的输出更接近训练模型所使用的数据，也可以通过修改guidance_loss_scale参数来增强颜色的迁移效果。
CLIP控制图像生成
虽然上述方式可以引导和控制图像生成某种颜色，但现在LLM更主流的方式是通过Prompt（仅仅打几行字描述需求）来得到自己想要的图像，那么CLIP是一个不错的选择。CLIP是有OpenAI开发的图文匹配大模型，由于这个过程是可微分的，所以可以将其作为损失函数来引导扩散模型。
使用CLIP控制图像生成的基本流程如下：
- 使用CLIP模型对Prompt表示为512embedding向量；
- 在扩散模型的生成过程中需要多次执行如下步骤：1）生成多个“去噪”图像；2）对生成的每个“去噪”图像用CLIP模型进行embedding，并对Prompt embedding和图像的embedding进行对比；3）计算Prompt和“去噪”后图像的梯度，使用这个梯度先更新输入图像X，然后再使用调度器更新X；加载CLIP模型
import open_clip
clip_model, _, preprocess = open_clip.create_model_and_transforms(
"ViT-B-32", pretrained="openai"
)
clip_model.to(device)
# 图像变换：用于修改图像尺寸和增广数据，同时归一化数据，以使数据能够适配CLIP模型
tfms = torchvision.transforms.Compose(
[
torchvision.transforms.RandomResizedCrop(224),# 随机裁剪
torchvision.transforms.RandomAffine(5), # 随机扭曲图片
torchvision.transforms.RandomHorizontalFlip(),# 随机左右镜像，
# 你也可以使用其他增广方法
torchvision.transforms.Normalize(
mean=(0.48145466, 0.4578275, 0.40821073),
std=(0.26862954, 0.26130258, 0.27577711),
),
]
)
# 定义一个损失函数，用于获取图片的特征，然后与提示文字的特征进行对比
def clip_loss(image, text_features):
image_features = clip_model.encode_image(
tfms(image)
) # 注意施加上面定义好的变换
input_normed = torch.nn.functional.normalize(image_features.
unsqueeze(1), dim=2)
embed_normed = torch.nn.functional.normalize(text_features.
unsqueeze(0), dim=2)
dists = (
input_normed.sub(embed_normed).norm(dim=2).div(2).
arcsin().pow(2).mul(2)
) # 使用Squared Great Circle Distance计算距离
return dists.mean()下面是引导模型生成图像的过程，步骤与上述类似，只需要把color_loss()替换成CLIP的损失函数
prompt = "Red Rose (still life), red flower painting"
# 读者可以探索一下这些超参数的影响
guidance_scale = 8
n_cuts = 4
# 这里使用稍微多一些的步数
scheduler.set_timesteps(50)
# 使用CLIP从提示文字中提取特征
text = open_clip.tokenize([prompt]).to(device)
with torch.no_grad(), torch.cuda.amp.autocast():
text_features = clip_model.encode_text(text)
x = torch.randn(4, 3, 256, 256).to(
device
)
for i, t in tqdm(enumerate(scheduler.timesteps)):
model_input = scheduler.scale_model_input(x, t)
# 预测噪声
with torch.no_grad():
noise_pred = image_pipe.unet(model_input, t)["sample"]
cond_grad = 0
for cut in range(n_cuts):
# 设置输入图像的requires_grad属性为True
x = x.detach().requires_grad_()
# 获得“去噪”后的图像
x0 = scheduler.step(noise_pred, t, x).pred_original_sample
# 计算损失值
loss = clip_loss(x0, text_features) * guidance_scale
# 获取梯度并使用n_cuts进行平均
cond_grad -= torch.autograd.grad(loss, x)[0] / n_cuts
if i % 25 == 0:
print("Step:", i, ", Guidance loss:", loss.item())
# 根据这个梯度更新x
alpha_bar = scheduler.alphas_cumprod[i]
x = (
x.detach() + cond_grad * alpha_bar.sqrt()
) # 注意这里的缩放因子
# 使用调度器更新x
x = scheduler.step(noise_pred, t, x).prev_sample
grid = torchvision.utils.make_grid(x.detach(), nrow=4)
im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
Image.fromarray(np.array(im * 255).astype(np.uint8))# 输出
Step: 0 , Guidance loss: 7.418107986450195
Step: 25 , Guidance loss: 7.085518836975098
上述生成的图像虽然不够完美，但可以调整一些超参数，比如梯度缩放因子alpha_bar.sqrt()，虽然理论上存在所谓的正确的缩放这些梯度方法，但在实践中仍需要实验来检验，下面介绍一些常用的方案：
plt.plot([1 for a in scheduler.alphas_cumprod], label="no scaling")
plt.plot([a for a in scheduler.alphas_cumprod], label="alpha_bar")
plt.plot([a.sqrt() for a in scheduler.alphas_cumprod],
label="alpha_bar.sqrt()")
plt.plot(
[(1 - a).sqrt() for a in scheduler.alphas_cumprod], label="(1-
alpha_bar).sqrt()"
)
plt.legend()
文章转自微信公众号@ArronAI
