Datawhale AI夏令营第四期魔搭- AIGC方向 task02笔记

大纲

一、前言

本文主要是针对baseline代码进行逐行解读，当然啦，是在大模型的帮助下进行尽可能详细的解读，以帮助我学习理解baseline代码。

二、代码块解读

整个代码的逻辑如下图

图片来源：DataWhale

2.1 依赖包安装

!pip install simple-aesthetics-predictor
!pip install -v -e data-juicer
!pip uninstall pytorch-lightning -y
!pip install peft lightning pandas torchvision
!pip install -e DiffSynth-Studio

安装依赖包及卸载依赖包
其中：-v 选项意味着 verbose 输出，即显示详细的安装信息。
-e 选项表示以可编辑模式安装包，这意味着安装的包可以从指定的位置被直接修改，而不需要重新安装。
关于

data-juicer 是一个用于数据清洗和质量评估的工具。它可以帮助您检查和清理数据集中的各种问题，比如重复项、噪声、格式错误等。data-juicer 的主要作用包括但不限于以下几点：
数据质量评估：
可以评估数据集中每个样本的质量，例如通过检测样本是否包含足够的信息、是否存在异常值、是否与其他样本重复等。
可以生成报告，帮助您了解数据集的整体健康状况。
数据清洗：
data-juicer 可以自动识别并修复数据集中的一些常见问题，比如文本格式不一致、图像分辨率过低等。
可以帮助您删除或标记有问题的数据点。
数据增强：
对于一些任务，data-juicer 可能还提供了数据增强的功能，以便生成更多样化的训练数据。
数据预处理：
在数据进入模型之前，data-juicer 可以帮助您对数据进行标准化、归一化等预处理操作。
可定制性：
data-juicer 允许用户根据自己的需求定义特定的数据质量指标和清洗策略。
在安装 data-juicer 并以可编辑模式 (-e) 安装后，您可以直接从源代码目录运行和调试 data-juicer，这使得在开发过程中更容易进行修改和测试。

2.2 数据集下载

from modelscope.msdatasets import MsDataset

ds = MsDataset.load(
'AI-ModelScope/lowres_anime',
subset_name='default',
split='train',
cache_dir="/mnt/workspace/kolors/data"
)

这段代码是从 ModelScope 框架中加载数据集的一个示例。ModelScope 是阿里云推出的一个模型开放平台，它不仅提供了大量的预训练模型，还支持数据集管理。
首先通过from…import…导入 ModelScope 数据集模块中的 MsDataset 类。MsDataset 类可以用来加载 ModelScope 中的数据集。
接着创建了一个 MsDataset 实例并加载了一个名为 AI-ModelScope/lowres_anime 的数据集。加载完成之后，ds 变量将包含加载的数据集。可以通过这个变量访问数据集中的各项属性和方法。例如，可以遍历数据集中的每一个样本，或者获取数据集中样本的数量等。

各参数详解：

‘AI-ModelScope/lowres_anime’: 这是指定要加载的数据集的名字。通常在 ModelScope中，数据集的名字由一个唯一的标识符来表示，形如 <组织名>/<数据集名>。在这个例子中，数据集名为 lowres_anime，而组织名为AI-ModelScope。
subset_name=‘default’:指定了数据集中的子集名称。有些数据集可能会被分成不同的子集，例如不同的版本或不同领域的数据。这里使用的是默认子集。
split=‘train’: 指定了要加载的数据集分割部分。通常数据集会被分为训练集（train）、验证集（validation）和测试集（test）。这里加载的是训练集部分。
cache_dir=“/mnt/workspace/kolors/data”: 指定了数据集下载后存储的本地目录。这样可以避免重复下载相同的数据集，并且可以加速数据集的加载速度。

2.3 metadata.jsonl文件生成（图片及对应标签）

import json, os
from data_juicer.utils.mm_utils import SpecialTokens
from tqdm import tqdm

os.makedirs("./data/lora_dataset/train", exist_ok=True)
os.makedirs("./data/data-juicer/input", exist_ok=True)
with open("./data/data-juicer/input/metadata.jsonl", "w") as f:
	for data_id, data in enumerate(tqdm(ds)):
		image = data["image"].convert("RGB")
		image.save(f"/mnt/workspace/kolors/data/lora_dataset/train/{data_id}.jpg")
		metadata = {"text": "二次元", "image": [f"/mnt/workspace/kolors/data/lora_dataset/train/{data_id}.jpg"]}
		f.write(json.dumps(metadata))
		f.write("\n")

