如何训练一个简单的stable diffusion模型(附详细注释）

注：代码来自https://github.com/darcula1993/diffusion-models-class-CN/blob/main/unit1/01_introduction_to_diffusers_CN.ipynb 

本文是本人学习后的的尝试以及注解

一、准备工作

"""

这行命令使用pip工具来安装或升级多个Python包。具体来说，它执行以下操作：

-qq：这是pip的安静模式选项，它会减少输出信息，只显示关键信息，使安装过程更为简洁。

-U：这是pip的升级选项，它指示pip升级已经安装的包到最新版本（如果存在新版本）。

接下来，列出了要安装或升级的包：

diffusers：一个Python包

datasets：Hugging Face Transformers库的一部分，用于提供和管理各种自然语言处理（NLP）数据集的工具。

transformers：Hugging Face Transformers库，提供了预训练的深度学习模型，用于自然语言处理和文本生成任务。

accelerate：Hugging Face库的一部分，用于加速深度学习模型的训练和推理。

ftfy：一个用于处理Unicode文本的Python库，用于修复和清理不规范的Unicode文本。

pyarrow：正如前面所提到的，pyarrow是一个用于高效处理大规模数据集的Python库，支持列式存储和跨语言互操作性。
"""

%pip install -qq -U diffusers datasets transformers accelerate ftfy pyarrow

# 登录hugging face
from huggingface_hub import notebook_login

notebook_login()

显示下图则登陆成功：

# 安装 Git LFS 来上传模型检查点：

%%capture
!sudo apt -qq install git-lfs
!git config --global credential.helper store

# 导入将要使用的库，并定义一些方便函数，稍后将会使用这些函数
import numpy as np
import torch
import torch.nn.functional as F
from matplotlib import pyplot as plt
from PIL import Image


def show_images(x):
  """Given a batch of images x, make a grid and convert to PIL"""
  '''
  输入参数：x，一个批量的图像数据（通常是PyTorch张量）。
  功能：将输入的图像数据从范围(-1, 1)映射到(0, 1)，然后将这些图像排列成一个网格，并将网格转换为PIL图像。
  返回值：返回一个PIL图像，其中包含了排列好的输入图像网格。
  '''
  x = x * 0.5 + 0.5  # Map from (-1, 1) back to (0, 1)
  grid = torchvision.utils.make_grid(x)
  grid_im = grid.detach().cpu().permute(1, 2, 0).clip(0, 1) * 255
  grid_im = Image.fromarray(np.array(grid_im).astype(np.uint8))
  return grid_im


def make_grid(images, size=64):
  """Given a list of PIL images, stack them together into a line for easy viewing"""
  '''
  输入参数：images，一个包含多个PIL图像的列表，以及一个可选的size参数，用于指定图像的大小。
  功能：将多个PIL图像按照指定的大小堆叠在一行上，以便于查看。
  返回值：返回一个新的PIL图像，其中包含了堆叠在一行上的输入图像。
  '''
  output_im = Image.new("RGB", (size * len(images), size))
  for i, im in enumerate(images):
      output_im.paste(im.resize((size, size)), (i * size, 0))
  return output_im


# Mac users may need device = 'mps' (untested)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

二、下载训练数据集

# 下载一个来自 Hugging Face Hub 的图像集。具体来说，是个 1000 张蝴蝶图像收藏集
import torchvision  # 导入PyTorch的torchvision库
from datasets import load_dataset  # 导入数据集
from torchvision import transforms  # 导入PyTorch的transforms模块

# 使用load_dataset函数加载名为"huggan/smithsonian_butterflies_subset"的数据集中的训练集数据
dataset = load_dataset("huggan/smithsonian_butterflies_subset", split="train")

# 或者从本地文件夹加载图像数据
# dataset = load_dataset("imagefolder", data_dir="path/to/folder")

# 指定图像大小为32x32像素
image_size = 32
# 如果GPU内存不足，可以降低批次大小
batch_size = 64

# 定义数据增强操作
preprocess = transforms.Compose(
    [
        transforms.Resize((image_size, image_size)),  # 调整图像大小
        transforms.RandomHorizontalFlip(),  # 随机水平翻转（数据增强）
        transforms.ToTensor(),  # 将图像转换为张量（数值范围从0到1）
        transforms.Normalize([0.5], [0.5]),  # 将图像像素值归一化到(-1, 1)的范围
    ]
)

# 定义一个用于对数据进行转换的函数
def transform(examples):
    images = [preprocess(image.convert("RGB")) for image in examples["image"]]
    return {"images": images}

# 将数据集的转换函数设置为上述定义的transform函数
dataset.set_transform(transform)

# 创建一个数据加载器，用于以批次方式提供转换后的图像数据
train_dataloader = torch.utils.data.DataLoader(
    dataset, batch_size=batch_size, shuffle=True  # 使用指定的批次大小和随机打乱数据
)

# 我们可以从中取出一批图像数据来看一看他们是什么样子:

xb = next(iter(train_dataloader))["images"].to(device)[:8]
print("X shape:", xb.shape)
show_images(xb).resize((8 * 64, 64), resample=Image.NEAREST)

结果如下图所示：

三、定义管理器

我们的训练计划是，取出这些输入图片然后对它们增添噪声，在这之后把带噪的图片送入模型。在推理阶段，我们将用模型的预测值来不断迭代去除这些噪点。在diffusers中，这两个步骤都是由 管理器（调度器） 来处理的。噪声管理器决定在不同的迭代周期时分别加入多少噪声。

from diffusers import DDPMScheduler

noise_scheduler = DDPMScheduler(num_train_timesteps=1000)

# 绘图查看输入 (x) 与噪声是如何在不同迭代周期中量化和叠加的
plt.plot(noise_scheduler.alphas_cumprod.cpu() ** 0.5, label=r"${\sqrt{\bar{\alpha}_t}}$")
plt.plot((1 - noise_scheduler.alphas_cumprod.cpu()) ** 0.5, label=r"$\sqrt{(1 - \bar{\alpha}_t)}$")
plt.legend(fontsize="x-large");

# 使用noise_scheduler.add_noise功能来添加不同程度的噪声
timesteps = torch.linspace(0, 999, 8).long().to(device)
noise = torch.randn_like(xb)
noisy_xb = noise_scheduler.add_noise(xb, noise, timesteps)
print("Noisy X shape", noisy_xb.shape)
show_images(noisy_xb).resize((8 * 64, 64), resample=Image.NEAREST)

四、定义、训练模型

# 定义模型
from diffusers import UNet2DModel

# Create a model
model = UNet2DModel(
  sample_size=image_size,  # the target image resolution
  in_channels=3,  # the number of input channels, 3 for RGB images
  out_channels=3,  # the number of output channels
  layers_per_block=2,  # how many ResNet layers to use per UNet block
  block_out_channels=(64, 128, 128, 256),  # More channels -> more parameters
  down_block_types=(
      "DownBlock2D",  # a regular ResNet downsampling block
      "DownBlock2D",
      "AttnDownBlock2D",  # a ResNet downsampling block with spatial self-attention
      "AttnDownBlock2D",
  ),
  up_block_types=(
      "AttnUpBlock2D",
      "AttnUpBlock2D",  # a ResNet upsampling block with spatial self-attention
      "UpBlock2D",
      "UpBlock2D",  # a regular ResNet upsampling block
  ),
)
model.to(device);

开始训练模型 

# 创建了一个对象，用于调度噪声的添加。参数指定了训练的总步数和噪声的变化规律
noise_scheduler = DDPMScheduler(
    num_train_timesteps=1000, beta_schedule="squaredcos_cap_v2"
)

# 创建了一个AdamW优化器，用于更新模型的参数。返回模型的可训练参数，参数指定了学习率。
optimizer = torch.optim.AdamW(model.parameters(), lr=4e-4)

# 创建一个空列表，用于存储每个步骤的损失值。
losses = []

for epoch in range(30):
  for step, batch in enumerate(train_dataloader):
    # 从批次中获取干净的图像数据，并将其移动到指定的设备（例如GPU）上进行计算。
    clean_images = batch["images"].to(device)
    # 生成与干净图像相同形状的噪声张量，该噪声将被添加到图像中。
    noise = torch.randn(clean_images.shape).to(clean_images.device)
    # 获取批次的大小。
    bs = clean_images.shape[0]

    # 为每个图像随机生成一个时间步长，该时间步长将用于确定噪声的变化程度。
    timesteps = torch.randint(
      0, noise_scheduler.num_train_timesteps, (bs,), device=clean_images.device
    ).long()

    # 根据噪声调度器中的每个噪声幅度和时间步长，将噪声添加到干净图像中，生成带有噪声的图像。
    noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)

    # 使用模型对带有噪声的图像进行预测，得到去噪后的图像。
    noise_pred = model(noisy_images, timesteps, return_dict=False)[0]

    # 计算预测图像与真实噪声之间的均方误差损失，计算损失相对于模型参数的梯度，并将当前步骤的损失添加到列表中。
    loss = F.mse_loss(noise_pred, noise)
    loss.backward(loss)
    losses.append(loss.item())

    # 使用优化器更新模型的参数，并将梯度置零。
    optimizer.step()
    optimizer.zero_grad()

  # 每隔5个epoch，计算最近一个epoch的平均损失，并打印出来。
  if (epoch + 1) % 5 == 0:
      loss_last_epoch = sum(losses[-len(train_dataloader) :]) / len(train_dataloader)
      print(f"Epoch:{epoch+1}, loss: {loss_last_epoch}")

