diff --git a/README.md b/README.md index 641ccbe..dec058c 100644 --- a/README.md +++ b/README.md @@ -69,11 +69,23 @@ python tools/diffusers2ckpt.py ./new_model ./ckpt_models/newModel_half.ckpt --ha 我在tools/handle_images.py中提供了一份批量处理的代码用于参考 自动center crop图像,并缩放尺寸 ``` -python tools/handle_images.py ./datasets/test ./datasets/test2 --width=512 --height=512 +!python tools/handle_images.py --origin_image_path ./datasets/test \ + --output_image_path ./datasets/test2 \ + --width=512 --height=512 ``` test为未处理的原始图像文件夹,test2为输出处理图像的路径 如需处理透明背景png图为黑色/白色底jpg,可以添加--png参数。 +### 图像裁剪为3:2 && 2:3的矩形 +不需要修改width和height(输出自动分成两个文件夹) +``` +!python tools/handle_images.py --origin_image_path ./datasets/test \ + --output_image_path_0 ./datasets/a1 --output_image_path_1 ./datasets/a2 \ + --width=768 --height=512 +``` +test为未处理的原始图像文件夹, +a1,a2为裁剪后输出的两个文件夹,分别保存768x512和512x768的图像。 + ### 图像自动标注 使用deepdanbooru生成tags label. ``` @@ -107,12 +119,14 @@ accelerate config ``` sh train_object.sh +# sh train_object_rect.sh --width 768 --height 512 #输入为矩形图像 ``` 如果要Finetune训练自己的大模型: (推荐准备3000+张图片,包含尽可能的多样性,数据决定训练出的模型质量) ``` sh train_style.sh +# sh train_style_rect.sh --width 768 --height 512 #输入为矩形图像 ``` A5000的训练速度大概8分钟/1000步 diff --git a/tools/scale_images.py b/tools/scale_images.py new file mode 100644 index 0000000..60fbc2e --- /dev/null +++ b/tools/scale_images.py @@ -0,0 +1,39 @@ +from PIL import Image +import os +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument("--origin_image_path", default=None, type=str, help="Path to the images to convert.") +parser.add_argument("--output_image_path", default=None, type=str, help="Path to the output images.") +parser.add_argument("--max_side_length", default=None, type=int, help="Path to the images to convert.") +args = parser.parse_args() + + +def resize_images(folder_path,output_folder_path, max_side_length): + for filename in os.listdir(folder_path): + if filename.endswith(('.jpg', '.png')): + image_path = os.path.join(folder_path, filename) + output_path = os.path.join(output_folder_path, filename) + + image = Image.open(image_path) + image = image.convert('RGB') + + width, height = image.size + + scale = min(max_side_length / width, max_side_length / height) + + new_width = int(width * scale) + new_height = int(height * scale) + + resized_image = image.resize((new_width, new_height), Image.ANTIALIAS) + + resized_image.save(output_path) + +input_path = args.origin_image_path +output_path = args.output_image_path +if output_path!=None: + if not os.path.exists(output_path): + os.makedirs(output_path) + +max_side_length = args.max_side_length # 最大边长 24G <= 786 +resize_images(input_path,output_path, max_side_length) diff --git a/train_object_all.sh b/train_object_all.sh new file mode 100644 index 0000000..01134a1 --- /dev/null +++ b/train_object_all.sh @@ -0,0 +1,77 @@ +# 用于训练特定物体/人物的方法(只需单一标签) +export MODEL_NAME="./model" +export INSTANCE_DIR="./datasets/test2" +export OUTPUT_DIR="./new_model" +export CLASS_DIR="./datasets/class" # 用于存放模型生成的先验知识的图片文件夹,请勿改动 +export LOG_DIR="/root/tf-logs" +export TEST_PROMPTS_FILE="./test_prompts_object.txt" + +rm -rf $CLASS_DIR/* # 如果你要训练与上次不同的特定物体/人物,需要先清空该文件夹。其他时候可以注释掉这一行(前面加#) +rm -rf $LOG_DIR/* + +accelerate launch tools/train_dreambooth_all.py \ + --train_text_encoder \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --mixed_precision="fp16" \ + --instance_data_dir=$INSTANCE_DIR \ + --instance_prompt="dog" \ + --with_prior_preservation --prior_loss_weight=1.0 \ + --class_prompt="dog" \ + --class_data_dir=$CLASS_DIR \ + --num_class_images=200 \ + --output_dir=$OUTPUT_DIR \ + --logging_dir=$LOG_DIR \ + --train_batch_size=1 \ + --gradient_accumulation_steps=1 --gradient_checkpointing \ + --use_8bit_adam \ + --learning_rate=2e-6 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --auto_test_model \ + --test_prompts_file=$TEST_PROMPTS_FILE \ + --test_seed=123 \ + --test_num_per_prompt=3 \ + --max_train_steps=1000 \ + --save_model_every_n_steps=500 + +# 如果max_train_steps改大了,请记得把save_model_every_n_steps也改大 +# 不然磁盘很容易中间就满了 + +# 以下是核心参数介绍: +# 主要的几个 +# --train_text_encoder 训练文本编码器 +# --mixed_precision="fp16" 混合精度训练 +# - center_crop +# 是否裁剪图片,一般如果你的数据集不是正方形的话,需要裁剪 +# - resolution +# 图片的分辨率,一般是512,使用该参数会自动缩放输入图像 +# 可以配合center_crop使用,达到裁剪成正方形并缩放到512*512的效果 +# - instance_prompt +# 如果你希望训练的是特定的人物,使用该参数 +# 如 --instance_prompt="a photo of girl" +# - class_prompt +# 如果你希望训练的是某个特定的类别,使用该参数可能提升一定的训练效果 +# - use_txt_as_label +# 是否读取与图片同名的txt文件作为label +# 如果你要训练的是整个大模型的图像风格,那么可以使用该参数 +# 该选项会忽略instance_prompt参数传入的内容 +# - learning_rate +# 学习率,一般是2e-6,是训练中需要调整的关键参数 +# 太大会导致模型不收敛,太小的话,训练速度会变慢 +# - lr_scheduler, 可选项有constant, linear, cosine, cosine_with_restarts, cosine_with_hard_restarts +# 学习率调整策略,一般是constant,即不调整,如果你的数据集很大,可以尝试其他的,但是可能会导致模型不收敛,需要调整学习率 +# - lr_warmup_steps,如果你使用的是constant,那么这个参数可以忽略, +# 如果使用其他的,那么这个参数可以设置为0,即不使用warmup +# 也可以设置为其他的值,比如1000,即在前1000个step中,学习率从0慢慢增加到learning_rate的值 +# 一般不需要设置, 除非你的数据集很大,训练收敛很慢 +# - max_train_steps +# 训练的最大步数,一般是1000,如果你的数据集比较大,那么可以适当增大该值 +# - save_model_every_n_steps +# 每多少步保存一次模型,方便查看中间训练的结果找出最优的模型,也可以用于断点续训 + +# --with_prior_preservation,--prior_loss_weight=1.0,分别是使用先验知识保留和先验损失权重 +# 如果你的数据样本比较少,那么可以使用这两个参数,可以提升训练效果,还可以防止过拟合(即生成的图片与训练的图片相似度过高) + +# --auto_test_model, --test_prompts_file, --test_seed, --test_num_per_prompt +# 分别是自动测试模型(每save_model_every_n_steps步后)、测试的文本、随机种子、每个文本测试的次数 +# 测试的样本图片会保存在模型输出目录下的test文件夹中 \ No newline at end of file diff --git a/train_style_all.sh b/train_style_all.sh new file mode 100644 index 0000000..a5933d7 --- /dev/null +++ b/train_style_all.sh @@ -0,0 +1,60 @@ +# 主要用于训练风格、作画能力(需要每张图片都有对应的标签描述) +export MODEL_NAME="./model" +export INSTANCE_DIR="./datasets/test2" +export OUTPUT_DIR="./new_model" +export LOG_DIR="/root/tf-logs" +export TEST_PROMPTS_FILE="./test_prompts_style.txt" + +rm -rf $LOG_DIR/* + +accelerate launch tools/train_dreambooth_all.py \ + --pretrained_model_name_or_path=$MODEL_NAME \ + --mixed_precision="fp16" \ + --instance_data_dir=$INSTANCE_DIR \ + --use_txt_as_label \ + --output_dir=$OUTPUT_DIR \ + --logging_dir=$LOG_DIR \ + --train_batch_size=1 \ + --use_8bit_adam \ + --gradient_accumulation_steps=1 --gradient_checkpointing \ + --learning_rate=2e-6 \ + --lr_scheduler="constant" \ + --lr_warmup_steps=0 \ + --max_train_steps=1000 \ + --save_model_every_n_steps=500 \ + --auto_test_model \ + --test_prompts_file=$TEST_PROMPTS_FILE \ + --test_seed=123 \ + --test_num_per_prompt=3 + +# 如果max_train_steps改大了,请记得把save_model_every_n_steps也改大,不然磁盘容易中间就满了 + +# 以下是核心参数介绍: +# 主要的几个 +# --train_text_encoder 训练文本编码器 +# --mixed_precision="fp16" 混合精度训练 +# - center_crop +# 是否裁剪图片,一般如果你的数据集不是正方形的话,需要裁剪 +# - resolution +# 图片的分辨率,一般是512,使用该参数会自动缩放输入图像 +# 可以配合center_crop使用,达到裁剪成正方形并缩放到512*512的效果 +# - instance_prompt +# 如果你希望训练的是特定的人物,使用该参数 +# 如 --instance_prompt="a photo of girl" +# - use_txt_as_label +# 是否读取与图片同名的txt文件作为label +# 如果你要训练的是整个大模型的图像风格,那么可以使用该参数 +# 该选项会忽略instance_prompt参数传入的内容 +# - learning_rate +# 学习率,一般是2e-6,是训练中需要调整的关键参数 +# 太大会导致模型不收敛,太小的话,训练速度会变慢 +# - max_train_steps +# 训练的最大步数,一般是1000,如果你的数据集比较大,那么可以适当增大该值 +# - save_model_every_n_steps +# 每多少步保存一次模型,方便查看中间训练的结果找出最优的模型,也可以用于断点续训 + +# --train_text_encoder # 除了图像生成器,也训练文本编码器 + +# --auto_test_model, --test_prompts_file, --test_seed, --test_num_per_prompt +# 分别是自动测试模型(每save_model_every_n_steps步后)、测试的文本、随机种子、每个文本测试的次数 +# 测试的样本图片会保存在模型输出目录下的test文件夹中 \ No newline at end of file diff --git "a/\350\277\220\350\241\214.ipynb" "b/\350\277\220\350\241\214.ipynb" index d2bd112..0661311 100644 --- "a/\350\277\220\350\241\214.ipynb" +++ "b/\350\277\220\350\241\214.ipynb" @@ -136,20 +136,28 @@ "id": "2633307a", "metadata": {}, "source": [ - "[可选] 保留更高质量的裁剪(矩形裁剪) \n", - "-- 不需要修改width和height,自适应" + "最大边同比缩放(后续输入任意比例的图片) " ] }, { "cell_type": "code", - "execution_count": 20, + "execution_count": 9, "id": "10d2bb3d-9002-4d3b-a4be-f5f74a008b9c", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "d:\\PC\\cyvision\\dreambooth-for-diffusion\\tools\\scale_images.py:28: DeprecationWarning: ANTIALIAS is deprecated and will be removed in Pillow 10 (2023-07-01). Use Resampling.LANCZOS instead.\n", + " resized_image = image.resize((new_width, new_height), Image.ANTIALIAS)\n" + ] + } + ], "source": [ - "!python tools/handle_images.py --origin_image_path ./datasets/test \\\n", - " --output_image_path_0 ./datasets/a1 --output_image_path_1 ./datasets/a2 \\\n", - " --width=768 --height=512 " + "!python tools/scale_images.py --origin_image_path ./datasets/test \\\n", + " --output_image_path ./datasets/test5 \\\n", + " --max_side_length 712" ] }, { @@ -172,7 +180,7 @@ "outputs": [], "source": [ "# 该步根据需要标注文件数量不同,需要运行一段时间(测试6000张图片需要10分钟)\n", - "!python tools/label_images.py --path=./datasets/test2 --model_path='./deepdanbooru/'" + "!python tools/label_images.py --path=./datasets/cyAnime --model_path='./deepdanbooru/'" ] }, { @@ -289,12 +297,12 @@ { "cell_type": "code", "execution_count": null, - "id": "c7b8f2e4", + "id": "78b21c73", "metadata": {}, "outputs": [], "source": [ - "# 选择矩形图像数据集训练\n", - "!sh train_object_rect.sh --width 512 --height 768 ## " + "# 事先进行最大边同比缩放后,支持任意尺寸输入(batch_size=1)\n", + "!sh train_object_all.sh " ] }, { @@ -327,8 +335,8 @@ "metadata": {}, "outputs": [], "source": [ - "# 如果输入图像是同等的矩形\n", - "!sh train_object_rect.sh --width 768 --height 512 " + "# 事先进行最大边同比缩放后,支持任意尺寸输入(batch_size=1)\n", + "!sh train_style_all.sh " ] }, {