# v2: select precision for saved checkpoint
# v3: add logging for tensorboard, fix to shuffle=False in DataLoader (shuffling is in dataset)
# v4: support SD2.0, add lr scheduler options, supports save_every_n_epochs and save_state for DiffUsers model
# v5: refactor to use model_util, support safetensors, add settings to use Diffusers' xformers, add log prefix
# v6: model_util update
# v7: support Diffusers 0.10.0 (v-parameterization training, safetensors in Diffusers) and accelerate 0.15.0, support full path in metadata
# v8: experimental full fp16 training.
# v9: add keep_tokens and save_model_as option, flip augmentation

# このスクリプトのライセンスは、train_dreambooth.pyと同じくApache License 2.0とします
# License:
# Copyright 2022 Kohya S. @kohya_ss
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# License of included scripts:

# Diffusers: ASL 2.0 https://github.com/huggingface/diffusers/blob/main/LICENSE

# Memory efficient attention:
# based on https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/memory_efficient_attention_pytorch/flash_attention.py
# MIT https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/LICENSE

import argparse
import math
import os
import random
import json
import importlib
import time

from tqdm import tqdm
import torch
from accelerate import Accelerator
from accelerate.utils import set_seed
from transformers import CLIPTokenizer
import diffusers
from diffusers import DDPMScheduler, StableDiffusionPipeline
import numpy as np
from einops import rearrange
from torch import einsum

import library.model_util as model_util

# Tokenizer: checkpointから読み込むのではなくあらかじめ提供されているものを使う
TOKENIZER_PATH = 'openai/clip-vit-large-patch14'
V2_STABLE_DIFFUSION_PATH = 'stabilityai/stable-diffusion-2'     # ここからtokenizerだけ使う v2とv2.1はtokenizer仕様は同じ

# checkpointファイル名
EPOCH_STATE_NAME = 'epoch-{:06d}-state'
LAST_STATE_NAME = 'last-state'

LAST_DIFFUSERS_DIR_NAME = 'last'
EPOCH_DIFFUSERS_DIR_NAME = 'epoch-{:06d}'


def collate_fn(examples):
    return examples[0]


class FineTuningDataset(torch.utils.data.Dataset):
    def __init__(
        self,
        metadata,
        train_data_dir,
        batch_size,
        tokenizer,
        max_token_length,
        shuffle_caption,
        shuffle_keep_tokens,
        dataset_repeats,
        debug,
    ) -> None:
        super().__init__()

        self.metadata = metadata
        self.train_data_dir = train_data_dir
        self.batch_size = batch_size
        self.tokenizer: CLIPTokenizer = tokenizer
        self.max_token_length = max_token_length
        self.shuffle_caption = shuffle_caption
        self.shuffle_keep_tokens = shuffle_keep_tokens
        self.debug = debug

        self.tokenizer_max_length = (
            self.tokenizer.model_max_length
            if max_token_length is None
            else max_token_length + 2
        )

        print('make buckets')

        # 最初に数を数える
        self.bucket_resos = set()
        for img_md in metadata.values():
            if 'train_resolution' in img_md:
                self.bucket_resos.add(tuple(img_md['train_resolution']))
        self.bucket_resos = list(self.bucket_resos)
        self.bucket_resos.sort()
        print(f'number of buckets: {len(self.bucket_resos)}')

        reso_to_index = {}
        for i, reso in enumerate(self.bucket_resos):
            reso_to_index[reso] = i

        # bucketに割り当てていく
        self.buckets = [[] for _ in range(len(self.bucket_resos))]
        n = 1 if dataset_repeats is None else dataset_repeats
        images_count = 0
        for image_key, img_md in metadata.items():
            if 'train_resolution' not in img_md:
                continue
            if not os.path.exists(self.image_key_to_npz_file(image_key)):
                continue

            reso = tuple(img_md['train_resolution'])
            for _ in range(n):
                self.buckets[reso_to_index[reso]].append(image_key)
            images_count += n

        # 参照用indexを作る
        self.buckets_indices = []
        for bucket_index, bucket in enumerate(self.buckets):
            batch_count = int(math.ceil(len(bucket) / self.batch_size))
            for batch_index in range(batch_count):
                self.buckets_indices.append((bucket_index, batch_index))

        self.shuffle_buckets()
        self._length = len(self.buckets_indices)
        self.images_count = images_count

    def show_buckets(self):
        for i, (reso, bucket) in enumerate(
            zip(self.bucket_resos, self.buckets)
        ):
            print(f'bucket {i}: resolution {reso}, count: {len(bucket)}')

    def shuffle_buckets(self):
        random.shuffle(self.buckets_indices)
        for bucket in self.buckets:
            random.shuffle(bucket)

    def image_key_to_npz_file(self, image_key):
        npz_file_norm = os.path.splitext(image_key)[0] + '.npz'
        if os.path.exists(npz_file_norm):
            if random.random() < 0.5:
                npz_file_flip = os.path.splitext(image_key)[0] + '_flip.npz'
                if os.path.exists(npz_file_flip):
                    return npz_file_flip
            return npz_file_norm

        npz_file_norm = os.path.join(self.train_data_dir, image_key + '.npz')
        if random.random() < 0.5:
            npz_file_flip = os.path.join(
                self.train_data_dir, image_key + '_flip.npz'
            )
            if os.path.exists(npz_file_flip):
                return npz_file_flip
        return npz_file_norm

    def load_latent(self, image_key):
        return np.load(self.image_key_to_npz_file(image_key))['arr_0']

    def __len__(self):
        return self._length

    def __getitem__(self, index):
        if index == 0:
            self.shuffle_buckets()

        bucket = self.buckets[self.buckets_indices[index][0]]
        image_index = self.buckets_indices[index][1] * self.batch_size

        input_ids_list = []
        latents_list = []
        captions = []
        for image_key in bucket[image_index : image_index + self.batch_size]:
            img_md = self.metadata[image_key]
            caption = img_md.get('caption')
            tags = img_md.get('tags')

            if caption is None:
                caption = tags
            elif tags is not None and len(tags) > 0:
                caption = caption + ', ' + tags
            assert (
                caption is not None and len(caption) > 0
            ), f'caption or tag is required / キャプションまたはタグは必須です:{image_key}'

            latents = self.load_latent(image_key)

            if self.shuffle_caption:
                tokens = caption.strip().split(',')
                if self.shuffle_keep_tokens is None:
                    random.shuffle(tokens)
                else:
                    if len(tokens) > self.shuffle_keep_tokens:
                        keep_tokens = tokens[: self.shuffle_keep_tokens]
                        tokens = tokens[self.shuffle_keep_tokens :]
                        random.shuffle(tokens)
                        tokens = keep_tokens + tokens
                caption = ','.join(tokens).strip()

            captions.append(caption)

            input_ids = self.tokenizer(
                caption,
                padding='max_length',
                truncation=True,
                max_length=self.tokenizer_max_length,
                return_tensors='pt',
            ).input_ids

            if self.tokenizer_max_length > self.tokenizer.model_max_length:
                input_ids = input_ids.squeeze(0)
                iids_list = []
                if self.tokenizer.pad_token_id == self.tokenizer.eos_token_id:
                    # v1
                    # 77以上の時は "<BOS> .... <EOS> <EOS> <EOS>" でトータル227とかになっているので、"<BOS>...<EOS>"の三連に変換する
                    # 1111氏のやつは , で区切る、とかしているようだが　とりあえず単純に
                    for i in range(
                        1,
                        self.tokenizer_max_length
                        - self.tokenizer.model_max_length
                        + 2,
                        self.tokenizer.model_max_length - 2,
                    ):  # (1, 152, 75)
                        ids_chunk = (
                            input_ids[0].unsqueeze(0),
                            input_ids[
                                i : i + self.tokenizer.model_max_length - 2
                            ],
                            input_ids[-1].unsqueeze(0),
                        )
                        ids_chunk = torch.cat(ids_chunk)
                        iids_list.append(ids_chunk)
                else:
                    # v2
                    # 77以上の時は "<BOS> .... <EOS> <PAD> <PAD>..." でトータル227とかになっているので、"<BOS>...<EOS> <PAD> <PAD> ..."の三連に変換する
                    for i in range(
                        1,
                        self.tokenizer_max_length
                        - self.tokenizer.model_max_length
                        + 2,
                        self.tokenizer.model_max_length - 2,
                    ):
                        ids_chunk = (
                            input_ids[0].unsqueeze(0),  # BOS
                            input_ids[
                                i : i + self.tokenizer.model_max_length - 2
                            ],
                            input_ids[-1].unsqueeze(0),
                        )      # PAD or EOS
                        ids_chunk = torch.cat(ids_chunk)

                        # 末尾が <EOS> <PAD> または <PAD> <PAD> の場合は、何もしなくてよい
                        # 末尾が x <PAD/EOS> の場合は末尾を <EOS> に変える（x <EOS> なら結果的に変化なし）
                        if (
                            ids_chunk[-2] != self.tokenizer.eos_token_id
                            and ids_chunk[-2] != self.tokenizer.pad_token_id
                        ):
                            ids_chunk[-1] = self.tokenizer.eos_token_id
                        # 先頭が <BOS> <PAD> ... の場合は <BOS> <EOS> <PAD> ... に変える
                        if ids_chunk[1] == self.tokenizer.pad_token_id:
                            ids_chunk[1] = self.tokenizer.eos_token_id

                        iids_list.append(ids_chunk)

                input_ids = torch.stack(iids_list)      # 3,77

            input_ids_list.append(input_ids)
            latents_list.append(torch.FloatTensor(latents))

        example = {}
        example['input_ids'] = torch.stack(input_ids_list)
        example['latents'] = torch.stack(latents_list)
        if self.debug:
            example['image_keys'] = bucket[
                image_index : image_index + self.batch_size
            ]
            example['captions'] = captions
        return example


