add any aspect input

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步
 

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)

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 <xxx> 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文件夹中

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 <xxx> 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文件夹中

运行.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  "
    ]
   },
   {