DataWhale夏令营第四期魔搭- AIGC方向 task01笔记

30分钟体验一站式baseline

 1、下载baseline文件

终端输入：

git lfs install

git clone https://www.modelscope.cn/datasets/maochase/kolors.git

2、环境安装

安装 Data-Juicer 和 DiffSynth-Studio（安装完成后，重启Notebook kernel）

Data-Juicer：数据处理和转换工具，旨在简化数据的提取、转换和加载过程

DiffSynth-Studio：高效微调训练大模型工具

安装代码如下：

!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

3、下载数据集

使用 ModelScope 库中的 MsDataset 类加载数据集的训练集部分，并将数据集缓存到指定的目录中 /mnt/workspace/kolors/data。接下来处理数据集 `ds` 中的图像数据，将图像保存到指定目录，并将相关的元数据（包括文本描述和图像路径）以 JSON 格式写入到 `metadata.jsonl` 文件中

from modelscope.msdatasets import MsDataset

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

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")

4、数据处理

使用 data-juicer 处理数据

保存处理好的数据

5、训练模型

下载模型

from diffsynth import download_models

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

开始训练：

通过操作系统调用执行一个命令行命令，从而启动一个Python脚本 train_kolors_lora.py，并通过命令行参数配置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)

加载模型：

加载和配置多个深度学习模型，包括文本编码器、diffusion模型和VAE模型。使用 load_lora 函数加载LoRA适配器到UNet模型中，以增强模型的性能和适应性。

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"
)

生成图像：

prompt为正面描述，negative_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")

6、小结

虽然在学校时候有接触过深度学习的相关知识，但因为了解不深的缘故仍算一个小白，这一次作业也是，基本全程跟着教程走，对AI生成图像的流程有了一个基本了解，但对代码的理解还不够，希望仍借着接下来的学习加强自己的能力。