def save_hypernetwork(output_file, hypernetwork):
    state_dict = hypernetwork.get_state_dict()
    torch.save(state_dict, output_file)


def train(args):
    fine_tuning = (
        args.hypernetwork_module is None
    )            # fine tuning or hypernetwork training

    # その他のオプション設定を確認する
    if args.v_parameterization and not args.v2:
        print(
            'v_parameterization should be with v2 / v1でv_parameterizationを使用することは想定されていません'
        )
    if args.v2 and args.clip_skip is not None:
        print(
            'v2 with clip_skip will be unexpected / v2でclip_skipを使用することは想定されていません'
        )

    # モデル形式のオプション設定を確認する
    load_stable_diffusion_format = os.path.isfile(
        args.pretrained_model_name_or_path
    )

    if load_stable_diffusion_format:
        src_stable_diffusion_ckpt = args.pretrained_model_name_or_path
        src_diffusers_model_path = None
    else:
        src_stable_diffusion_ckpt = None
        src_diffusers_model_path = args.pretrained_model_name_or_path

    if args.save_model_as is None:
        save_stable_diffusion_format = load_stable_diffusion_format
        use_safetensors = args.use_safetensors
    else:
        save_stable_diffusion_format = (
            args.save_model_as.lower() == 'ckpt'
            or args.save_model_as.lower() == 'safetensors'
        )
        use_safetensors = args.use_safetensors or (
            'safetensors' in args.save_model_as.lower()
        )

    # 乱数系列を初期化する
    if args.seed is not None:
        set_seed(args.seed)

    # メタデータを読み込む
    if os.path.exists(args.in_json):
        print(f'loading existing metadata: {args.in_json}')
        with open(args.in_json, 'rt', encoding='utf-8') as f:
            metadata = json.load(f)
    else:
        print(f'no metadata / メタデータファイルがありません: {args.in_json}')
        return

    # tokenizerを読み込む
    print('prepare tokenizer')
    if args.v2:
        tokenizer = CLIPTokenizer.from_pretrained(
            V2_STABLE_DIFFUSION_PATH, subfolder='tokenizer'
        )
    else:
        tokenizer = CLIPTokenizer.from_pretrained(TOKENIZER_PATH)

    if args.max_token_length is not None:
        print(f'update token length: {args.max_token_length}')

    # datasetを用意する
    print('prepare dataset')
    train_dataset = FineTuningDataset(
        metadata,
        args.train_data_dir,
        args.train_batch_size,
        tokenizer,
        args.max_token_length,
        args.shuffle_caption,
        args.keep_tokens,
        args.dataset_repeats,
        args.debug_dataset,
    )

    print(f'Total dataset length / データセットの長さ: {len(train_dataset)}')
    print(f'Total images / 画像数: {train_dataset.images_count}')

    if len(train_dataset) == 0:
        print(
            'No data found. Please verify the metadata file and train_data_dir option. / 画像がありません。メタデータおよびtrain_data_dirオプションを確認してください。'
        )
        return

    if args.debug_dataset:
        train_dataset.show_buckets()
        i = 0
        for example in train_dataset:
            print(f"image: {example['image_keys']}")
            print(f"captions: {example['captions']}")
            print(f"latents: {example['latents'].shape}")
            print(f"input_ids: {example['input_ids'].shape}")
            print(example['input_ids'])
            i += 1
            if i >= 8:
                break
        return

    # acceleratorを準備する
    print('prepare accelerator')
    if args.logging_dir is None:
        log_with = None
        logging_dir = None
    else:
        log_with = 'tensorboard'
        log_prefix = '' if args.log_prefix is None else args.log_prefix
        logging_dir = (
            args.logging_dir
            + '/'
            + log_prefix
            + time.strftime('%Y%m%d%H%M%S', time.localtime())
        )
    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        mixed_precision=args.mixed_precision,
        log_with=log_with,
        logging_dir=logging_dir,
    )

    # accelerateの互換性問題を解決する
    accelerator_0_15 = True
    try:
        accelerator.unwrap_model('dummy', True)
        print('Using accelerator 0.15.0 or above.')
    except TypeError:
        accelerator_0_15 = False

    def unwrap_model(model):
        if accelerator_0_15:
            return accelerator.unwrap_model(model, True)
        return accelerator.unwrap_model(model)

    # mixed precisionに対応した型を用意しておき適宜castする
    weight_dtype = torch.float32
    if args.mixed_precision == 'fp16':
        weight_dtype = torch.float16
    elif args.mixed_precision == 'bf16':
        weight_dtype = torch.bfloat16

    save_dtype = None
    if args.save_precision == 'fp16':
        save_dtype = torch.float16
    elif args.save_precision == 'bf16':
        save_dtype = torch.bfloat16
    elif args.save_precision == 'float':
        save_dtype = torch.float32

    # モデルを読み込む
    if load_stable_diffusion_format:
        print('load StableDiffusion checkpoint')
        (
            text_encoder,
            vae,
            unet,
        ) = model_util.load_models_from_stable_diffusion_checkpoint(
            args.v2, args.pretrained_model_name_or_path
        )
    else:
        print('load Diffusers pretrained models')
        pipe = StableDiffusionPipeline.from_pretrained(
            args.pretrained_model_name_or_path,
            tokenizer=None,
            safety_checker=None,
        )
        # , torch_dtype=weight_dtype) ここでtorch_dtypeを指定すると学習時にエラーになる
        text_encoder = pipe.text_encoder
        unet = pipe.unet
        vae = pipe.vae
        del pipe
    vae.to('cpu')                     # 保存時にしか使わないので、メモリを開けるためCPUに移しておく

    # Diffusers版のxformers使用フラグを設定する関数
    def set_diffusers_xformers_flag(model, valid):
        #   model.set_use_memory_efficient_attention_xformers(valid)            # 次のリリースでなくなりそう
        # pipeが自動で再帰的にset_use_memory_efficient_attention_xformersを探すんだって(;´Д｀)
        # U-Netだけ使う時にはどうすればいいのか……仕方ないからコピって使うか
        # 0.10.2でなんか巻き戻って個別に指定するようになった(;^ω^)

        # Recursively walk through all the children.
        # Any children which exposes the set_use_memory_efficient_attention_xformers method
        # gets the message
        def fn_recursive_set_mem_eff(module: torch.nn.Module):
            if hasattr(module, 'set_use_memory_efficient_attention_xformers'):
                module.set_use_memory_efficient_attention_xformers(valid)

            for child in module.children():
                fn_recursive_set_mem_eff(child)

        fn_recursive_set_mem_eff(model)

    # モデルに xformers とか memory efficient attention を組み込む
    if args.diffusers_xformers:
        print('Use xformers by Diffusers')
        set_diffusers_xformers_flag(unet, True)
    else:
        # Windows版のxformersはfloatで学習できないのでxformersを使わない設定も可能にしておく必要がある
        print("Disable Diffusers' xformers")
        set_diffusers_xformers_flag(unet, False)
        replace_unet_modules(unet, args.mem_eff_attn, args.xformers)

    if not fine_tuning:
        # Hypernetwork
        print('import hypernetwork module:', args.hypernetwork_module)
        hyp_module = importlib.import_module(args.hypernetwork_module)

        hypernetwork = hyp_module.Hypernetwork()

        if args.hypernetwork_weights is not None:
            print('load hypernetwork weights from:', args.hypernetwork_weights)
            hyp_sd = torch.load(args.hypernetwork_weights, map_location='cpu')
            success = hypernetwork.load_from_state_dict(hyp_sd)
            assert success, 'hypernetwork weights loading failed.'

        print('apply hypernetwork')
        hypernetwork.apply_to_diffusers(None, text_encoder, unet)

    # 学習を準備する：モデルを適切な状態にする
    training_models = []
    if fine_tuning:
        if args.gradient_checkpointing:
            unet.enable_gradient_checkpointing()
        training_models.append(unet)

        if args.train_text_encoder:
            print('enable text encoder training')
            if args.gradient_checkpointing:
                text_encoder.gradient_checkpointing_enable()
            training_models.append(text_encoder)
        else:
            text_encoder.to(accelerator.device, dtype=weight_dtype)
            text_encoder.requires_grad_(False)             # text encoderは学習しない
            text_encoder.eval()
    else:
        unet.to(
            accelerator.device
        )  # , dtype=weight_dtype)     # dtypeを指定すると学習できない
        unet.requires_grad_(False)
        unet.eval()
        text_encoder.to(accelerator.device, dtype=weight_dtype)
        text_encoder.requires_grad_(False)
        text_encoder.eval()
        training_models.append(hypernetwork)

    for m in training_models:
        m.requires_grad_(True)
    params = []
    for m in training_models:
        params.extend(m.parameters())
    params_to_optimize = params

    # 学習に必要なクラスを準備する
    print('prepare optimizer, data loader etc.')

    # 8-bit Adamを使う
    if args.use_8bit_adam:
        try:
            import bitsandbytes as bnb
        except ImportError:
            raise ImportError(
                'No bitsand bytes / bitsandbytesがインストールされていないようです'
            )
        print('use 8-bit Adam optimizer')
        optimizer_class = bnb.optim.AdamW8bit
    else:
        optimizer_class = torch.optim.AdamW

    # betaやweight decayはdiffusers DreamBoothもDreamBooth SDもデフォルト値のようなのでオプションはとりあえず省略
    optimizer = optimizer_class(params_to_optimize, lr=args.learning_rate)

    # dataloaderを準備する
    # DataLoaderのプロセス数：0はメインプロセスになる
    n_workers = min(8, os.cpu_count() - 1)      # cpu_count-1 ただし最大8
    train_dataloader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=1,
        shuffle=False,
        collate_fn=collate_fn,
        num_workers=n_workers,
    )

    # lr schedulerを用意する
    lr_scheduler = diffusers.optimization.get_scheduler(
        args.lr_scheduler,
        optimizer,
        num_warmup_steps=args.lr_warmup_steps,
        num_training_steps=args.max_train_steps
        * args.gradient_accumulation_steps,
    )

    # acceleratorがなんかよろしくやってくれるらしい
    if args.full_fp16:
        assert (
            args.mixed_precision == 'fp16'
        ), "full_fp16 requires mixed precision='fp16' / full_fp16を使う場合はmixed_precision='fp16'を指定してください。"
        print('enable full fp16 training.')

    if fine_tuning:
        # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
        if args.full_fp16:
            unet.to(weight_dtype)
            text_encoder.to(weight_dtype)

        if args.train_text_encoder:
            (
                unet,
                text_encoder,
                optimizer,
                train_dataloader,
                lr_scheduler,
            ) = accelerator.prepare(
                unet, text_encoder, optimizer, train_dataloader, lr_scheduler
            )
        else:
            (
                unet,
                optimizer,
                train_dataloader,
                lr_scheduler,
            ) = accelerator.prepare(
                unet, optimizer, train_dataloader, lr_scheduler
            )
    else:
        if args.full_fp16:
            unet.to(weight_dtype)
            hypernetwork.to(weight_dtype)

        (
            unet,
            hypernetwork,
            optimizer,
            train_dataloader,
            lr_scheduler,
        ) = accelerator.prepare(
            unet, hypernetwork, optimizer, train_dataloader, lr_scheduler
        )

    # 実験的機能：勾配も含めたfp16学習を行う　PyTorchにパッチを当ててfp16でのgrad scaleを有効にする
    if args.full_fp16:
        org_unscale_grads = accelerator.scaler._unscale_grads_

        def _unscale_grads_replacer(
            optimizer, inv_scale, found_inf, allow_fp16
        ):
            return org_unscale_grads(optimizer, inv_scale, found_inf, True)

        accelerator.scaler._unscale_grads_ = _unscale_grads_replacer

    # TODO accelerateのconfigに指定した型とオプション指定の型とをチェックして異なれば警告を出す

    # resumeする
    if args.resume is not None:
        print(f'resume training from state: {args.resume}')
        accelerator.load_state(args.resume)

    # epoch数を計算する
    num_update_steps_per_epoch = math.ceil(
        len(train_dataloader) / args.gradient_accumulation_steps
    )
    num_train_epochs = math.ceil(
        args.max_train_steps / num_update_steps_per_epoch
    )

    # 学習する
    total_batch_size = (
        args.train_batch_size
        * accelerator.num_processes
        * args.gradient_accumulation_steps
    )
    print('running training / 学習開始')
    print(f'  num examples / サンプル数: {train_dataset.images_count}')
    print(f'  num batches per epoch / 1epochのバッチ数: {len(train_dataloader)}')
    print(f'  num epochs / epoch数: {num_train_epochs}')
    print(f'  batch size per device / バッチサイズ: {args.train_batch_size}')
    print(
        f'  total train batch size (with parallel & distributed & accumulation) / 総バッチサイズ（並列学習、勾配合計含む）: {total_batch_size}'
    )
    print(
        f'  gradient ccumulation steps / 勾配を合計するステップ数 = {args.gradient_accumulation_steps}'
    )
    print(f'  total optimization steps / 学習ステップ数: {args.max_train_steps}')

    progress_bar = tqdm(
        range(args.max_train_steps),
        smoothing=0,
        disable=not accelerator.is_local_main_process,
        desc='steps',
    )
    global_step = 0

    # v4で更新：clip_sample=Falseに
    # Diffusersのtrain_dreambooth.pyがconfigから持ってくるように変更されたので、clip_sample=Falseになるため、それに合わせる
    # 既存の1.4/1.5/2.0/2.1はすべてschdulerのconfigは（クラス名を除いて）同じ
    # よくソースを見たら学習時はclip_sampleは関係ないや(;'∀')
    noise_scheduler = DDPMScheduler(
        beta_start=0.00085,
        beta_end=0.012,
        beta_schedule='scaled_linear',
        num_train_timesteps=1000,
        clip_sample=False,
    )

    if accelerator.is_main_process:
        accelerator.init_trackers(
            'finetuning' if fine_tuning else 'hypernetwork'
        )

    # 以下 train_dreambooth.py からほぼコピペ
    for epoch in range(num_train_epochs):
        print(f'epoch {epoch+1}/{num_train_epochs}')
        for m in training_models:
            m.train()

        loss_total = 0
        for step, batch in enumerate(train_dataloader):
            with accelerator.accumulate(
                training_models[0]
            ):  # 複数モデルに対応していない模様だがとりあえずこうしておく
                latents = batch['latents'].to(accelerator.device)
                latents = latents * 0.18215
                b_size = latents.shape[0]

                # with torch.no_grad():
                with torch.set_grad_enabled(args.train_text_encoder):
                    # Get the text embedding for conditioning
                    input_ids = batch['input_ids'].to(accelerator.device)
                    input_ids = input_ids.reshape(
                        (-1, tokenizer.model_max_length)
                    )     # batch_size*3, 77

                    if args.clip_skip is None:
                        encoder_hidden_states = text_encoder(input_ids)[0]
                    else:
                        enc_out = text_encoder(
                            input_ids,
                            output_hidden_states=True,
                            return_dict=True,
                        )
                        encoder_hidden_states = enc_out['hidden_states'][
                            -args.clip_skip
                        ]
                        encoder_hidden_states = (
                            text_encoder.text_model.final_layer_norm(
                                encoder_hidden_states
                            )
                        )

                    # bs*3, 77, 768 or 1024
                    encoder_hidden_states = encoder_hidden_states.reshape(
                        (b_size, -1, encoder_hidden_states.shape[-1])
                    )

                    if args.max_token_length is not None:
                        if args.v2:
                            # v2: <BOS>...<EOS> <PAD> ... の三連を <BOS>...<EOS> <PAD> ... へ戻す　正直この実装でいいのかわからん
                            states_list = [
                                encoder_hidden_states[:, 0].unsqueeze(1)
                            ]                              # <BOS>
                            for i in range(
                                1,
                                args.max_token_length,
                                tokenizer.model_max_length,
                            ):
                                chunk = encoder_hidden_states[
                                    :, i : i + tokenizer.model_max_length - 2
                                ]              # <BOS> の後から 最後の前まで
                                if i > 0:
                                    for j in range(len(chunk)):
                                        if (
                                            input_ids[j, 1]
                                            == tokenizer.eos_token
                                        ):                                      # 空、つまり <BOS> <EOS> <PAD> ...のパターン
                                            chunk[j, 0] = chunk[
                                                j, 1
                                            ]                                                     # 次の <PAD> の値をコピーする
                                states_list.append(
                                    chunk
                                )  # <BOS> の後から <EOS> の前まで
                            states_list.append(
                                encoder_hidden_states[:, -1].unsqueeze(1)
                            )                         # <EOS> か <PAD> のどちらか
                            encoder_hidden_states = torch.cat(
                                states_list, dim=1
                            )
                        else:
                            # v1: <BOS>...<EOS> の三連を <BOS>...<EOS> へ戻す
                            states_list = [
                                encoder_hidden_states[:, 0].unsqueeze(1)
                            ]                              # <BOS>
                            for i in range(
                                1,
                                args.max_token_length,
                                tokenizer.model_max_length,
                            ):
                                states_list.append(
                                    encoder_hidden_states[
                                        :,
                                        i : i + tokenizer.model_max_length - 2,
                                    ]
                                )  # <BOS> の後から <EOS> の前まで
                            states_list.append(
                                encoder_hidden_states[:, -1].unsqueeze(1)
                            )                         # <EOS>
                            encoder_hidden_states = torch.cat(
                                states_list, dim=1
                            )

                # Sample noise that we'll add to the latents
                noise = torch.randn_like(latents, device=latents.device)

                # Sample a random timestep for each image
                timesteps = torch.randint(
                    0,
                    noise_scheduler.config.num_train_timesteps,
                    (b_size,),
                    device=latents.device,
                )
                timesteps = timesteps.long()

                # Add noise to the latents according to the noise magnitude at each timestep
                # (this is the forward diffusion process)
                noisy_latents = noise_scheduler.add_noise(
                    latents, noise, timesteps
                )

                # Predict the noise residual
                noise_pred = unet(
                    noisy_latents, timesteps, encoder_hidden_states
                ).sample

                if args.v_parameterization:
                    # v-parameterization training
                    # Diffusers 0.10.0からv_parameterizationの学習に対応したのでそちらを使う
                    target = noise_scheduler.get_velocity(
                        latents, noise, timesteps
                    )
                else:
                    target = noise

                loss = torch.nn.functional.mse_loss(
                    noise_pred.float(), target.float(), reduction='mean'
                )

                accelerator.backward(loss)
                if accelerator.sync_gradients:
                    params_to_clip = []
                    for m in training_models:
                        params_to_clip.extend(m.parameters())
                    accelerator.clip_grad_norm_(
                        params_to_clip, 1.0
                    )  # args.max_grad_norm)

                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad(set_to_none=True)

            # Checks if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:
                progress_bar.update(1)
                global_step += 1

            current_loss = loss.detach().item()        # 平均なのでbatch sizeは関係ないはず
            if args.logging_dir is not None:
                logs = {
                    'loss': current_loss,
                    'lr': lr_scheduler.get_last_lr()[0],
                }
                accelerator.log(logs, step=global_step)

            loss_total += current_loss
            avr_loss = loss_total / (step + 1)
            logs = {'loss': avr_loss}  # , "lr": lr_scheduler.get_last_lr()[0]}
            progress_bar.set_postfix(**logs)

            if global_step >= args.max_train_steps:
                break

        if args.logging_dir is not None:
            logs = {'epoch_loss': loss_total / len(train_dataloader)}
            accelerator.log(logs, step=epoch + 1)

        accelerator.wait_for_everyone()

        if args.save_every_n_epochs is not None:
            if (epoch + 1) % args.save_every_n_epochs == 0 and (
                epoch + 1
            ) < num_train_epochs:
                print('saving checkpoint.')
                os.makedirs(args.output_dir, exist_ok=True)
                ckpt_file = os.path.join(
                    args.output_dir,
                    model_util.get_epoch_ckpt_name(use_safetensors, epoch + 1),
                )

                if fine_tuning:
                    if save_stable_diffusion_format:
                        model_util.save_stable_diffusion_checkpoint(
                            args.v2,
                            ckpt_file,
                            unwrap_model(text_encoder),
                            unwrap_model(unet),
                            src_stable_diffusion_ckpt,
                            epoch + 1,
                            global_step,
                            save_dtype,
                            vae,
                        )
                    else:
                        out_dir = os.path.join(
                            args.output_dir,
                            EPOCH_DIFFUSERS_DIR_NAME.format(epoch + 1),
                        )
                        os.makedirs(out_dir, exist_ok=True)
                        model_util.save_diffusers_checkpoint(
                            args.v2,
                            out_dir,
                            unwrap_model(text_encoder),
                            unwrap_model(unet),
                            src_diffusers_model_path,
                            vae=vae,
                            use_safetensors=use_safetensors,
                        )
                else:
                    save_hypernetwork(ckpt_file, unwrap_model(hypernetwork))

                if args.save_state:
                    print('saving state.')
                    accelerator.save_state(
                        os.path.join(
                            args.output_dir, EPOCH_STATE_NAME.format(epoch + 1)
                        )
                    )

    is_main_process = accelerator.is_main_process
    if is_main_process:
        if fine_tuning:
            unet = unwrap_model(unet)
            text_encoder = unwrap_model(text_encoder)
        else:
            hypernetwork = unwrap_model(hypernetwork)

    accelerator.end_training()

    if args.save_state:
        print('saving last state.')
        accelerator.save_state(os.path.join(args.output_dir, LAST_STATE_NAME))

    del accelerator                         # この後メモリを使うのでこれは消す

    if is_main_process:
        os.makedirs(args.output_dir, exist_ok=True)
        ckpt_file = os.path.join(
            args.output_dir, model_util.get_last_ckpt_name(use_safetensors)
        )

        if fine_tuning:
            if save_stable_diffusion_format:
                print(
                    f'save trained model as StableDiffusion checkpoint to {ckpt_file}'
                )
                model_util.save_stable_diffusion_checkpoint(
                    args.v2,
                    ckpt_file,
                    text_encoder,
                    unet,
                    src_stable_diffusion_ckpt,
                    epoch,
                    global_step,
                    save_dtype,
                    vae,
                )
            else:
                # Create the pipeline using using the trained modules and save it.
                print(f'save trained model as Diffusers to {args.output_dir}')
                out_dir = os.path.join(
                    args.output_dir, LAST_DIFFUSERS_DIR_NAME
                )
                os.makedirs(out_dir, exist_ok=True)
                model_util.save_diffusers_checkpoint(
                    args.v2,
                    out_dir,
                    text_encoder,
                    unet,
                    src_diffusers_model_path,
                    vae=vae,
                    use_safetensors=use_safetensors,
                )
        else:
            print(f'save trained model to {ckpt_file}')
            save_hypernetwork(ckpt_file, hypernetwork)

        print('model saved.')


