扩散模型的微调与引导,教程地址:微调与引导
讲解扩散模型两种主要的基于现有模型实现改造的方法
- 方法一:微调(fine-tuning)
- 我们将在新的数据集上重新训练已有的模型,来改变它原有的输出类型。
- 在新数据集上微调一个现有的扩散模型,这包括:
- 使用累积梯度的方法去应对训练的 batch 太小所带来的一些问题
- 在训练过程中,将样本上传到 Weights and Biases 来记录日志,以此来监控训练过程(通过附 加的实例脚本程序)
- 将最终结果管线(pipeline)保存下来,并上传到Hub
- 首先登陆hugging face,输入access tokens,方便上传模型。
# Code to log in to the Hugging Face Hub, needed for sharing models
# Make sure you use a token with WRITE access
from huggingface_hub import notebook_login
notebook_login()
- 导入必要库,torch (高性能的数据处理和深度学习网络),diffusers(扩散模型库),matplotlib(画图库),PIL(图像处理官方库),datasets(数据集),transforms(数据处理),tqdm(进度条显示),torchvision(pytorch的一部分)
import 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 tqdm
我们从一个训练管道DDPMPipeline,加载一个预训练模型google/ddpm-celebahq-256,生成最终的图像,过程比较慢,在生成的每一步中,模型都是接收一个加噪声的输入,并被要求去预测这个噪声(以此来估计完全没噪声的图片是什么样)。最初这些预测都还不算太好,因此我们把这一过程分解成很多步。但另一方面,使用1000+步来实现整个生成过程也不是必要的,因为最近的很多研究已经找到了使用尽可能少的步数来生成好的结果的方法。
- 现在引出我们的主角
DDIM
,通过减少迭代周期,来实现加速。
Diffusers 库中,这些采样方法是通过调度器(scheduler)来操控的,每次更新通过step()函数完成。为了生成图片,我们从随机噪声 𝑥
开始,每一个迭代周期(timestep)我们都送入模型一个带噪声的输入 𝑥
并把模型预测结果再次输入step()函数。这里返回的输出都被命名为prev_sample —— 之所以是“previous”,是因为我们是在时间上“后退”,即从高噪声到低噪声(这和前向扩散过程是相反的)。
我们实践一下看看:先载入一个 scheduler,这里用的是 DDIMScheduler(基于这篇论文:Denoising Diffusion Implicit Models)。与原版的 DDPM 相比,它可以用少得多的迭代周期来产生很不错的采样样本。
# Create new scheduler and set num inference steps
scheduler = DDIMScheduler.from_pretrained("google/ddpm-celebahq-256")
scheduler.set_timesteps(num_inference_steps=40)
我们来造出 4 个随机噪声图片,用它们跑一跑采样循环,同时观察一下每一步的 𝑥
和预测出的去噪版本:
# The random starting point
x = torch.randn(4, 3, 256, 256).to(device) # Batch of 4, 3-channel 256 x 256 px images
# Loop through the sampling timesteps
for i, t in tqdm(enumerate(scheduler.timesteps)):
# Prepare model input
model_input = scheduler.scale_model_input(x, t)
# Get the prediction
with torch.no_grad():
noise_pred = image_pipe.unet(model_input, t)["sample"]
# Calculate what the updated sample should look like with the scheduler
scheduler_output = scheduler.step(noise_pred, t, x)
# Update x
x = scheduler_output.prev_sample
# Occasionally display both x and the predicted denoised images
if i % 10 == 0 or i == len(scheduler.timesteps) - 1:
fig, axs = plt.subplots(1, 2, figsize=(12, 5))
grid = torchvision.utils.make_grid(x, nrow=4).permute(1, 2, 0)
axs[0].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)
axs[0].set_title(f"Current x (step {i})")
pred_x0 = (
scheduler_output.pred_original_sample
) # Not available for all schedulers
grid = torchvision.utils.make_grid(pred_x0, nrow=4).permute(1, 2, 0)
axs[1].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)
axs[1].set_title(f"Predicted denoised images (step {i})")
plt.show()
你也可以直接用新的调度器替换原有管线(pipeline)中的调度器,然后采样。像这样做:
image_pipe.scheduler = scheduler
images = image_pipe(num_inference_steps=40).images
images[0]
给我们一个预训练过的管线(pipeline),我们怎样使用新的训练数据重训模型来生成图片?
看起来这和我们从头训练模型是几乎一样的(正如我们在第一单元所见的一样),除了我们这里是用现有模型作为初始化的。让我们实践一下看看,并额外考虑几点我们要注意的东西。
首先,数据方面,你可以尝试用 Vintage Faces 数据集 或者这些动漫人脸图片来获取和这个人脸模型的原始训练数据类似的数据。但我们现在还是先用和第一单元一样的蝴蝶数据集吧。通过以下代码来下载蝴蝶数据集,并建立一个能按批(batch)采样图片的dataloader
我们这里使用的 batch size 很小(只有 4),因为我们的训练是基于较大的图片尺寸的(256px),并且我们的模型也很大,如果我们的 batch size 太高,GPU 的内存可能会不够用了。你可以减小图片尺寸,换取更大的 batch size 来加速训练,但这里的模型一开始都是基于生成 256px 尺寸的图片来设计和训练的。
# @markdown load and prepare a dataset:
# Not on Colab? Comments with #@ enable UI tweaks like headings or user inputs
# but can safely be ignored if you're working on a different platform.
dataset_name = "huggan/smithsonian_butterflies_subset" # @param
dataset = load_dataset(dataset_name, split="train")
image_size = 256 # @param
batch_size = 4 # @param
preprocess = transforms.Compose(
[
transforms.Resize((image_size, image_size)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
]
)
def transform(examples):
images = [preprocess(image.convert("RGB")) for image in examples["image"]]
return {"images": images}
dataset.set_transform(transform)
train_dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True
)
print("Previewing batch:")
batch = next(iter(train_dataloader))
grid = torchvision.utils.make_grid(batch["images"], nrow=4)
plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5);
现在我们看看训练循环。我们把要优化的目标参数设定为image_pipe.unet.parameters(),以此来更新预训练过的模型的权重。其它部分的代码基本上和第一单元例子中的对应部分一样。在 Colab 上跑的话,大约需要10分钟,你可以趁这个时间喝杯茶休息一下。
num_epochs = 2 # @param
lr = 1e-5 # 2param
grad_accumulation_steps = 2 # @param
optimizer = torch.optim.AdamW(image_pipe.unet.parameters(), lr=lr)
losses = []
for epoch in range(num_epochs):
for step, batch in tqdm(enumerate(train_dataloader), total=len(train_dataloader)):
clean_images = batch["images"].to(device)
# Sample noise to add to the images
noise = torch.randn(clean_images.shape).to(clean_images.device)
bs = clean_images.shape[0]
# Sample a random timestep for each image
timesteps = torch.randint(
0,
image_pipe.scheduler.num_train_timesteps,
(bs,),
device=clean_images.device,
).long()
# Add noise to the clean images according to the noise magnitude at each timestep
# (this is the forward diffusion process)
noisy_images = image_pipe.scheduler.add_noise(clean_images, noise, timesteps)
# Get the model prediction for the noise
noise_pred = image_pipe.unet(noisy_images, timesteps, return_dict=False)[0]
# Compare the prediction with the actual noise:
loss = F.mse_loss(
noise_pred, noise
) # NB - trying to predict noise (eps) not (noisy_ims-clean_ims) or just (clean_ims)
# Store for later plotting
losses.append(loss.item())
# Update the model parameters with the optimizer based on this loss
loss.backward(loss)
# Gradient accumulation:
if (step + 1) % grad_accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
print(
f"Epoch {epoch} average loss: {sum(losses[-len(train_dataloader):])/len(train_dataloader)}"
)
# Plot the loss curve:
plt.plot(losses)
损失函数
我们的损失值曲线简直像噪声一样混乱!这是因为每一次迭代我们都只用了四个训练样本,而且加到它们上面的噪声水平还都是随机挑选的。这对于训练来讲并不理想。一种弥补的措施是,我们使用一个非常小的学习率,限制每次更新的幅度。但我们还有一个更好的方法,既能得到和使用更大的 batch size 一样的收益,又不需要让我们的内存爆掉。
点击这里看看:gradient accumulation。如果我们多运行几次loss.backward()后再调用optimizer.step()和optimizer.zero_grad(),PyTorch 就会把梯度累积(加和)起来,这样多个批次的数据产生的更新信号就会被高效地融合在一起,产出一个单独的(更好的)梯度估计用于参数更新。这样做会减少参数更新的总次数,就正如我们使用更大的 batch size 时希望看到的一样。梯度累积是一个很多框架都会替你做的事情(比如这里:🤗 Accelerate makes this easy),但这里我们从头实现一遍也挺好的,因为这对你在 GPU 内存受限时训练模型非常有帮助。正如你在上面代码中看到的那样(在注释 # Gradient accumulation 后),其实也不需要你写很多代码。
考虑因素3: 即使这样,我们的训练还是挺慢的,而且每遍历完一轮数据集才打印出一行更新,这也不足以让我们知道我们的训练到底怎样了。我们也许还应该:
训练过程中时不时地生成点图像样本,供我们检查模型性能
在训练过程中,把诸如损失值和生成的图片样本在内的一些东西记录到日志里。你可以使用诸如 Weights and Biases 或 tensorboard 之类的工具
我创建了一个快速的脚本程序(finetune_model.py),使用了上述的训练代码并加入了少量日志记录功能
用这个模型生成点图
# @markdown Generate and plot some images:
x = torch.randn(8, 3, 256, 256).to(device) # Batch of 8
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
grid = torchvision.utils.make_grid(x, nrow=4)
plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5);
考虑因素4: 微调这个过程可能是难以预知的。如果我们训练很长时间,我们也许能看见一些生成得很完美的蝴蝶,但中间过程从模型自身讲也极其有趣,尤其是你对艺术风格感兴趣时!你可以试试短时间或长时间地观察一下训练过程,并试着该百年学习率,看看这会怎么影响模型的最终输出。
代码:用我们在WikiArt上面使用最小化样例脚本去微调一个模型
## 下载微调用的脚本:
!wget https://github.com/huggingface/diffusion-models-class/raw/main/unit2/finetune_model.py
## 运行脚本,在Vintage Face数据集上训练脚本
## (最好在终端里跑):
!python finetune_model.py --image_size 128 --batch_size 8 --num_epochs 16\
--grad_accumulation_steps 2 --start_model "google/ddpm-celebahq-256"\
--dataset_name "Norod78/Vintage-Faces-FFHQAligned" --wandb_project 'dm-finetune'\
--log_samples_every 100 --save_model_every 1000 --model_save_name 'vintageface'
#保存和载入微调过的管线
image_pipe.save_pretrained("my-finetuned-model")
#仿照第一单元 Introduction to Diffusers 中的内容把模型上传到Hub中
# @title Upload a locally saved pipeline to the hub
# Code to upload a pipeline saved locally to the hub
from huggingface_hub import HfApi, ModelCard, create_repo, get_full_repo_name
# Set up repo and upload files
model_name = "ddpm-celebahq-finetuned-butterflies-2epochs" # @param What you want it called on the hub
local_folder_name = "my-finetuned-model" # @param Created by the script or one you created via image_pipe.save_pretrained('save_name')
description = "Describe your model here" # @param
hub_model_id = get_full_repo_name(model_name)
create_repo(hub_model_id)
api = HfApi()
api.upload_folder(
folder_path=f"{local_folder_name}/scheduler", path_in_repo="", repo_id=hub_model_id
)
api.upload_folder(
folder_path=f"{local_folder_name}/unet", path_in_repo="", repo_id=hub_model_id
)
api.upload_file(
path_or_fileobj=f"{local_folder_name}/model_index.json",
path_in_repo="model_index.json",
repo_id=hub_model_id,
)
# Add a model card (optional but nice!)
content = f"""
---
license: mit
tags:
- pytorch
- diffusers
- unconditional-image-generation
- diffusion-models-class
---
# Example Fine-Tuned Model for Unit 2 of the [Diffusion Models Class 🧨](https://github.com/huggingface/diffusion-models-class)
{description}
## Usage
from diffusers import DDPMPipeline
pipeline = DDPMPipeline.from_pretrained('{hub_model_id}')
image = pipeline().images[0]
image
"""
card = ModelCard(content)
card.push_to_hub(hub_model_id)
现在微调完了你的第一个扩散模型
- 方法二 引导
一般而言,判断一次微调到底有多管用并不容易,而且“足够好的性能”在不同应用场景下代表什么水平也会有所不同。比如,如果你在一个很小的数据集上微调一个文本条件模型,比如stable diffusion模型,你可能希望模型极可能保留它原始训练所学习的东西以便于它能理解你的数据集没有涵盖的各种文本提示;同时你又希望它适配你的数据,以便它生成的东西和你的数据风格一致。这可能意味着,你需要使用一个很低的学习率,并配合对模型进行指数平均,就像这个关于创建一个宝可梦版 stable diffusion 模型的博客中的做法一样。别的情况下,你可能还想在一个新数据集上完全重训一个模型(就像我们前面从卧室图片到 wikiart 图片微调一样),那你就需要大的学习率和长时间训练了。即使从这里的损失值曲线看不出模型在进步,但生成样本已经很清楚地显示出了一个从原始数据到更有艺术范的风格迁移的过程了,虽然看着还是不太协调。
这就带领我们来到了下一部分,让我们看看怎样对这种模型施加额外的引导,来更好地控制模型的输出
如果我们想对生成的样本施加点控制,那需要怎么做呢?例如,我们想让生成的图片偏向于靠近某种颜色。该怎么做呢?这里我们要介绍引导(guidance),它可以用来在采样的过程中施加额外控制。
第一步,我们先创建一个函数,定义我们希望优化的一个指标(损失值)。这里是一个让生成的图片趋向于某种颜色的例子,它将图片像素值和目标颜色(这里用的是一种浅蓝绿色)对比,返回平均的误差
def color_loss(images, target_color=(0.1, 0.9, 0.5)):
"""Given a target color (R, G, B) return a loss for how far away on average
the images' pixels are from that color. Defaults to a light teal: (0.1, 0.9, 0.5)"""
target = (
torch.tensor(target_color).to(images.device) * 2 - 1
) # Map target color to (-1, 1)
target = target[
None, :, None, None
] # Get shape right to work with the images (b, c, h, w)
error = torch.abs(
images - target
).mean() # Mean absolute difference between the image pixels and the target color
return error
接下来,我们要修改采样循环,在每一步,我们要做这些事情:
- 创建一个新版的 x,并且 requires_grad = True
- 算出去噪后的版本(x0)
- 将预测出的x0送入我们的损失函数中
- 找到这个损失函数对于 x 的梯度
- 在我们使用调度器前,用这个梯度去修改 x ,希望 x 朝着能减低损失值的方向改进
这里有两种实现方法,你可以探索一下哪一种更好。第一,我们是在从 UNet 得到噪声预测后才给 x 设置 requires_grad 的,这样对内存来讲更高效一点(因为我们不用穿过扩散模型去追踪梯度),但这样做梯度的精度会低一点。第二种方法是,我们先给 x 设置 requires_grad,然后再送入 UNet 并计算预测出的 x0。
# Variant 1: shortcut method
# The guidance scale determines the strength of the effect
guidance_loss_scale = 40 # Explore changing this to 5, or 100
x = torch.randn(8, 3, 256, 256).to(device)
for i, t in tqdm(enumerate(scheduler.timesteps)):
# Prepare the model input
model_input = scheduler.scale_model_input(x, t)
# predict the noise residual
with torch.no_grad():
noise_pred = image_pipe.unet(model_input, t)["sample"]
# Set x.requires_grad to True
x = x.detach().requires_grad_()
# Get the predicted x0
x0 = scheduler.step(noise_pred, t, x).pred_original_sample
# Calculate loss
loss = color_loss(x0) * guidance_loss_scale
if i % 10 == 0:
print(i, "loss:", loss.item())
# Get gradient
cond_grad = -torch.autograd.grad(loss, x)[0]
# Modify x based on this gradient
x = x.detach() + cond_grad
# Now step with scheduler
x = scheduler.step(noise_pred, t, x).prev_sample
# View the output
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))
这里的第二种实现方法需要几乎第一种的两倍的 GPU 内存,即使这里我们用的 batch size 是 4 而不是 8。试试看你能不能看出点不同?想想为什么这里梯度更精确?
# Variant 2: setting x.requires_grad before calculating the model predictions
guidance_loss_scale = 40
x = torch.randn(4, 3, 256, 256).to(device)
for i, t in tqdm(enumerate(scheduler.timesteps)):
# Set requires_grad before the model forward pass
x = x.detach().requires_grad_()
model_input = scheduler.scale_model_input(x, t)
# predict (with grad this time)
noise_pred = image_pipe.unet(model_input, t)["sample"]
# Get the predicted x0:
x0 = scheduler.step(noise_pred, t, x).pred_original_sample
# Calculate loss
loss = color_loss(x0) * guidance_loss_scale
if i % 10 == 0:
print(i, "loss:", loss.item())
# Get gradient
cond_grad = -torch.autograd.grad(loss, x)[0]
# Modify x based on this gradient
x = x.detach() + cond_grad
# Now step with scheduler
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))
在第二种实现方法中,内存的需求更高了,但颜色迁移的效果却减弱了。你可能觉得这种方法其实不如第一种方法。但是,可以说这里的输出更接近于训练模型所使用的数据,你也永远可以通过增大 guidance_loss_scale 来加强颜色迁移的效果。使用哪种方案其实最终取决于哪种方案在实验中效果更好。
- CLIP 引导
CLIP 是一个由 OpenAI 开发的模型,它可以让我们拿图片和文字说明去作比较。这是个非常强大的功能,因为它让我们能量化一张图和一句提示语有多匹配。另外,由于这个过程是可微分的,我们可以使用它作为损失函数去引导我们的扩散模型。
引导生成的图片向某种颜色倾斜确实让我们多少对生成有所控制,但如果我们能仅仅打几行字描述一下就得到我们想要的图片呢?
这里我们不深究细节。基本的方法是:
- 给文字提示语做嵌入(embedding),为 CLIP 获取一个 512 维的 embedding
- 对于扩散模型的生成过程的每一步:
- 做出多个不同版本的预测出来的去噪图片(不同的变种可以提供一个更干净的损失信号)
- 对每一个预测出的去噪图片,用 CLIP 给图片做嵌入(embedding),并将这个嵌入和文字的嵌入做对比(用一种叫 Great Circle Distance Squared 的度量方法)
- 计算这个损失对于当前带噪的 x 的梯度,并在用调度器(scheduler)更新它之前用这个梯度去修改 x
如果你想看关于 CLIP 的更深入的讲解,你可以看看这个课程或这个关于 OpenCLIP 的报告,我们就是用 OpenCLIP 载入 CLIP 模型的。运行下列代码就可以载入一个 CLIP 模型:
# @markdown load a CLIP model and define the loss function
import open_clip
clip_model, _, preprocess = open_clip.create_model_and_transforms(
"ViT-B-32", pretrained="openai"
)
clip_model.to(device)
# Transforms to resize and augment an image + normalize to match CLIP's training data
tfms = torchvision.transforms.Compose(
[
torchvision.transforms.RandomResizedCrop(224), # Random CROP each time
torchvision.transforms.RandomAffine(
5
), # One possible random augmentation: skews the image
torchvision.transforms.RandomHorizontalFlip(), # You can add additional augmentations if you like
torchvision.transforms.Normalize(
mean=(0.48145466, 0.4578275, 0.40821073),
std=(0.26862954, 0.26130258, 0.27577711),
),
]
)
# And define a loss function that takes an image, embeds it and compares with
# the text features of the prompt
def clip_loss(image, text_features):
image_features = clip_model.encode_image(
tfms(image)
) # Note: applies the above transforms
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的损失函数:
# @markdown applying guidance using CLIP
prompt = "Red Rose (still life), red flower painting" # @param
# Explore changing this
guidance_scale = 8 # @param
n_cuts = 4 # @param
# More steps -> more time for the guidance to have an effect
scheduler.set_timesteps(50)
# We embed a prompt with CLIP as our target
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
) # RAM usage is high, you may want only 1 image at a time
for i, t in tqdm(enumerate(scheduler.timesteps)):
model_input = scheduler.scale_model_input(x, t)
# predict the noise residual
with torch.no_grad():
noise_pred = image_pipe.unet(model_input, t)["sample"]
cond_grad = 0
for cut in range(n_cuts):
# Set requires grad on x
x = x.detach().requires_grad_()
# Get the predicted x0:
x0 = scheduler.step(noise_pred, t, x).pred_original_sample
# Calculate loss
loss = clip_loss(x0, text_features) * guidance_scale
# Get gradient (scale by n_cuts since we want the average)
cond_grad -= torch.autograd.grad(loss, x)[0] / n_cuts
if i % 25 == 0:
print("Step:", i, ", Guidance loss:", loss.item())
# Modify x based on this gradient
alpha_bar = scheduler.alphas_cumprod[i]
x = (
x.detach() + cond_grad * alpha_bar.sqrt()
) # Note the additional scaling factor here!
# Now step with scheduler
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))
看起来有点像玫瑰!虽然还不够完美,但你如果接着调一调设定参数,你可以得到一些更令人满意的图片。
如果你仔细看上面的代码,你会发现我在使用alpha_bar.sqrt()作为因子去缩放梯度。虽然理论上有所谓正确的缩放这些梯度的方法,但在实践中,你可以用实验去验证。对于有些引导来说,你可能希望大部分的引导作用都集中在刚开始的几步里,对于另一些(比如一些关注点在纹理方面的风格损失函数)来讲,你可能希望它仅在生成过程的结束部分加入进来。对此,下面展示一些可能的方案.
你可以做点实验,针对不同的调整方案、引导规模大小(guidance_scale),以及任何你能想到的小技巧(比如用一个范围去截断梯度也很常见),试试你能把效果做到多好。也可以在其它模型上试试,比如我们最开始使用的那个人脸模型,你能很可靠地让它生成男性脸吗?如果你把 CLIP 引导和前面我们用过地基于颜色的损失函数结合起来用呢?
如果你看一些实操的 CLIP 引导的扩散模型代码,你会发现一种更复杂的方法:使用更好的类别去选取随机图像裁剪,并对损失函数进行许多额外的调整以获得更好的性能。在文本条件扩散模型出现前,这可是最好的文本到图像转换系统!我们这里的“玩具级”项目有很多可改进的空间,但不管怎样,它抓住了核心要点:借助于我们的引导和 CLIP 惊人的能力,我们可以给一个没有条件约束的扩散模型加上文本级的控制 🎨
- 把自定义的采样训练做成 Gradio 上的展示来分享
也许你现在已经想出了一个很好玩的损失函数去引导生成过程,现在你想把你的微调模型和自定义的采样策略分享给全世界…
点击这里了解 Gradio。Gradio 是一个免费的开源工具,让用户可以方便地通过一个简单的网页界面来创建和分享交互式的机器学习模型。使用Gradio,用户可以为自己的机器学习模型自定义接口,然后通过一个唯一的URL共享给他人。Gradio 也别集成入了 🤗 Spaces,使得创建 Demo 和共享给他人变得更容易。
我们将把我们需要的核心逻辑放在一个函数中,这个函数接收一些输入然后输出一张图片作为输出。然后这个函数会被封装在一个简单的接口中,让用户能自己定义一些参数(这些参数是作为输入提供给 generate 函数的)。这里有很多组件 可以用 —— 在这个例子中我们就加入了一个滑杆来控制引导的力度(guidance scale),以及一个颜色选择器供我们定义目标颜色。
import gradio as gr
from PIL import Image, ImageColor
# The function that does the hard work
def generate(color, guidance_loss_scale):
target_color = ImageColor.getcolor(color, "RGB") # Target color as RGB
target_color = [a / 255 for a in target_color] # Rescale from (0, 255) to (0, 1)
x = torch.randn(1, 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 = x.detach().requires_grad_()
x0 = scheduler.step(noise_pred, t, x).pred_original_sample
loss = color_loss(x0, target_color) * guidance_loss_scale
cond_grad = -torch.autograd.grad(loss, x)[0]
x = x.detach() + cond_grad
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
im = Image.fromarray(np.array(im * 255).astype(np.uint8))
im.save("test.jpeg")
return im
# See the gradio docs for the types of inputs and outputs available
inputs = [
gr.ColorPicker(label="color", value="55FFAA"), # Add any inputs you need here
gr.Slider(label="guidance_scale", minimum=0, maximum=30, value=3),
]
outputs = gr.Image(label="result")
# And the minimal interface
demo = gr.Interface(
fn=generate,
inputs=inputs,
outputs=outputs,
examples=[
["#BB2266", 3],
["#44CCAA", 5], # You can provide some example inputs to get people started
],
)
demo.launch(debug=True) # debug=True allows you to see errors and output in Colab
当然你也可以做出点复杂得多的接口,并加入点炫酷的风格和很长的输入序列。这里我们只做最简单的演示。
在 🤗 Spaces 上的演示 demo 都是默认用 CPU 跑的,所以在你移交之前,把你的接口在 Colab 上制作原型(如上所示)也是很不错的选择。当你准备好你的 demo 时,你需要创建一个 space,按照 requirements.txt 安装程序所需的库,然后把所都代码都放在一个名为 app.py 的文件里,这个文件是用来定义相关函数和接口的。
幸运的是,你也可以复制一个 space。你可以在这里看看我的 demo,然后点击“Duplicate this Space”用我的代码作为一个模板,用于你后续修改代码来添加你自己的模型和引导函数。
在设置选项中,你也可以配置你的 space,让它在更厉害的硬件上跑(当然会按小时收费)。如果你确实做出了一些很惊艳的东西、想在更好的硬件上分享它,可是你却不想花钱?那你可以在 Discord 上告诉我们,看我们能不能提供点帮助。
总结和下一步的工作
本节笔记本涵盖了好多内容啊!让我们回顾一下核心要点:
- 载入一个现有模型并用不同的调度器(scheduler)去采样其实很容易
- 微调(fine-tuning)看起来很像从头训练一个模型,唯一不同的是我们用已有的模型做初始化,以此我们希望能快点得到更好效果
- 如果在大尺寸图片上微调大模型,我们可以用诸如梯度累积(gradient accumulation)的方法去应对训练时batch size太小的问题
- 把采样的图片保存到日志里对微调很重要,而损失值曲线却可能无法反映有用的信息
- 引导(guidance)可以让我们在使用一个没有条件约束的模型时,通过一些引导或损失函数来掌控生成过程。这里我们每步都会找一个损失对于带噪图片x的梯度,然后用梯度更新这个带噪图片,之后再进入下一个生成迭代
- 用CLIP引导让我们可以用文字描述去控制一个没有条件约束的模型的生成过程!
如果你想在实践中运用这些知识,你还可以做这些:
- 微调你自己的模型并把它上传到 Hub。这包括:首先找一个起始点(比如一个在 [faces](https://huggingface.co/google/ddpm-celebahq-256%EF%BC%89%E3%80%81%5Bbedrooms%5D(https://huggingface.co/fusing/ddpm-lsun-bedroom)%E3%80%81%5Bcats%5D(https://huggingface.co/fusing/ddpm-lsun-cat) 或 wikiart example above 数据集上训练过的模型)和一个新的数据集(也许你可以用这个 animal faces,或者你自己的图片);然后跑一跑这个笔记本中的代码或样例脚本程序。
- 用你的微调过的模型探索一下引导,你可以用我们例子中的引导函数(颜色损失函数或 CLIP),也可以自己创造一个。
- 把你的 demo 分享到 Gradio 上,你可以通过修改这里的 space示例,也可以创建你自己的有更多功能的 space