# 绘制 loss 曲线，我们能看到模型在一开始快速的收敛，接下来以一个较慢的速度持续优化（我们用右边 log 坐标轴的视图可以看的更清楚）：

fig, axs = plt.subplots(1, 2, figsize=(12, 4))
axs[0].plot(losses)
axs[1].plot(np.log(losses))
plt.show()

五、生成图像 

下面开始生成图像 

# 方法 1：建立一个管道：
from diffusers import DDPMPipeline

image_pipe = DDPMPipeline(unet=model, scheduler=noise_scheduler)
pipeline_output = image_pipe()
pipeline_output.images[0]

# 我们可以在本地文件夹这样保存一个管道：
image_pipe.save_pretrained("my_pipeline")

# 检查文件夹的内容：
!ls my_pipeline/

 这里scheduler与unet子文件夹中包含了生成图像所需的全部组件。比如，在unet文件中能看到模型参数 (diffusion_pytorch_model.bin) 与描述模型结构的配置文件。

# 方法 2：写一个取样循环

# 从随机噪声开始，遍历管理器的迭代周期来看从最嘈杂直到最微小的噪声变化，基于模型的预测一步步减少一些噪声：

# Random starting point (8 random images):
sample = torch.randn(8, 3, 32, 32).to(device)

for i, t in enumerate(noise_scheduler.timesteps):

    # Get model pred
    with torch.no_grad():
        residual = model(sample, t).sample

    # Update sample with step
    sample = noise_scheduler.step(residual, t, sample).prev_sample
    
