クラウドでLoRA作成(追加学習)がしたいときModalを利用するとすごく捗る(画像生成AI Stable diffusion)-14にゴー

前回Stable Diffusion WebUIをModalで動かす記事を書きましたが、学習がやりやすい気配を感じたので今回はLoRAの学習をModalでやってみました。

追記(2023-07-28):
- optimizerを設定できるようにコードを変更
- AdamWとprodigyを使えるように。他は未確認
- よく使う変数をコード上部に持ってきた
- 記事中の画像は以前のコードで作ったものなので現在のコードで作ると結果が(多分)変わります
- 以前のコードは末尾に追記

方針

kohya氏のsd-scriptsを利用
学習について、共通編のDreamBooth、キャプション方式を参考に学習データを準備。ただし正則化画像は使ってません
東北ずん子のAI画像モデル用学習データを利用。タグなどはこれをそのまま使う
学習素材をmodal.Mountでアップロードし、作ったLoRAのデータをリモートからローカルに受け渡して保存

ModalでLoRAを作ることしか考えていないのでキャラクターの再現度を高める工夫などは特にしていないです。

事前準備

Modalを利用可能にしておく

Modalのアカウントを取得し、modal-clientのインストールとmodalのトークンの取得を済ませておいてください。Modalを使ってStable Diffusion WebUIを動かすのModalの準備と同じ手順です。

virtualenv venv
venv\Scripts\activate.bat

pip install modal-client

modal token new

東北ずん子のAI画像モデル用学習データをダウンロード

https://zunko.jp/con_illust.htmlの下の方に画像学習用のデータへのリンクがあります。Googleドライブに保存されているので01_LoRA学習用データ_A氏提供版_背景白をダウンロードしてください。

この中にあるzunkoフォルダを使います。

学習用フォルダのディレクトリ構成

学習用フォルダとして適当(スペースや全角文字などが入らない)な場所にzunkotrainというフォルダを作り、そこにzunkoフォルダと設定ファイルを置きます。設定ファイル(trainconfig.toml)は後述。

zunkotrain
- zunko
  - zko (1).png
  - zko (1).txt
  - 以下略
- trainconfig.toml

設定ファイル

以下をtrainconfig.tomlとして学習用フォルダに保存。

[general]
enable_bucket = true                        # Aspect Ratio Bucketingを使うか否か

[[datasets]]
resolution = 512                            # 学習解像度
batch_size = 1                              # バッチサイズ

  [[datasets.subsets]]
  image_dir = '/traindata/zunko'                     # 学習用画像を入れたフォルダをModal上でのディレクトリで指定。
  caption_extension = '.txt'            # キャプションファイルの拡張子
  num_repeats = 1                          # 学習用画像の繰り返し回数

pythonを記述して学習を実行

コード

以下をtrainlora.pyとして適当なディレクトリに保存。

import os
import sys
import subprocess
import shlex
import datetime
import shutil

import toml
import modal


GPU = "T4"  # "T4", "A10G", "A100", modal.gpu.A100(memory=20)
TIMEOUT = 10800

# loraの学習に使うモデルのURL
PRETRAINED_MODEL_URL = "https://huggingface.co/WarriorMama777/OrangeMixs/resolve/main/Models/AbyssOrangeMix3/AOM3_orangemixs.safetensors"
PRETRAINED_MODEL_NAME = PRETRAINED_MODEL_URL.split("/")[-1]

LOCAL_TRAIN_DIR = "C:/path/to/trainimages/zunkotrain"  # 学習用のフォルダへのパス
TOML_FILE = "trainconfig.toml"  # 設定ファイルの名前

REMOTE_TRAIN_DIR = "/traindata"  #  学習用フォルダのマウント先のディレクトリ
REMOTE_OUTPUT_DIR = "/trainedoutputs"  # 生成されたloraを保存する場所
OUTPUT_NAME = "zunko"  # loraの名前

dt_now = datetime.datetime.now()
timestamp = dt_now.strftime('%Y%m%d%H%M%S')
LOCAL_SAVE_TO = timestamp + OUTPUT_NAME  # 成果物のローカルの保存先


# よく使う設定
NETWORK_DIM = 32
NETWORK_ALPHA = 16

MAX_TRAIN_EPOCHS = 30
SAVE_EVERY_N_EPOCHS = 10
OPTIMIZER = "AdamW"  # AdamW, prodigy
LEARNING_RATE = "1e-4"  # 1e-4
LR_SCHEDULER = "constant"  # constant, cosine, cosine_with_restarts


stub = modal.Stub("trainlora")
mount = modal.Mount.from_local_dir(LOCAL_TRAIN_DIR, remote_path=REMOTE_TRAIN_DIR)

