見出し画像

WinローカルでSDXLのControlNet学習

とりにく

自分用まとめ。SDXLのControlNet学習についてのメモです。
Canny画像→ノーマルマップ画像を出力するモデル作成を目的としたいと思います。

PC環境

OS：Windows10
CPU：i9-9900K
メモリ：64GB
GPU：RTX3090

①データセットの撮影（Unity使用）

Unityで自作スクリプト
https://github.com/tori29umai0123/CotrolNet_Shadow
を使って撮影。今回は順光とノーマルマップだけ撮影。

myminifactory等で公開されている2214体の3D彫像モデルをあらゆる角度から順光画像とノーマルマップ画像を30,000枚ずつ撮影しました。

順光
ノーマルマップ

②データセットの前処理

適当なディレクトリにデータセット作成用のディレクトリを作り、そこで前処理します

cd C:\
mkdir "DatasetPreprocessing"
cd DatasetPreprocessing
python -m venv venv
venv\Scripts\activate
pip install opencv-python Pillow datasets

canny_convert.py

import cv2
import os
import sys
from pathlib import Path

def canny_process_image(image_path, output_path):
    # 画像を読み込む
    img = cv2.imread(image_path)
    # グレースケール変換
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # Cannyエッジ検出
    edges = cv2.Canny(gray, 15, 50)
    # 処理済みの画像をPNG形式で保存
    cv2.imwrite(output_path, edges)

def process_directory(input_dir, output_dir):
    for root, dirs, files in os.walk(input_dir):
        for file in files:
            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')):
                input_path = os.path.join(root, file)
                relative_path = os.path.relpath(input_path, input_dir)
                # 元のファイル名を保持しつつ、拡張子を.pngに変更
                relative_path_no_ext = os.path.splitext(relative_path)[0] + '.png'
                output_path = os.path.join(output_dir, relative_path_no_ext)
                output_file_dir = os.path.dirname(output_path)
                
                # 必要に応じて出力ディレクトリを作成
                if not os.path.exists(output_file_dir):
                    os.makedirs(output_file_dir)
                
                # 画像をCanny処理
                canny_process_image(input_path, output_path)

if __name__ == "__main__":
    # コマンドライン引数の取得時にエラー処理を追加
    if len(sys.argv) < 3:
        print("Usage: python script.py <input_dir> <output_dir>")
        sys.exit(1)

    input_dir = sys.argv[1]
    output_dir = sys.argv[2]

    # ディレクトリを処理
    process_directory(input_dir, output_dir)

【使い方】

python canny_convert.py <入力フォルダ> <出力フォルダ>

ノーマルマップ画像→Canny画像になる。

ノーマルマップ画像
Canny画像

この二つの関係性を学習していきます。

③promptの解析

画像のprompt解析ソフトを作成してビルドしたので配布します。

↑を使って画像フォルダ以下にpromptのテキストファイルを保存します。

上記の保存されたテキストファイルをベースにpromptを結合して一つのjsonにします。

create_json.py


import json
import os
import sys

def create_json_from_txt(folder_path, output_folder):
    # 出力ファイル名を prompt.json に固定
    output_file = os.path.join(output_folder, "prompt.json")

    # 出力用のフォルダが存在しない場合は作成
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    with open(output_file, 'w', encoding='utf-8') as outfile:
        # 指定されたフォルダ内のファイルリストを取得し、ソートする
        sorted_filenames = sorted(os.listdir(folder_path))
        # ソートされたファイルリストをループ処理
        for filename in sorted_filenames:
            if filename.endswith(".txt"):
                # .txtファイルの完全なパスを取得
                txt_file_path = os.path.join(folder_path, filename)
                # ファイル名から拡張子を取り除いて基本名を取得
                base_filename = os.path.splitext(filename)[0]
                # sourceとtargetのパスを決め打ちで生成
                source_png = f"source/{base_filename}.png"
                target_png = f"target/{base_filename}.png"

                # .txtファイルを開いて内容を読み込む
                with open(txt_file_path, 'r', encoding='utf-8') as file:
                    prompt = file.read().strip()  # ファイルの余計な空白を削除

                    # JSONオブジェクトを生成
                    json_object = {
                        "source": source_png,
                        "target": target_png,
                        "prompt": prompt
                    }

                    # JSONオブジェクトをファイルに書き込む
                    json.dump(json_object, outfile, ensure_ascii=False)
                    outfile.write("\n")  # 各オブジェクトを新しい行に書き込む

