はじめに
ControlNet は Stable Diffusion に追加の制御層を提供し、公式の実装では深度、エッジライン、OpenPose などの複数の次元から生成画像を制御できます。今回は、明るさ（brightness / grayscale）による画像生成の制御を目指し、古い写真のカラー復元や既存画像の再着色などのニーズに対応します。
本記事では、HuggingFace Diffusers を使用して Brightness ControlNet を訓練するプロセスを記録し、紹介します。
データセットの準備
データソース：
- LAION-Aesthetics V1（LAION の美的スコアが 7 以上のサブセット）
- COYO-700M（aesthetic_score_laion_v2 スコアを含む）
データのダウンロード：
from img2dataset import download
import shutil
import multiprocessing
def main():
download(
processes_count=16,
thread_count=64,
url_list="laion2B-en-aesthetic",
resize_mode="center_crop",
image_size=512,
output_folder="../laion-en-aesthetic",
output_format="files",
input_format="parquet",
skip_reencode=True,
save_additional_columns=["similarity","hash","punsafe","pwatermark","aesthetic"],
url_col="URL",
caption_col="TEXT",
distributor="multiprocessing",
)
if __name__ == "__main__":
multiprocessing.freeze_support()
main()
HuggingFace Datasets の構築、ローカルに保存し Hub にプッシュ：
import os
from datasets import Dataset
from pathlib import Path
from PIL import Image
data_dir = Path(r"H:\DataScience\laion-en-aesthetic")
def entry_for_id(image_folder, filename):
img = Image.open(image_folder / filename)
gray_img = img.convert('L')
caption_filename = filename.replace('.jpg', '.txt')
with open(image_folder / caption_filename) as f:
caption = f.read()
return {
"image": img,
"grayscale_image": gray_img,
"caption": caption,
}
max_images = 1000000
def generate_entries():
index = 0
# cc3m のすべてのサブフォルダ
image_folders = [f.path for f in os.scandir(data_dir) if f.is_dir()]
for image_folder in image_folders:
image_folder = Path(image_folder)
print(image_folder)
# cc3m サブフォルダのすべてのファイル
for filename in os.listdir(image_folder):
if not filename.endswith('.jpg'):
continue
yield entry_for_id(image_folder, filename)
index += 1
if index >= max_images:
break
if index >= max_images:
break
ds = Dataset.from_generator(generate_entries, cache_dir="./.cache")
ds.save_to_disk("./grayscale_image_aesthetic_1M")
ds.push_to_hub('ioclab/grayscale_image_aesthetic_1M', private=True)
トレーニングプロセス
ControlNet training example スクリプトを使用してトレーニングを行い、具体的なパラメータは以下の通りです：
accelerate launch train_controlnet_local.py \
--pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5" \
--output_dir="./output_v1a2u" \
--dataset_name="./grayscale_image_aesthetic_100k" \
--resolution=512 \
--learning_rate=1e-5 \
--image_column=image \
--caption_column=caption \
--conditioning_image_column=grayscale_image \
--train_batch_size=16 \
--gradient_accumulation_steps=4 \
--num_train_epochs=2 \
--tracker_project_name="control_v1a2u_sd15_brightness" \
--enable_xformers_memory_efficient_attention \
--checkpointing_steps=5000 \
--hub_model_id="ioclab/grayscale_controlnet" \
--report_to wandb \
--push_to_hub
wandb バックエンドデータ：
A6000 GPU トレーニング時間：13h、sample_count：100k、epoch：1、batch_size：16、gradient_accumulation_steps：1。
TPU v4-8 GPU トレーニング時間：25h、sample_count：3m、epoch：1、batch_size：2、gradient_accumulation_steps：25。
トレーニングレポート：
Google が提供する TPU v4-8 マシンは、240 コア 480 スレッド CPU、400GB メモリ、128GB TPU メモリ、2000Mbps 帯域幅、3TB ディスクを搭載しています。
簡単な計算では、TPU v4-8 bf16 は単一の A6000 fp32 と比較して 15 倍の速度向上を示しています。
ControlNet 論文で言及された "Sudden Convergence" 現象：
効果
参考資料
-
Adding Conditional Control to Text-to-Image Diffusion Models
ControlNet 論文。原理説明、トレーニングパラメータ、比較画像などの重要な情報を含みます。
-
ControlNet 公式リポジトリ、トレーニングチュートリアルを含みます。
-
ControlNet 1.1 Nightly。
-
Train your ControlNet with diffusers 🧨
HuggingFace 公式の Diffusers を使用した ControlNet トレーニングチュートリアル、非常に詳細です。
-
HuggingFace ControlNet トレーニングスクリプトの例。
-
JAX/Diffusers community sprint 🧨
HuggingFace × Google コミュニティスプリントイベントのドキュメント。