# region モジュール入れ替え部
"""
高速化のためのモジュール入れ替え
"""

# FlashAttentionを使うCrossAttention
# based on https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/memory_efficient_attention_pytorch/flash_attention.py
# LICENSE MIT https://github.com/lucidrains/memory-efficient-attention-pytorch/blob/main/LICENSE

# constants

EPSILON = 1e-6

# helper functions


def exists(val):
    return val is not None


def default(val, d):
    return val if exists(val) else d


# flash attention forwards and backwards

# https://arxiv.org/abs/2205.14135


class FlashAttentionFunction(torch.autograd.function.Function):
    @staticmethod
    @torch.no_grad()
    def forward(ctx, q, k, v, mask, causal, q_bucket_size, k_bucket_size):
        """Algorithm 2 in the paper"""

        device = q.device
        dtype = q.dtype
        max_neg_value = -torch.finfo(q.dtype).max
        qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)

        o = torch.zeros_like(q)
        all_row_sums = torch.zeros(
            (*q.shape[:-1], 1), dtype=dtype, device=device
        )
        all_row_maxes = torch.full(
            (*q.shape[:-1], 1), max_neg_value, dtype=dtype, device=device
        )

        scale = q.shape[-1] ** -0.5

        if not exists(mask):
            mask = (None,) * math.ceil(q.shape[-2] / q_bucket_size)
        else:
            mask = rearrange(mask, 'b n -> b 1 1 n')
            mask = mask.split(q_bucket_size, dim=-1)

        row_splits = zip(
            q.split(q_bucket_size, dim=-2),
            o.split(q_bucket_size, dim=-2),
            mask,
            all_row_sums.split(q_bucket_size, dim=-2),
            all_row_maxes.split(q_bucket_size, dim=-2),
        )

        for ind, (qc, oc, row_mask, row_sums, row_maxes) in enumerate(
            row_splits
        ):
            q_start_index = ind * q_bucket_size - qk_len_diff

            col_splits = zip(
                k.split(k_bucket_size, dim=-2),
                v.split(k_bucket_size, dim=-2),
            )

            for k_ind, (kc, vc) in enumerate(col_splits):
                k_start_index = k_ind * k_bucket_size

                attn_weights = (
                    einsum('... i d, ... j d -> ... i j', qc, kc) * scale
                )

                if exists(row_mask):
                    attn_weights.masked_fill_(~row_mask, max_neg_value)

                if causal and q_start_index < (
                    k_start_index + k_bucket_size - 1
                ):
                    causal_mask = torch.ones(
                        (qc.shape[-2], kc.shape[-2]),
                        dtype=torch.bool,
                        device=device,
                    ).triu(q_start_index - k_start_index + 1)
                    attn_weights.masked_fill_(causal_mask, max_neg_value)

                block_row_maxes = attn_weights.amax(dim=-1, keepdims=True)
                attn_weights -= block_row_maxes
                exp_weights = torch.exp(attn_weights)

                if exists(row_mask):
                    exp_weights.masked_fill_(~row_mask, 0.0)

                block_row_sums = exp_weights.sum(dim=-1, keepdims=True).clamp(
                    min=EPSILON
                )

                new_row_maxes = torch.maximum(block_row_maxes, row_maxes)

                exp_values = einsum(
                    '... i j, ... j d -> ... i d', exp_weights, vc
                )

                exp_row_max_diff = torch.exp(row_maxes - new_row_maxes)
                exp_block_row_max_diff = torch.exp(
                    block_row_maxes - new_row_maxes
                )

                new_row_sums = (
                    exp_row_max_diff * row_sums
                    + exp_block_row_max_diff * block_row_sums
                )

                oc.mul_((row_sums / new_row_sums) * exp_row_max_diff).add_(
                    (exp_block_row_max_diff / new_row_sums) * exp_values
                )

                row_maxes.copy_(new_row_maxes)
                row_sums.copy_(new_row_sums)

        ctx.args = (causal, scale, mask, q_bucket_size, k_bucket_size)
        ctx.save_for_backward(q, k, v, o, all_row_sums, all_row_maxes)

        return o

    @staticmethod
    @torch.no_grad()
    def backward(ctx, do):
        """Algorithm 4 in the paper"""

        causal, scale, mask, q_bucket_size, k_bucket_size = ctx.args
        q, k, v, o, l, m = ctx.saved_tensors

        device = q.device

        max_neg_value = -torch.finfo(q.dtype).max
        qk_len_diff = max(k.shape[-2] - q.shape[-2], 0)

        dq = torch.zeros_like(q)
        dk = torch.zeros_like(k)
        dv = torch.zeros_like(v)

        row_splits = zip(
            q.split(q_bucket_size, dim=-2),
            o.split(q_bucket_size, dim=-2),
            do.split(q_bucket_size, dim=-2),
            mask,
            l.split(q_bucket_size, dim=-2),
            m.split(q_bucket_size, dim=-2),
            dq.split(q_bucket_size, dim=-2),
        )

        for ind, (qc, oc, doc, row_mask, lc, mc, dqc) in enumerate(row_splits):
            q_start_index = ind * q_bucket_size - qk_len_diff

            col_splits = zip(
                k.split(k_bucket_size, dim=-2),
                v.split(k_bucket_size, dim=-2),
                dk.split(k_bucket_size, dim=-2),
                dv.split(k_bucket_size, dim=-2),
            )

            for k_ind, (kc, vc, dkc, dvc) in enumerate(col_splits):
                k_start_index = k_ind * k_bucket_size

                attn_weights = (
                    einsum('... i d, ... j d -> ... i j', qc, kc) * scale
                )

                if causal and q_start_index < (
                    k_start_index + k_bucket_size - 1
                ):
                    causal_mask = torch.ones(
                        (qc.shape[-2], kc.shape[-2]),
                        dtype=torch.bool,
                        device=device,
                    ).triu(q_start_index - k_start_index + 1)
                    attn_weights.masked_fill_(causal_mask, max_neg_value)

                exp_attn_weights = torch.exp(attn_weights - mc)

                if exists(row_mask):
                    exp_attn_weights.masked_fill_(~row_mask, 0.0)

                p = exp_attn_weights / lc

                dv_chunk = einsum('... i j, ... i d -> ... j d', p, doc)
                dp = einsum('... i d, ... j d -> ... i j', doc, vc)

                D = (doc * oc).sum(dim=-1, keepdims=True)
                ds = p * scale * (dp - D)

                dq_chunk = einsum('... i j, ... j d -> ... i d', ds, kc)
                dk_chunk = einsum('... i j, ... i d -> ... j d', ds, qc)

                dqc.add_(dq_chunk)
                dkc.add_(dk_chunk)
                dvc.add_(dv_chunk)

        return dq, dk, dv, None, None, None, None