if __name__ == '__main__':
    if len(sys.argv) < 3:
        print("Usage: python script.py <source_folder> <output_folder>")
        sys.exit(1)

    source_folder = sys.argv[1]
    output_folder = sys.argv[2]

    # スクリプトを実行するための関数を呼び出す
    create_json_from_txt(source_folder, output_folder)

ちなみにノーマルマップ画像をprompt分析させるより、順光画像分析の方が解析精度が良かったので順光画像を使っています。
また上記はundesired_tagsが空欄ですが、実際自分が使う時は
undesired_tags = ["monochrome", "greyscale"]
のように除去ワードをいれています。（ノーマルマップを学習させる時、グレスケやモノクロのタグ情報は不要なので）

【使い方】

python create_json.py <入力フォルダ> <出力フォルダ>

以下のような形式の出力フォルダ/prompt.jsonを出力する。

sourceフォルダ：Canny画像
targetフォルダ：Normalmap画像
prompt.json

これでデータセットのベースができました。

dataset_save.py

import os
import sys
from datasets import load_dataset, Image, DatasetDict


def map_fn(source, target, prompt, input_dir):
    """
    map関数は、与えられたsource、target、promptを受け取り、画像と条件画像のフルパスとテキストを返します。
    """
    image = os.path.join(input_dir, target)
    conditioning_image = os.path.join(input_dir, source)
    text = prompt
    return dict(image=image, conditioning_image=conditioning_image, text=text)

if __name__ == "__main__":
    # コマンドライン引数の取得時にエラー処理を追加
    if len(sys.argv) < 3:
        print("Usage: python script.py <input_dir> <output_dir>")
        sys.exit(1)

    input_dir = sys.argv[1]
    output_dir = sys.argv[2]

    # JSONファイルのパスを構築
    json_path = os.path.join(input_dir, 'prompt.json')

    # データセットの読み込み
    dataset = load_dataset('json', data_files=json_path)['train']

    # map関数を適用して、必要な変換を行う
    dataset = dataset.map(lambda example: map_fn(example['source'], example['target'], example['prompt'], input_dir))

    # 不要な列の削除
    dataset = dataset.remove_columns(['source', 'target', 'prompt'])

    # カラムタイプを画像にキャスト
    dataset = dataset.cast_column("image", Image(decode=True))
    dataset = dataset.cast_column("conditioning_image", Image(decode=True))

    # DatasetDictに変換
    dataset_dict = DatasetDict(train=dataset)

    # 出力ディレクトリに保存
    dataset_dict.save_to_disk(output_dir)
    print(f"Dataset saved to {output_dir}")

【使い方】

python dataset_save.py "ベースデータ入力元" "データセット出力先"

diffusersで学習できる形に変換されました

ついでに変換されたデータをHugging Faceにアップロードしてしまいます
やり方はこちらを参照

上記の手順で
1.ログイン
2.リポジトリの作成
を済ませた後、以下のスクリプトを実行
dataset_upload.py

import os
import sys
from datasets import load_dataset
from huggingface_hub import HfApi, HfFolder

def load_dataset_from_arrow_files(base_path):
    """
    Loads a dataset from a specified path containing Arrow files and returns it.
    Assumes Arrow files are part of a training dataset split.
    """
    # Define the path pattern for the training Arrow files
    arrow_files_pattern = os.path.join(base_path, "train", "data-*.arrow")
    
    # Load the dataset from Arrow files
    # Assuming all Arrow files belong to the 'train' split
    dataset = load_dataset('arrow', data_files={'train': arrow_files_pattern})

    return dataset

def upload_dataset_to_hub(dataset_dict, dataset_name, user_name):
    """
    DatasetDictをHugging Face Hubにアップロードする。
    """
    api = HfApi()
    token = HfFolder.get_token()
    if token is None:
        raise ValueError("Hugging Faceの認証トークンが見つかりません。huggingface-cliでログインしてください。")

    # データセットをアップロード
    dataset_dict.push_to_hub(f"{user_name}/{dataset_name}", token=token, private=True)

if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("Usage: python dataset_upload.py <input_dir> <dataset_name> <user_name>")
        sys.exit(1)

    input_dir = sys.argv[1]
    dataset_name = sys.argv[2]
    user_name = sys.argv[3]

    # DatasetDictをロード
    dataset_dict = load_dataset_from_arrow_files(input_dir)

    # DatasetDictをHugging Face Hubにアップロード
    upload_dataset_to_hub(dataset_dict, dataset_name, user_name)

以下のように実行します

python dataset_upload.py データセットパス リポジトリ名 ユーザー名