这段代码主要用于从前面加载的数据集 ds 中提取图像，并为每张图像创建相应的元数据文件。
首先导入了几个必要的模块：

json: 用于处理 JSON 格式的数据。
os: 提供了与操作系统交互的功能，如创建目录。
data_juicer.utils.mm_utils.SpecialTokens: 从 data-juicer 库中导入了SpecialTokens 类
tqdm: 一个进度条库，用于显示循环进度。

接着用os.makedirs创建了两个目录：

“./data/lora_dataset/train”: 用于存放从数据集中提取出的图像。
“./data/data-juicer/input”: 用于存放元数据文件。

使用 with 语句以写入模式 (“w”) 打开一个名为 metadata.jsonl 的文件，jsonl 文件是一种常见的存储 JSON 对象列表的文件格式，其中每个对象都独占一行。
在for语句中，使用 enumerate 函数遍历数据集 ds，并使用 tqdm 显示一个进度条。data_id 是当前数据点的索引，data 是对应的数据点。
通过data["image’]提取数据点 data 中提取图像，并用convert将其转换为 RGB 格式。这是因为某些图像可能是其他格式，例如灰度图或包含 alpha 通道的图像。image.save()将图像保存到 “./data/lora_dataset/train” 目录下，并命名为 {data_id}.jpg。
接着创建了一个字典 metadata，其中包含两个键：

“text”: 一个字符串 “二次元”，可能是描述图像内容的标签。
“image”: 一个列表，包含一个字符串，即图像的路径。

采用json.dumps()将 metadata 字典序列化为 JSON 格式，并写入到 metadata.jsonl 文件中。每条记录后面追加一个换行符，以便每个 JSON 对象独立成行。

2.4 设置data-juicer 配置文件并执行，处理metadata.jsonl文件生成result.jsonl文件

data_juicer_config = """

#global parameters
project_name: 'data-process'
dataset_path: './data/data-juicer/input/metadata.jsonl'  # path to your dataset directory or file
np: 4  # number of subprocess to process your dataset

text_keys: 'text'
image_key: 'image'
image_special_token: '<__dj__image>'

export_path: './data/data-juicer/output/result.jsonl'

#process schedule
#a list of several process operators with their arguments
process:
	- image_shape_filter:
		min_width: 1024
		min_height: 1024
		any_or_all: any
	- image_aspect_ratio_filter:
		min_ratio: 0.5
		max_ratio: 2.0
		any_or_all: any
"""
with open("data/data-juicer/data_juicer_config.yaml", "w") as file:
file.write(data_juicer_config.strip())

!dj-process --config data/data-juicer/data_juicer_config.yaml

这段代码定义了一个 YAML格式的字符串data_juicer_config作为data-juicer 的配置文件，并使用该配置文件来处理之前创建的 metadata.jsonl 文件。

以下是 YAML 格式的字符串 data_juicer_config各参数的信息：

全局参数：
project_name: 项目的名称。
dataset_path: 指定要处理的数据集路径。
np: 处理数据集时使用的子进程数量。
text_keys: 文本字段的键名。
image_key: 图像字段的键名。
image_special_token: 特殊标记，用于表示图像的存在。
导出路径： export_path: 处理后的结果输出路径。
处理流程：
process: 一系列数据处理操作的定义。
image_shape_filter: 过滤不符合尺寸要求的图像。
min_width: 最小宽度。
min_height: 最小高度。
any_or_all: 如果是any，则只要满足其中一个条件即可；如果是 all，则需要同时满足所有条件。 image_aspect_ratio_filter: 过滤不符合宽高比要求的图像。
min_ratio: 最小宽高比。
max_ratio: 最大宽高比。
any_or_all: 同上。

接着将 data_juicer_config 写入到一个名为 data_juicer_config.yaml 的文件中。使用 strip() 去除字符串开头和结尾的空白字符。
最后使用 data-juicer 的命令行工具根据配置文件来处理数据集并将处理结果输出到指定的文件result.jsonl。

2.5 基于result.jsonl文件生成训练集图片及文件

import pandas as pd
import os, json
from PIL import Image
from tqdm import tqdm

texts, file_names = [], []
os.makedirs("./data/lora_dataset_processed/train", exist_ok=True)
with open("./data/data-juicer/output/result.jsonl", "r") as file:
    for data_id, data in enumerate(tqdm(file.readlines())):
        data = json.loads(data)
        text = data["text"]
        texts.append(text)
        image = Image.open(data["image"][0])
        image_path = f"./data/lora_dataset_processed/train/{data_id}.jpg"
        image.save(image_path)
        file_names.append(f"{data_id}.jpg")
data_frame = pd.DataFrame()
data_frame["file_name"] = file_names
data_frame["text"] = texts
data_frame.to_csv("./data/lora_dataset_processed/train/metadata.csv", index=False, encoding="utf-8-sig")
data_frame

用于读取经过 data-juicer 处理后的数据result.jsonl，并将其整理成 Pandas DataFrame，最终保存为 CSV 文件metadata.csv。

2.6 模型训练

2.6.1 模型加载

from diffsynth import download_models
download_models(["Kolors", "SDXL-vae-fp16-fix"])

使用 diffsynth 库中的 download_models 函数来下载指定的模型（“Kolors” 和 “SDXL-vae-fp16-fix”）。

!python DiffSynth-Studio/examples/train/kolors/train_kolors_lora.py -h

这段命令行代码使用 DiffSynth-Studio 中的一个脚本来展示其帮助信息。下面是逐行解析：

!python: 这是在 Jupyter Notebook 或者类似的环境中执行 shell 命令的方式，它会调用 Python
解释器来运行命令。
DiffSynth-Studio/examples/train/kolors/train_kolors_lora.py: 这是脚本的路径，它位于 DiffSynth-Studio 项目的 examples/train/kolors 目录下。
-h: 这个选项通常用于显示命令的帮助信息。

当执行这条命令时，Python 解释器会运行 train_kolors_lora.py 脚本，并因为 -h 参数的存在，脚本会输出其帮助信息，而不是执行实际的训练过程。

2.6.2 LoRA训练参数设置及训练

import os

cmd = """
python DiffSynth-Studio/examples/train/kolors/train_kolors_lora.py \
  --pretrained_unet_path models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors \
  --pretrained_text_encoder_path models/kolors/Kolors/text_encoder \
  --pretrained_fp16_vae_path models/sdxl-vae-fp16-fix/diffusion_pytorch_model.safetensors \
  --lora_rank 16 \
  --lora_alpha 4.0 \
  --dataset_path data/lora_dataset_processed \
  --output_path ./models \
  --max_epochs 1 \
  --center_crop \
  --use_gradient_checkpointing \
  --precision "16-mixed" 
""".strip()

os.system(cmd)

这段代码首先定义了一个字符串变量 cmd，cmd中包含了一个 shell 命令，用于运行 DiffSynth-Studio 中的训练脚本。使用 \ 符号进行换行，以保持命令的可读性。.strip() 方法用于移除字符串首尾的空白字符。通过os.system(cmd)来运行cmd字符串的shell命令。
主要作用是运行 DiffSynth-Studio 中的一个脚本，用于训练 LoRA 模型。它使用了预训练的 U-Net、文本编码器和 VAE 模型，并且指定了训练数据集的路径以及输出路径。此外，还设置了一些训练参数，如最大轮数、中心裁剪、梯度检查点和精度模式。

命令参数解析

python DiffSynth-Studio/examples/train/kolors/train_kolors_lora.py: 运行DiffSynth-Studio 中名为 train_kolors_lora.py 的训练脚本。
–pretrained_unet_path models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors: 指定预训练的U-Net 模型路径。
–pretrained_text_encoder_path models/kolors/Kolors/text_encoder: 指定预训练的文本编码器路径。
–pretrained_fp16_vae_path models/sdxl-vae-fp16-fix/diffusion_pytorch_model.safetensors: 指定预训练的VAE 模型路径，使用 FP16 精度。
–lora_rank 16: 设置 LoRA 适应的秩。
–lora_alpha 4.0: 设置 LoRA 的缩放因子。
–dataset_path data/lora_dataset_processed: 指定训练数据集的路径。
–output_path ./models: 指定训练结果的输出路径。
–max_epochs 1: 设置训练的最大轮数。
–center_crop: 启用中心裁剪。
–use_gradient_checkpointing: 启用梯度检查点，以减少内存占用。
–precision “16-mixed”: 设置精度模式为混合精度训练（FP16 和 FP32 混合使用）。

2.7 将训练好的LoRA加载到模型中，图片生成

from diffsynth import ModelManager, SDXLImagePipeline
from peft import LoraConfig, inject_adapter_in_model
import torch

def load_lora(model, lora_rank, lora_alpha, lora_path):
    lora_config = LoraConfig(
        r=lora_rank,
        lora_alpha=lora_alpha,
        init_lora_weights="gaussian",
        target_modules=["to_q", "to_k", "to_v", "to_out"],
    )
    model = inject_adapter_in_model(lora_config, model)
    state_dict = torch.load(lora_path, map_location="cpu")
    model.load_state_dict(state_dict, strict=False)
    return model

# Load models
model_manager = ModelManager(torch_dtype=torch.float16, device="cuda",
                             file_path_list=[
                                 "models/kolors/Kolors/text_encoder",
                                 "models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors",
                                 "models/kolors/Kolors/vae/diffusion_pytorch_model.safetensors"
                             ])
pipe = SDXLImagePipeline.from_model_manager(model_manager)

# Load LoRA
pipe.unet = load_lora(
    pipe.unet,
    lora_rank=16, # This parameter should be consistent with that in your training script.
    lora_alpha=2.0, # lora_alpha can control the weight of LoRA.
    lora_path="models/lightning_logs/version_0/checkpoints/epoch=0-step=500.ckpt"
)

这段代码用于加载预训练的模型，并应用 LoRA（Low-Rank Adaptation）适配器到 U-Net 模型中。
以下逐行解释代码：

2.7.1 导入加载所需依赖包

from diffsynth import ModelManager, SDXLImagePipeline
from peft import LoraConfig, inject_adapter_in_model
import torch

这里导入了必要的模块：

diffsynth.ModelManager: 用于管理模型的类。
diffsynth.SDXLImagePipeline: 用于图像生成的管道。
peft.LoraConfig: 用于配置 LoRA 适配器的类。
peft.inject_adapter_in_model: 用于注入 LoRA 适配器到模型中的函数。
torch: PyTorch库，用于深度学习。

2.7.2 定义LoRA加载函数

def load_lora(model, lora_rank, lora_alpha, lora_path):

这里定义了一个名为 load_lora 的函数，用于加载 LoRA 适配器到给定的模型中。

    lora_config = LoraConfig(
        r=lora_rank,
        lora_alpha=lora_alpha,
        init_lora_weights="gaussian",
        target_modules=["to_q", "to_k", "to_v", "to_out"],
    )

这行代码创建了一个 LoraConfig 对象，指定了 LoRA 适配器的配置：

r: LoRA 的秩。
lora_alpha: LoRA 的缩放因子。
init_lora_weights: 初始化 LoRA 权重的方法。
target_modules: 需要注入 LoRA 适配器的模块名称。

    model = inject_adapter_in_model(lora_config, model)

这行代码将 LoRA 适配器注入到模型中。

    state_dict = torch.load(lora_path, map_location="cpu")
    model.load_state_dict(state_dict, strict=False)

这里加载了 LoRA 适配器的状态字典，并将其加载到模型中。map_location=“cpu” 表示状态字典应该加载到 CPU 上，但如果模型已经在 GPU 上，则会自动加载到 GPU 上。

    return model

返回加载了 LoRA 适配器的模型。

2.7.3 创建ModelManager 实例，并加载文本编码器、 U-Net 模型、VAE 模型

# Load models
model_manager = ModelManager(torch_dtype=torch.float16, device="cuda",
                             file_path_list=[
                                 "models/kolors/Kolors/text_encoder",
                                 "models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors",
                                 "models/kolors/Kolors/vae/diffusion_pytorch_model.safetensors"
                             ])

这里创建了一个 ModelManager 实例，并加载了三个模型：

文本编码器、 U-Net 模型、VAE 模型。

模型的类型和设备被指定为 torch.float16 和 cuda，这意味着模型将以半精度浮点格式加载到 CUDA 设备上。

2.7.4 创建SDXLImagePipeline 实例，并使用 ModelManager 加载的模型来初始化管道

pipe = SDXLImagePipeline.from_model_manager(model_manager)

这行代码创建了一个 SDXLImagePipeline 实例，并使用 ModelManager 加载的模型来初始化管道。

# Load LoRA
pipe.unet = load_lora(
    pipe.unet,
    lora_rank=16, # This parameter should be consistent with that in your training script.
    lora_alpha=2.0, # lora_alpha can control the weight of LoRA.
    lora_path="models/lightning_logs/version_0/checkpoints/epoch=0-step=500.ckpt"
)

这里将 LoRA 适配器加载到 U-Net 模型中：

lora_rank: LoRA 的秩，应与训练脚本中的一致。
lora_alpha: 控制 LoRA 的权重。
lora_path: LoRA 适配器的路径。

总结一下，这段代码的主要作用是：

加载预训练的文本编码器、U-Net 和 VAE 模型。
创建一个 SDXLImagePipeline 实例。
将 LoRA 适配器加载到 U-Net 模型中。

2.7.5 基于prompt生成图片

torch.manual_seed(0)
image = pipe(
prompt="二次元，一个紫色短发小女孩，在家中沙发上坐着，双手托着腮，很无聊，全身，粉色连衣裙",
negative_prompt="丑陋、变形、嘈杂、模糊、低对比度",
cfg_scale=4,
num_inference_steps=50, height=1024, width=1024,
)
image.save("1.jpg")

在这段代码中，您首先设置了随机种子以确保实验的可重复性，并随后使用 SDXLImagePipeline 管道生成了一张图像。这里是代码的详细解释：

torch.manual_seed(0)

设置 PyTorch 的随机种子为 0，以确保生成的结果在每次运行时都相同。

image = pipe(
    prompt="二次元，一个紫色短发小女孩，在家中沙发上坐着，双手托着腮，很无聊，全身，粉色连衣裙",
    negative_prompt="丑陋、变形、嘈杂、模糊、低对比度",
    cfg_scale=4,
    num_inference_steps=50, height=1024, width=1024,
)

使用 pipe 生成一张图像，其中：

prompt 是生成图像的正面描述。
negative_prompt 是需要避免的特征或风格。
cfg_scale (Classifier-Free Guidance Scale) 是指导强度，较高的值会让图像更贴近提示内容。
num_inference_steps 是扩散过程中的步数，更多的步骤通常会产生更高质量的图像。 height 和 width 设置输出图像的尺寸。

image.save("1.jpg")

将生成的图像保存为文件 “1.jpg”。

2.8 图片生成展示

mport numpy as np
from PIL import Image

images = [np.array(Image.open(f"{i}.jpg")) for i in range(1, 9)]
image = np.concatenate([
np.concatenate(images[0:2], axis=1),
np.concatenate(images[2:4], axis=1),
np.concatenate(images[4:6], axis=1),
np.concatenate(images[6:8], axis=1),
], axis=0)
image = Image.fromarray(image).resize((1024, 2048))
image

展示了如果八张图片按照 2x4 的网格布局排列并调整到 1024x2048 像素时，最终图像的外观。请注意，这只是一个占位符，实际图像的内容取决于文件 1.jpg 到 8.jpg 中的内容。

2.9 复制训练好的LoRA文件和其中一张生成的图片到output文件夹

mkdir /mnt/workspace/kolors/output & cd
cp /mnt/workspace/kolors/models/lightning_logs/version_0/checkpoints/epoch=0-step=500.ckpt /mnt/workspace/kolors/output/
cp /mnt/workspace/kolors/1.jpg /mnt/workspace/kolors/output/

1、mkdir /mnt/workspace/kolors/output & cd: 创建一个名为 output 的目录，位于 /mnt/workspace/kolors/ 下，之后切换至刚创建的output 目录下。
2、cp /mnt/workspace/kolors/models/lightning_logs/version_0/checkpoints/epoch=0-step=500.ckpt /mnt/workspace/kolors/output/: 将LoRA文件 epoch=0-step=500.ckpt 复制到 output 目录下。
3、cp /mnt/workspace/kolors/1.jpg /mnt/workspace/kolors/output/: 将图片文件 1.jpg 复制到 output 目录下。

总结