def replace_unet_modules(
    unet: diffusers.models.unet_2d_condition.UNet2DConditionModel,
    mem_eff_attn,
    xformers,
):
    if mem_eff_attn:
        replace_unet_cross_attn_to_memory_efficient()
    elif xformers:
        replace_unet_cross_attn_to_xformers()


def replace_unet_cross_attn_to_memory_efficient():
    print(
        'Replace CrossAttention.forward to use FlashAttention (not xformers)'
    )
    flash_func = FlashAttentionFunction

    def forward_flash_attn(self, x, context=None, mask=None):
        q_bucket_size = 512
        k_bucket_size = 1024

        h = self.heads
        q = self.to_q(x)

        context = context if context is not None else x
        context = context.to(x.dtype)

        if hasattr(self, 'hypernetwork') and self.hypernetwork is not None:
            context_k, context_v = self.hypernetwork.forward(x, context)
            context_k = context_k.to(x.dtype)
            context_v = context_v.to(x.dtype)
        else:
            context_k = context
            context_v = context

        k = self.to_k(context_k)
        v = self.to_v(context_v)
        del context, x

        q, k, v = map(
            lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v)
        )

        out = flash_func.apply(
            q, k, v, mask, False, q_bucket_size, k_bucket_size
        )

        out = rearrange(out, 'b h n d -> b n (h d)')

        # diffusers 0.7.0~  わざわざ変えるなよ (;´Д｀)
        out = self.to_out[0](out)
        out = self.to_out[1](out)
        return out

    diffusers.models.attention.CrossAttention.forward = forward_flash_attn