仮に以下のようなデータセットをアップロードする場合、

python dataset_upload.py E:\AI\Canny2Normalmap_basedata Canny2Normalmap tori29umai

になります。

ベースモデルのダウンロード（diffusers形式）

Dataset Preprocessing以下にスクリプトを作成。
diffusers_models_dl_SDXL.py

import os
import requests
import sys
from tqdm import tqdm

# ファイルをダウンロードする関数
def download_file(url, output_file):
    response = requests.get(url, stream=True)
    total_size = int(response.headers.get("content-length", 0))

    with open(output_file, "wb") as f, tqdm(
        desc=output_file, total=total_size, unit="B", unit_scale=True
    ) as pbar:
        for data in response.iter_content(chunk_size=1024):
            f.write(data)
            pbar.update(len(data))

# ファイルをダウンロードする関数（複数のファイルを含むフォルダ）
def download_files(repo_id, subfolder, files, cache_dir):
    for file in tqdm(files, desc="Downloading files"):
        url = f"https://huggingface.co/{repo_id}/resolve/main/{subfolder}/{file}"
        output_file = os.path.join(cache_dir, file)
        if not os.path.exists(output_file):
            print(f"Downloading {file} from {url} to {output_file}...")
            download_file(url, output_file)
            print(f"{file} downloaded successfully!")
        else:
            print(f"{file} is already downloaded")

# モデルデータをダウンロードする関数
def check_and_download_model(model_dir, model_id, sub_dirs, files):
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)
        print(f"Downloading model to {model_dir}. Model ID: {model_id}")

        for sub_dir, sub_dir_files in sub_dirs:
            sub_dir_path = os.path.join(model_dir, sub_dir)
            if not os.path.exists(sub_dir_path):
                os.makedirs(sub_dir_path)
            download_files(model_id, sub_dir, sub_dir_files, sub_dir_path)

        for file in files:
            url = f"https://huggingface.co/{model_id}/resolve/main/{file}"
            output_file = os.path.join(model_dir, file)
            if not os.path.exists(output_file):
                print(f"Downloading {file} from {url} to {output_file}...")
                download_file(url, output_file)
                print(f"{file} downloaded successfully!")
            else:
                print(f"{file} is already downloaded")

        print("Model download completed.")
    else:
        print("Model is already downloaded.")

def download_diffusion_SDXL(model_dir, MODEL_ID):
    SUB_DIRS = [
        ("scheduler", ["scheduler_config.json"]),
        ("text_encoder", ["config.json", "model.safetensors"]),
        ("text_encoder_2", ["config.json", "model.safetensors"]),
        ("tokenizer", ["merges.txt", "special_tokens_map.json", "tokenizer_config.json", "vocab.json"]),
        ("tokenizer_2", ["merges.txt", "special_tokens_map.json", "tokenizer_config.json", "vocab.json"]),
        ("unet", ["config.json", "diffusion_pytorch_model.safetensors"]),
        ("vae", ["config.json", "diffusion_pytorch_model.safetensors"]),
    ]
    FILES = ["model_index.json"]

    check_and_download_model(model_dir, MODEL_ID, SUB_DIRS, FILES)

if __name__ == "__main__":
    # コマンドライン引数の取得時にエラー処理を追加
    if len(sys.argv) < 3:
        print("Usage: script.py <model_dir> <MODEL_ID>")
        sys.exit(1)
    model_dir = os.path.join(sys.argv[1], sys.argv[2])  # MODEL_IDを含むディレクトリパスを作成
    MODEL_ID = sys.argv[2]
    download_diffusion_SDXL(model_dir, MODEL_ID)

【使い方】

python diffusers_models_dl_SDXL.py "保存先" "モデルのID（Hugging Face形式）"

【具体例】

python diffusers_models_dl_SDXL.py "D:\Stable-diffusion\diffusers" "cagliostrolab/animagine-xl-3.0"

指定されたディレクトリにdiffusers形式でモデルがDLされます。
あるいはsafetensors→Diffusersフォーマットに変換する手法もあるようです。（好きなモデルをベースにしたい時はこっちがいいかも）
https://touch-sp.hatenablog.com/entry/2023/08/10/225418

diffusersのインストール

git clone https://github.com/huggingface/diffusers
cd C:\diffusers
python -m venv venv
venv\Scripts\activate

pip install -e .

cd C:\diffusers\examples\controlnet