# noise_scheduler.step () 函数相应做了 sample（取样）时的数学运算。
show_images(sample)

 六、将模型上传到Hugging Face

在上面的例子中我们把管道保存在了本地。 把模型 push 到 hub 上，我们会需要建立模型和相应文件的仓库名。 我们根据你的选择（模型 ID）来决定仓库的名字（大胆的去替换掉model_name吧；需要包含你的用户名，get_full_repo_name ()会帮你做到）：

from huggingface_hub import get_full_repo_name

model_name = "sd-class-butterflies-32"
hub_model_id = get_full_repo_name(model_name)
hub_model_id

# 然后，在 🤗 Hub 上创建模型仓库并 push 它吧：

from huggingface_hub import HfApi, create_repo

create_repo(hub_model_id)
api = HfApi()
api.upload_folder(
    folder_path="my_pipeline/scheduler", path_in_repo="", repo_id=hub_model_id
)
api.upload_folder(folder_path="my_pipeline/unet", path_in_repo="", repo_id=hub_model_id)
api.upload_file(
    path_or_fileobj="my_pipeline/model_index.json",
    path_in_repo="model_index.json",
    repo_id=hub_model_id,
)

# 最后一件事是创建一个超棒的模型卡，如此，我们的蝴蝶生成器可以轻松的在 Hub 上被找到（请在描述中随意发挥！）：

from huggingface_hub import ModelCard

content = f"""
---
license: mit
tags:
- pytorch
- diffusers
- unconditional-image-generation
- diffusion-models-class
---

# Model Card for Unit 1 of the [Diffusion Models Class 🧨](https://github.com/huggingface/diffusion-models-class)

This model is a diffusion model for unconditional image generation of cute 🦋.

## Usage

```python
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)

# 现在模型已经在 Hub 上了，你可以这样从任何地方使用DDPMPipeline的from_pretrained ()方法来下来它：
from diffusers import DDPMPipeline

image_pipe = DDPMPipeline.from_pretrained(hub_model_id)
pipeline_output = image_pipe()
pipeline_output.images[0]