def replace_unet_cross_attn_to_xformers():
    print('Replace CrossAttention.forward to use xformers')
    try:
        import xformers.ops
    except ImportError:
        raise ImportError('No xformers / xformersがインストールされていないようです')

    def forward_xformers(self, x, context=None, mask=None):
        h = self.heads
        q_in = self.to_q(x)

        context = default(context, x)
        context = context.to(x.dtype)

        if hasattr(self, 'hypernetwork') and self.hypernetwork is not None:
            context_k, context_v = self.hypernetwork.forward(x, context)
            context_k = context_k.to(x.dtype)
            context_v = context_v.to(x.dtype)
        else:
            context_k = context
            context_v = context

        k_in = self.to_k(context_k)
        v_in = self.to_v(context_v)

        q, k, v = map(
            lambda t: rearrange(t, 'b n (h d) -> b n h d', h=h),
            (q_in, k_in, v_in),
        )
        del q_in, k_in, v_in

        q = q.contiguous()
        k = k.contiguous()
        v = v.contiguous()
        out = xformers.ops.memory_efficient_attention(
            q, k, v, attn_bias=None
        )        # 最適なのを選んでくれる

        out = rearrange(out, 'b n h d -> b n (h d)', h=h)

        # diffusers 0.7.0~
        out = self.to_out[0](out)
        out = self.to_out[1](out)
        return out

    diffusers.models.attention.CrossAttention.forward = forward_xformers