image = (
    modal.Image.conda(python_version="3.10")
    .apt_install("git", "wget", "libgl1-mesa-dev", "libglib2.0-0", "libsm6", "libxrender1", "libxext6")
    .conda_install(
        "cudatoolkit=11.7",
        "cuda-nvcc",
        channels=["conda-forge", "nvidia"],
    )
    .pip_install(
        "torch==1.13.1+cu117",
        "torchvision==0.14.1+cu117",
        extra_index_url="https://download.pytorch.org/whl/cu117"
    )
    .pip_install(
        "triton",
        pre=True
    )
    .pip_install(
        "accelerate==0.15.0",
        "transformers==4.26.0",
        "ftfy==6.1.1",
        "albumentations==1.3.0",
        "opencv-python==4.7.0.68",
        "einops==0.6.0",
        "diffusers[torch]==0.10.2",
        "pytorch-lightning==1.9.0",
        "bitsandbytes==0.35.0",
        "tensorboard==2.10.1",
        "safetensors==0.2.6",
        "altair==4.2.2",
        "easygui==0.98.3",
        "toml==0.10.2",
        "voluptuous==0.13.1",
        "requests==2.28.2",
        "timm==0.6.12",
        "fairscale==0.4.13",
        "tensorflow==2.10.1",
        "huggingface-hub==0.13.3",
        "xformers==0.0.16rc425",
        "prodigyopt",
    )
)


@stub.function(image=image,
               mounts=[mount],
               gpu="T4",
               timeout=10800)
def run_train_lora():
    subprocess.run("accelerate config default --mixed_precision fp16", shell=True)
    subprocess.run("git clone https://github.com/kohya-ss/sd-scripts", shell=True)
    os.chdir("sd-scripts")
    print("downloading model")
    subprocess.run(f"wget {PRETRAINED_MODEL_URL} --no-verbose", shell=True)

    print("run accelerate launch")
    subprocess.run("accelerate launch --num_cpu_threads_per_process 1 "
        "train_network.py "
        f"--pretrained_model_name_or_path={PRETRAINED_MODEL_NAME} "
        f"--dataset_config={REMOTE_TRAIN_DIR}/{TOML_FILE} "
        f"--output_dir={REMOTE_OUTPUT_DIR} "
        f"--output_name={OUTPUT_NAME} "
        f"--learning_rate={LEARNING_RATE} "
        f"--max_train_epochs={MAX_TRAIN_EPOCHS} --xformers "
        f"--optimizer_type={OPTIMIZER} "
        f"--network_dim={NETWORK_DIM} --network_alpha={NETWORK_ALPHA} "
        f"--lr_scheduler={LR_SCHEDULER} "
        # f"--lr_scheduler_num_cycles={LR_SCHEDULER_NUM_CYCLES} "
        # "--optimizer_args weight_decay=0 betas=0.9,0.99 d0=1e-6 "
        "--gradient_checkpointing --mixed_precision=fp16 "
        "--cache_latents "
        "--logging_dir=/tmp/logs "
        f"--save_every_n_epochs={SAVE_EVERY_N_EPOCHS} --save_precision=fp16 "
        "--save_model_as=safetensors  --network_module=networks.lora", shell=True)

    os.chdir(REMOTE_OUTPUT_DIR)
    for loraname in os.listdir(REMOTE_OUTPUT_DIR):
        f = open(loraname, "rb")
        lora_data = f.read()
        yield (lora_data, loraname)
        f.close()


@stub.local_entrypoint()
def main():
    os.makedirs(LOCAL_SAVE_TO, exist_ok=True)
    for lora_data, loraname in run_train_lora.call():
        saved_path = os.path.join(LOCAL_SAVE_TO, loraname)
        with open(saved_path, "wb") as f:
            f.write(lora_data)

どこかで不具合があると途中経過が全部消えるので、最初はMAX_TRAIN_EPOCHS=2などとして小さく実験してみるといいかもしれません。

(追記、変更部分です。気になる方は末尾にある以前のコードを参照してください。この後出てくるLoRAを使った画像は以前のコードで作られた画像です)

実行

modal run trainlora.py

問題がなければtrainlora.pyと同じディレクトリにoutputs/zunko.safetensorsが生成されます。

できたloraを使って画像を生成してみる

作ったLoRA(zunko.safetensors)を使って画像を生成します。

何かと言われたら東北ずん子ですね。LoRAは効いてそうです。さらにタグやLoRAの有無で軽く比較してみます。negative promptは(worst quality, low quality:1.4), nsfw。

タグもちゃんと効いていそうです。

おわり

Modalは多少の制限もありますが、ローカルと同じように作業ができるのが大きな特徴です。 GPUが使えるクラウドサービスを使ってLoRAを作るとき、学習素材のアップロードと出来上がったLoRAのダウンロードが少し手間になりがちですが、 Modalを使うとそれらの手順を一度にまとめることができました。

実際には作成したloraをshared volume(永続ストレージ)に入れたほうが事故が少なそうであったり、学習時のパラメータをもう少しちゃんと設定したほうがよさそうですが、とりあえずLoRAを作ってみるという目的は果たせたのでこれで良しとします。

追記(初期バージョン)

置いておきます。

import os
import sys
import subprocess
import shlex

import modal


# loraの学習に使うモデルのURL
PRETRAINED_MODEL_URL = "https://huggingface.co/WarriorMama777/OrangeMixs/resolve/main/Models/AbyssOrangeMix3/AOM3A1_orangemixs.safetensors"
PRETRAINED_MODEL_NAME = PRETRAINED_MODEL_URL.split("/")[-1]