pip install -r requirements_sdxl.txt
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118 -U
pip install peft==0.8.2 -U
pip install xformers==0.0.20 -U
pip install bitsandbytes==0.41.1 --prefer-binary --extra-index-url=https://jllllll.github.io/bitsandbytes-windows-webui -U

学習開始

学習前に検証用画像（validation_image）と検証用prompt（validation_prompt）を準備しておきます
validation_image


canny画像

validation_prompt


1girl, solo, looking_at_viewer, bangs, thighhighs, gloves, hair_between_eyes, closed_mouth, standing, full_body, boots, shorts, belt, hood, cape, hand_on_hip, thigh_boots, cloak

実行コード

accelerate launch train_controlnet_sdxl.py ^
 --pretrained_model_name_or_path="DiffusersフォーマットのSDXLモデルのパス" ^
 --output_dir="出力先のパス" ^
 --train_data_dir="データセットのパス" ^
 --mixed_precision="fp16" ^
 --resolution=512 ^
 --learning_rate=1e-5 ^
 --train_batch_size=1 ^
 --tracker_project_name="controlnet" ^
 --enable_xformers_memory_efficient_attention ^
 --use_8bit_adam ^
 --max_train_steps=50000 ^
 --checkpointing_steps=500 ^
 --num_validation_images=1 ^
 --validation_steps=500 ^
 --validation_image "検証用画像" ^
 --validation_prompt "検証用prompt" ^
 --gradient_accumulation_steps=4 ^
 --seed=42

実行コード（具体例）

accelerate launch train_controlnet_sdxl.py ^
 --pretrained_model_name_or_path="D:\Stable-diffusion\diffusers\cagliostrolab\animagine-xl-3.0" ^
 --output_dir="D:\Stable-diffusion\Canny2Normalmap" ^
 --train_data_dir="E:\AI\Canny2Normalmap_dataset" ^
 --mixed_precision="fp16" ^
 --resolution=512 ^
 --learning_rate=1e-5 ^
 --train_batch_size=1 ^
 --tracker_project_name="controlnet" ^
 --enable_xformers_memory_efficient_attention ^
 --use_8bit_adam ^
 --max_train_steps=50000 ^
 --checkpointing_steps=500 ^
 --num_validation_images=1 ^
 --validation_steps=500 ^
 --validation_image "E:\desktop\output\test01.png" ^
 --validation_prompt "1girl, solo, looking_at_viewer, bangs, thighhighs, gloves, hair_between_eyes, closed_mouth, standing, full_body, boots, shorts, belt, hood, cape, hand_on_hip, thigh_boots, cloak" ^
 --gradient_accumulation_steps=4 ^
 --seed=42

ログの確認

学習しているコマンドプロンプトウインドウとは別に
diffusersの仮想環境に入り、以下のコマンドを実行

tensorboard --logdir=<ログディレクトリのパス>

すると以下のようにURLが表示されるので表示されたURLにブラウザでアクセスするとログが見れます。

よくわからないけんどなんかでてきました。見方はよくわからん。検証用画像でなんとなく判断できそうです。

学習結果については長くなったので次回へ続く・・・！（ネタバレするとまぁ成功しました

追記：
結構、学習に時間がかかるため（）一旦学習を止めてまた再開、みたいなことがよく起きますその際のコマンドも自分用にメモ

accelerate launch train_controlnet_sdxl.py ^
 --pretrained_model_name_or_path="D:\Stable-diffusion\diffusers\cagliostrolab\animagine-xl-3.0" ^
 --output_dir="D:\Stable-diffusion\Canny2Normalmap" ^
 --train_data_dir="E:\AI\Canny2Normalmap_dataset" ^
 --mixed_precision="fp16" ^
 --resolution=512 ^
 --learning_rate=1e-5 ^
 --train_batch_size=1 ^
 --tracker_project_name="controlnet" ^
 --enable_xformers_memory_efficient_attention ^
 --use_8bit_adam ^
 --max_train_steps=50000 ^
 --checkpointing_steps=500 ^
 --num_validation_images=1 ^
 --validation_steps=500 ^
 --validation_image "E:\desktop\output\test01.png" ^
 --validation_prompt "1girl, solo, looking_at_viewer, bangs, thighhighs, gloves, hair_between_eyes, closed_mouth, standing, full_body, boots, shorts, belt, hood, cape, hand_on_hip, thigh_boots, cloak" ^
 --gradient_accumulation_steps=4 ^
 --seed=42 ^
 --resume_from_checkpoint="D:\Stable-diffusion\Canny2Normalmap\checkpoint-7000"


この記事が気に入ったらサポートをしてみませんか？