# endregion


if __name__ == '__main__':
    # torch.cuda.set_per_process_memory_fraction(0.48)
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--v2',
        action='store_true',
        help='load Stable Diffusion v2.0 model / Stable Diffusion 2.0のモデルを読み込む',
    )
    parser.add_argument(
        '--v_parameterization',
        action='store_true',
        help='enable v-parameterization training / v-parameterization学習を有効にする',
    )
    parser.add_argument(
        '--pretrained_model_name_or_path',
        type=str,
        default=None,
        help='pretrained model to train, directory to Diffusers model or StableDiffusion checkpoint / 学習元モデル、Diffusers形式モデルのディレクトリまたはStableDiffusionのckptファイル',
    )
    parser.add_argument(
        '--in_json',
        type=str,
        default=None,
        help='metadata file to input / 読みこむメタデータファイル',
    )
    parser.add_argument(
        '--shuffle_caption',
        action='store_true',
        help='shuffle comma-separated caption when fine tuning / fine tuning時にコンマで区切られたcaptionの各要素をshuffleする',
    )
    parser.add_argument(
        '--keep_tokens',
        type=int,
        default=None,
        help='keep heading N tokens when shuffling caption tokens / captionのシャッフル時に、先頭からこの個数のトークンをシャッフルしないで残す',
    )
    parser.add_argument(
        '--train_data_dir',
        type=str,
        default=None,
        help='directory for train images / 学習画像データのディレクトリ',
    )
    parser.add_argument(
        '--dataset_repeats',
        type=int,
        default=None,
        help='num times to repeat dataset / 学習にデータセットを繰り返す回数',
    )
    parser.add_argument(
        '--output_dir',
        type=str,
        default=None,
        help='directory to output trained model, save as same format as input / 学習後のモデル出力先ディレクトリ（入力と同じ形式で保存）',
    )
    parser.add_argument(
        '--save_precision',
        type=str,
        default=None,
        choices=[None, 'float', 'fp16', 'bf16'],
        help='precision in saving (available in StableDiffusion checkpoint) / 保存時に精度を変更して保存する（StableDiffusion形式での保存時のみ有効）',
    )
    parser.add_argument(
        '--save_model_as',
        type=str,
        default=None,
        choices=[
            None,
            'ckpt',
            'safetensors',
            'diffusers',
            'diffusers_safetensors',
        ],
        help='format to save the model (default is same to original) / モデル保存時の形式（未指定時は元モデルと同じ）',
    )
    parser.add_argument(
        '--use_safetensors',
        action='store_true',
        help='use safetensors format to save (if save_model_as is not specified) / checkpoint、モデルをsafetensors形式で保存する（save_model_as未指定時）',
    )
    parser.add_argument(
        '--train_text_encoder',
        action='store_true',
        help='train text encoder / text encoderも学習する',
    )
    parser.add_argument(
        '--hypernetwork_module',
        type=str,
        default=None,
        help='train hypernetwork instead of fine tuning, module to use / fine tuningの代わりにHypernetworkの学習をする場合、そのモジュール',
    )
    parser.add_argument(
        '--hypernetwork_weights',
        type=str,
        default=None,
        help='hypernetwork weights to initialize for additional training / Hypernetworkの学習時に読み込む重み（Hypernetworkの追加学習）',
    )
    parser.add_argument(
        '--save_every_n_epochs',
        type=int,
        default=None,
        help='save checkpoint every N epochs / 学習中のモデルを指定エポックごとに保存する',
    )
    parser.add_argument(
        '--save_state',
        action='store_true',
        help='save training state additionally (including optimizer states etc.) / optimizerなど学習状態も含めたstateを追加で保存する',
    )
    parser.add_argument(
        '--resume',
        type=str,
        default=None,
        help='saved state to resume training / 学習再開するモデルのstate',
    )
    parser.add_argument(
        '--max_token_length',
        type=int,
        default=None,
        choices=[None, 150, 225],
        help='max token length of text encoder (default for 75, 150 or 225) / text encoderのトークンの最大長（未指定で75、150または225が指定可）',
    )
    parser.add_argument(
        '--train_batch_size',
        type=int,
        default=1,
        help='batch size for training / 学習時のバッチサイズ',
    )
    parser.add_argument(
        '--use_8bit_adam',
        action='store_true',
        help='use 8bit Adam optimizer (requires bitsandbytes) / 8bit Adamオプティマイザを使う（bitsandbytesのインストールが必要）',
    )
    parser.add_argument(
        '--mem_eff_attn',
        action='store_true',
        help='use memory efficient attention for CrossAttention / CrossAttentionに省メモリ版attentionを使う',
    )
    parser.add_argument(
        '--xformers',
        action='store_true',
        help='use xformers for CrossAttention / CrossAttentionにxformersを使う',
    )
    parser.add_argument(
        '--diffusers_xformers',
        action='store_true',
        help="use xformers by diffusers (Hypernetworks doesn't work) / Diffusersでxformersを使用する（Hypernetwork利用不可）",
    )
    parser.add_argument(
        '--learning_rate',
        type=float,
        default=2.0e-6,
        help='learning rate / 学習率',
    )
    parser.add_argument(
        '--max_train_steps',
        type=int,
        default=1600,
        help='training steps / 学習ステップ数',
    )
    parser.add_argument(
        '--seed',
        type=int,
        default=None,
        help='random seed for training / 学習時の乱数のseed',
    )
    parser.add_argument(
        '--gradient_checkpointing',
        action='store_true',
        help='enable gradient checkpointing / grandient checkpointingを有効にする',
    )
    parser.add_argument(
        '--gradient_accumulation_steps',
        type=int,
        default=1,
        help='Number of updates steps to accumulate before performing a backward/update pass / 学習時に逆伝播をする前に勾配を合計するステップ数',
    )
    parser.add_argument(
        '--mixed_precision',
        type=str,
        default='no',
        choices=['no', 'fp16', 'bf16'],
        help='use mixed precision / 混合精度を使う場合、その精度',
    )
    parser.add_argument(
        '--full_fp16',
        action='store_true',
        help='fp16 training including gradients / 勾配も含めてfp16で学習する',
    )
    parser.add_argument(
        '--clip_skip',
        type=int,
        default=None,
        help='use output of nth layer from back of text encoder (n>=1) / text encoderの後ろからn番目の層の出力を用いる（nは1以上）',
    )
    parser.add_argument(
        '--debug_dataset',
        action='store_true',
        help='show images for debugging (do not train) / デバッグ用に学習データを画面表示する（学習は行わない）',
    )
    parser.add_argument(
        '--logging_dir',
        type=str,
        default=None,
        help='enable logging and output TensorBoard log to this directory / ログ出力を有効にしてこのディレクトリにTensorBoard用のログを出力する',
    )
    parser.add_argument(
        '--log_prefix',
        type=str,
        default=None,
        help='add prefix for each log directory / ログディレクトリ名の先頭に追加する文字列',
    )
    parser.add_argument(
        '--lr_scheduler',
        type=str,
        default='constant',
        help='scheduler to use for learning rate / 学習率のスケジューラ: linear, cosine, cosine_with_restarts, polynomial, constant (default), constant_with_warmup',
    )
    parser.add_argument(
        '--lr_warmup_steps',
        type=int,
        default=0,
        help='Number of steps for the warmup in the lr scheduler (default is 0) / 学習率のスケジューラをウォームアップするステップ数（デフォルト0）',
    )

    args = parser.parse_args()
    train(args)