LOCAL_TRAIN_DIR = "C:/path/to/trainimages/zunkotrain"  # 学習用のフォルダへのパス
TOML_FILE = "trainconfig.toml"  # 設定ファイルの名前

REMOTE_TRAIN_DIR = "/traindata"  #  学習用フォルダのマウント先のディレクトリ
REMOTE_OUTPUT_DIR = "/trainedoutputs"  # 生成されたloraを保存する場所
OUTPUT_NAME = "zunko"  # loraの名前

LOCAL_SAVE_TO = "outputs"  # 成果物のローカルの保存先


stub = modal.Stub("trainlora")
mount = modal.Mount.from_local_dir(LOCAL_TRAIN_DIR, remote_path=REMOTE_TRAIN_DIR)

image = (
    modal.Image.from_dockerhub("python:3.10-slim")
    .apt_install("git", "wget", "libgl1-mesa-dev", "libglib2.0-0", "libsm6", "libxrender1", "libxext6")
    .pip_install(
        "torch==1.13.1+cu117",
        "torchvision==0.14.1+cu117",
        extra_index_url="https://download.pytorch.org/whl/cu117"
    )
    .pip_install(
        "triton",
        pre=True
    )
    .pip_install(
        "accelerate==0.15.0",
        "transformers==4.26.0",
        "ftfy==6.1.1",
        "albumentations==1.3.0",
        "opencv-python==4.7.0.68",
        "einops==0.6.0",
        "diffusers[torch]==0.10.2",
        "pytorch-lightning==1.9.0",
        "bitsandbytes==0.35.0",
        "tensorboard==2.10.1",
        "safetensors==0.2.6",
        "altair==4.2.2",
        "easygui==0.98.3",
        "toml==0.10.2",
        "voluptuous==0.13.1",
        "requests==2.28.2",
        "timm==0.6.12",
        "fairscale==0.4.13",
        "tensorflow==2.10.1",
        "huggingface-hub==0.13.3",
        "xformers==0.0.16rc425",
    )
)


@stub.function(image=image,
               mounts=[mount],
               gpu="T4",
               timeout=10800)
def run_train_lora():
    subprocess.run("accelerate config default --mixed_precision fp16", shell=True)
    subprocess.run("git clone https://github.com/kohya-ss/sd-scripts", shell=True)
    os.chdir("sd-scripts")
    subprocess.run(f"wget {PRETRAINED_MODEL_URL} --no-verbose", shell=True)

    subprocess.run("accelerate launch --num_cpu_threads_per_process 1 "
        "train_network.py "
        f"--pretrained_model_name_or_path={PRETRAINED_MODEL_NAME} "
        f"--dataset_config={REMOTE_TRAIN_DIR}/{TOML_FILE} "
        f"--output_dir={REMOTE_OUTPUT_DIR} "
        f"--output_name={OUTPUT_NAME} "
        "--shuffle_caption --train_batch_size=1 --learning_rate=1e-4 "
        "--max_train_steps=2000 --xformers "
        "--gradient_checkpointing --mixed_precision=fp16 "
        "--save_every_n_epochs=4 --save_precision=fp16 "
        "--save_model_as=safetensors  --network_module=networks.lora", shell=True)

    # print(os.listdir(REMOTE_OUTPUT_DIR))
    os.chdir(REMOTE_OUTPUT_DIR)
    for loraname in os.listdir(REMOTE_OUTPUT_DIR):
        f = open(loraname, "rb")
        lora_data = f.read()
        yield (lora_data, loraname)
        f.close()


@stub.local_entrypoint()
def main():
    os.makedirs(LOCAL_SAVE_TO, exist_ok=True)
    for lora_data, loraname in run_train_lora.call():
        saved_path = os.path.join(LOCAL_SAVE_TO, loraname)
        with open(saved_path, "wb") as f:
            f.write(lora_data)

クラウドでLoRA作成(追加学習)がしたいときModalを利用するとすごく捗る(画像生成AI Stable diffusion)

方針

事前準備

Modalを利用可能にしておく

東北ずん子のAI画像モデル用学習データをダウンロード

学習用フォルダのディレクトリ構成

設定ファイル

pythonを記述して学習を実行

コード

実行

できたloraを使って画像を生成してみる

おわり

追記(初期バージョン)

0 件のコメント:

コメントを投稿

人気の投稿

ブログアーカイブ

クラウドでLoRA作成(追加学習)がしたいときModalを利用するとすごく捗る(画像生成AI Stable diffusion)

方針

事前準備

Modalを利用可能にしておく

東北ずん子のAI画像モデル用学習データをダウンロード

学習用フォルダのディレクトリ構成

設定ファイル

pythonを記述して学習を実行

コード

実行

できたloraを使って画像を生成してみる

おわり

追記(初期バージョン)

0 件のコメント:

コメントを投稿

人気の投稿

ブログ アーカイブ

ブログアーカイブ