DataWhale AI夏令营-从零入门AI生图原理&实践
文生图(text-to-image)简介
文生图有很多模型架构其实,比如GAN、Flow-based等,这里我主要介绍两种比较能打的分别是DALL·E和Diffusion。其中本次DataWhale的学习也是围绕Diffusion模型开展的,很喜欢一句话在这里引用一下" 高斯去噪,扩散称王 。"
DALL·E
文生图初期是由OpenAI开发的DALL·E模型,其实他的原理比较符合直接接触自然语言处理方向转AIGC的同学。
他的核心原理主要分为两部分:
- VAE(Variational AutoEncoder):正如他的名字,就像大多数自然语言处理任务那样,我需要先预训练得到一个词表,VAE的思想也是,压缩图像把图像压缩成image token(其实这部分大家可以去了解一下,我最开始接触是图像加密,要把图像编码,enmmmm,可能又不大一样,大家感兴趣可以自己去了解一下,我也是一知半解)。
- 把文本片段(text)和image token拼接到一起,然后训练一个自回归的Transformer,让模型学习图像和文本间的联合分布。
其实DALL·E的思路基本和我们做的一些nlp方向,大模型方向的思路比较接近,但是我曾经用过OpenAI的文生图感觉他的效果,ennnmmm,感觉生成的东西很抽象,很难评,可能像gpt那样需要不断的迭代训练参数到了一定的量级之后效果才能显著的提升吧。
Diffusion
扩散模型(Diffusion),说到扩散模型,其实最著名的还是Stable Diffusion,基本提到AIGC就不得不提Stable Diffusion,虽然这次竞赛并没有使用Stable Diffusion进行生图但是也建议后续大家可以自己玩玩这个模型,Stable Diffusion的体系非常全面,尤其是他的Web页面用起来很方便。
这里我就不多啰嗦了,Stable Diffusion模型资料网上非常多。
继续说我们的扩散模型基础:
-
正向图像扩散(forward image diffusion):学习过图像处理的同学都应该直到,给一张图片不断的添加噪声,逐渐的图片就会完全无法识别。
如图左图为原始图像(灰度),右图为加了噪声的图像。
-
反向图像扩散(reverse image diffusion):
其实根据正向过程,大家可以才想到,所有图像一直加噪声,最后不都应该变成一对类似像素点的图片吗,那我要是反过来,不就生成原图了吗。其实从概念上是不行的,而且你还要计算图像的噪声,这有点太高难度了。
不过,虽然无法精确计算,但可以近似! 核心思想是以可学习网络(learnable network)的形式, 近似反向扩散过程。
接下来就不在给大家细讲了,后面再说就都是公式了,我这里其实也是一知半解,更多的是把我的理解给写出来吧,毕竟我也没能力做算法,感兴趣可以去看一看论文、资料,这里我也推荐一篇博客,讲的很好译(文生图)简史
Baseline
这里只对代码进行注释,搭建环境的过程详见DataWhale文档。
下载需要的包
!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
下载数据集,这里也可以去github或者modelscope上找一些其他的数据集,当然赛事使用的数据集质量就很高了其实。
from modelscope.msdatasets import MsDataset
# 使用modelscope加载数据集
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 并将其转换为 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")
配置 data-juicer 的参数和处理流程
project_name: 项目名称
dataset_path: 数据集的路径
np: 并行处理时使用的子进程数量
text_keys: 指定文本字段的键名
image_key: 指定图像字段的键名
image_special_token: 用于表示图像的特殊标记
export_path: 处理后结果的导出路径
process: 这是一个处理流程的列表,每个元素表示一个处理步骤和其对应的参数。
image_shape_filter: 过滤图像形状,要求图像的宽度和高度至少为 1024 像素。
image_aspect_ratio_filter: 过滤图像的宽高比,要求宽高比在 0.5 到 2.0 之间。
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
保存处理好的数据
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
下载模型(准备开始train)
from diffsynth import download_models
download_models(["Kolors", "SDXL-vae-fp16-fix"])
开始训练:
python DiffSynth-Studio/examples/train/kolors/train_kolors_lora.py: 运行指定路径下的 train_kolors_lora.py 脚本,进行模型训练。
–pretrained_unet_path models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors: 指定预训练的 UNet 模型路径。
–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 的 rank 值,这个参数影响LoRA的压缩程度(Lora实际就是训练一个低秩矩阵,来影响原模型)。
–lora_alpha 4.0: 设置 LoRA 的 alpha 参数,影响LoRA训练的平衡。
–dataset_path data/lora_dataset_processed: 指定用于训练的数据集路径。
–output_path ./models: 指定训练结果保存的路径。
–max_epochs 1: 设置训练的最大 epoch 数为 1。
–center_crop: 使用中心裁剪方法处理图像。
–use_gradient_checkpointing: 启用梯度检查点,节省显存。
–precision “16-mixed”: 使用混合精度训练(FP16),提高训练效率并减少显存使用。
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)
加载合并Lora后的模型:
Lora参数:
lora_rank: LoRA 的秩,控制低秩近似的程度。
lora_alpha: LoRA 的缩放系数,影响 LoRA 的训练和应用效果。
init_lora_weights=“gaussian”: 初始化 LoRA 权重为高斯分布。
target_modules: 目标模块,指定 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"
)
一共是八张图片,这里我就偷个懒,只展示一张了:
prompt:正向提示词。
negative_prompt:反向提示词。
cfg_scale:模型在生成图像时对 prompt 的依赖程度。较高的值会更严格地按照提示生成图像,这个参数的范围一般是1-20(低强度:1-5;中等强度:5-10;高强度:10-20)大家可以修改体验一下。
num_inference_steps: 推理步骤的数量,代表生成图像时模型迭代的次数。更多的步骤通常可以带来更高质量的图像。
height 和 width: 设置生成图像的高度和宽度,单位为像素。
torch.manual_seed(7)
image = pipe(
prompt="二次元,一个白衣年轻修仙者,站在瀑布下方,身后是庞大的远古神兽虚影,他正在突破境界,身旁灵气形成的漩涡闪烁着光芒",
negative_prompt="丑陋、变形、嘈杂、模糊、低对比度",
cfg_scale=4,
num_inference_steps=50, height=1024, width=1024,
)
image.save("1.jpg")
展示一下我的图片吧,挑了张最好的,哈哈
把八张图片拼到一起
import 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
来点小技巧吧
好东西往往都放在最后,哈哈哈,这里给大家一个建议吧。我做大模型相关应用多一些(瞎玩),建议大家可以让大模型给你生成提示词,你可以把baseline的八条prompt给大模型然后给大模型一个场景让他给你按照baseline的格式输出结果,你再把大模型生成的prompt按照你想要的修改一下,能省不少时间,不然你自己闷头写也不好描述。
这是我写给的大模型的一个input,供给大家思考,最好别直接用我这个,不然生成图片之后你就会后悔了。
我现在在做一个文生图的大模型,可以生成八张连贯的图片。我想以:一个黑人篮球运动员科比,科比登上飞机,飞机出事故经历颠簸,醒来发现自己已经在天堂了,上帝给他介绍天堂让他加入天堂篮球队,在天堂和上帝打篮球,拿到天堂的冠军,在夜里拿着妻子孩子的合照怀念。这些为提示词来生成图片。请你模仿下面的提示词来帮我撰写这八张图片的提示词。
二次元,一个紫色短发小女孩,在家中沙发上坐着,双手托着腮,很无聊,全身,粉色连衣裙
二次元,日系动漫,演唱会的观众席,人山人海,一个紫色短发小女孩穿着粉色吊带漏肩连衣裙坐在演唱会的观众席,舞台上衣着华丽的歌星们在唱歌
二次元,一个紫色短发小女孩穿着粉色吊带漏肩连衣裙坐在演唱会的观众席,露出憧憬的神情
二次元,一个紫色短发小女孩穿着粉色吊带漏肩连衣裙,对着流星许愿,闭着眼睛,十指交叉,侧面
二次元,一个紫色中等长度头发小女孩穿着粉色吊带漏肩连衣裙,在练习室练习唱歌
二次元,一个紫色长发小女孩穿着粉色吊带漏肩连衣裙,在练习室练习唱歌,手持话筒
二次元,紫色长发少女,穿着黑色连衣裙,试衣间,心情忐忑
二次元,紫色长发少女,穿着黑色礼服,连衣裙,在台上唱歌