Kohya-ss-sd-scripts/sdxl_train.py

# training with captions

import argparse
from typing import List, Optional, Union

import torch
from accelerate import Accelerator
from library.device_utils import init_ipex, clean_memory_on_device

init_ipex()

from library import sdxl_model_util, strategy_sd, strategy_sdxl

import library.train_util as train_util

from library.utils import setup_logging, add_logging_arguments

import library.sdxl_train_util as sdxl_train_util
from library.sdxl_original_unet import SdxlUNet2DConditionModel
import train_native

setup_logging()
import logging

logger = logging.getLogger(__name__)

class SdxlNativeTrainer(train_native.NativeTrainer):
    def __init__(self):
        super().__init__()
        self.vae_scale_factor = sdxl_model_util.VAE_SCALE_FACTOR
        self.unet_num_blocks_for_block_lr = sdxl_model_util.UNET_NUM_BLOCKS_FOR_BLOCK_LR
        self.is_sdxl = True
        self.arb_min_steps = sdxl_model_util.ARB_MIN_STEPS

    def get_block_params_to_optimize(self, unet: SdxlUNet2DConditionModel, block_lrs: List[float]) -> List[dict]:
        block_params = [[] for _ in range(len(block_lrs))]

        for i, (name, param) in enumerate(unet.named_parameters()):
            if name.startswith("time_embed.") or name.startswith("label_emb."):
                block_index = 0  # 0
            elif name.startswith("input_blocks."):  # 1-9
                block_index = 1 + int(name.split(".")[1])
            elif name.startswith("middle_block."):  # 10-12
                block_index = 10 + int(name.split(".")[1])
            elif name.startswith("output_blocks."):  # 13-21
                block_index = 13 + int(name.split(".")[1])
            elif name.startswith("out."):  # 22
                block_index = 22
            else:
                raise ValueError(f"unexpected parameter name: {name}")

            block_params[block_index].append(param)

        params_to_optimize = []
        for i, params in enumerate(block_params):
            if block_lrs[i] == 0:  # 0のときは学習しない do not optimize when lr is 0
                continue
            params_to_optimize.append({"params": params, "lr": block_lrs[i]})

        return params_to_optimize

    def append_block_lr_to_logs(self, block_lrs, logs, lr_scheduler, optimizer_type):
        names = []
        block_index = 0
        while block_index < self.unet_num_blocks_for_block_lr + 2:
            if block_index < self.unet_num_blocks_for_block_lr:
                if block_lrs[block_index] == 0:
                    block_index += 1
                    continue
                names.append(f"block{block_index}")
            elif block_index == self.unet_num_blocks_for_block_lr:
                names.append("text_encoder1")
            elif block_index == self.unet_num_blocks_for_block_lr + 1:
                names.append("text_encoder2")

            block_index += 1

        train_util.append_lr_to_logs_with_names(logs, lr_scheduler, optimizer_type, names)

    def assert_extra_args(self, args, train_dataset_group: Union[train_util.DatasetGroup, train_util.MinimalDataset], val_dataset_group: Optional[train_util.DatasetGroup]):
        super().assert_extra_args(args, train_dataset_group, val_dataset_group)
        sdxl_train_util.verify_sdxl_training_args(args)

        if args.cache_text_encoder_outputs:
            assert (
                train_dataset_group.is_text_encoder_output_cacheable()
            ), "when caching Text Encoder output, either caption_dropout_rate, shuffle_caption, token_warmup_step or caption_tag_dropout_rate cannot be used / Text Encoderの出力をキャッシュするときはcaption_dropout_rate, shuffle_caption, token_warmup_step, caption_tag_dropout_rateは使えません"
        
        train_dataset_group.verify_bucket_reso_steps(self.arb_min_steps)
        if val_dataset_group is not None:
            val_dataset_group.verify_bucket_reso_steps(self.arb_min_steps)

    def get_tokenize_strategy(self, args):
        return strategy_sdxl.SdxlTokenizeStrategy(args.max_token_length, args.tokenizer_cache_dir)

    def get_tokenizers(self, tokenize_strategy: strategy_sdxl.SdxlTokenizeStrategy):
        return [tokenize_strategy.tokenizer1, tokenize_strategy.tokenizer2]  # will be removed in the future
    
    def get_latents_caching_strategy(self, args):
        latents_caching_strategy = strategy_sd.SdSdxlLatentsCachingStrategy(
            False, args.cache_latents_to_disk, args.vae_batch_size, args.skip_cache_check
        )
        return latents_caching_strategy
    
    def load_target_model(self, args, weight_dtype, accelerator):
        (
            load_stable_diffusion_format,
            text_encoder1,
            text_encoder2,
            vae,
            unet,
            logit_scale,
            ckpt_info,
        ) = sdxl_train_util.load_target_model(args, accelerator, sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, weight_dtype)

        self.load_stable_diffusion_format = load_stable_diffusion_format
        self.logit_scale = logit_scale
        self.ckpt_info = ckpt_info

        # モデルに xformers とか memory efficient attention を組み込む
        if args.diffusers_xformers:
            # もうU-Netを独自にしたので動かないけどVAEのxformersは動くはず
            # How about Text encoders?
            accelerator.print("Use xformers by Diffusers")
            if not self.is_sdxl:
                self.set_diffusers_xformers_flag(unet, True)
            self.set_diffusers_xformers_flag(vae, True)
            self.set_diffusers_xformers_flag(text_encoder1, True)
            self.set_diffusers_xformers_flag(text_encoder2, True)
        else:
            # Windows版のxformersはfloatで学習できなかったりするのでxformersを使わない設定も可能にしておく必要がある
            accelerator.print("Disable Diffusers' xformers")
            train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers, args.sdpa)
            if args.xformers and torch.__version__ >= "2.0.0":  # PyTorch 2.0.0 以上対応のxformersなら以下が使える
                #vae.set_use_memory_efficient_attention_xformers(args.xformers)
                self.set_diffusers_xformers_flag(vae, True)
                self.set_diffusers_xformers_flag(text_encoder1, True)
                self.set_diffusers_xformers_flag(text_encoder2, True)

        return sdxl_model_util.MODEL_VERSION_SDXL_BASE_V1_0, [text_encoder1, text_encoder2], vae, unet

    def get_text_encoding_strategy(self, args):
        return strategy_sdxl.SdxlTextEncodingStrategy()
    
    def get_models_for_text_encoding(self, args, accelerator, text_encoders):
        return text_encoders + [accelerator.unwrap_model(text_encoders[-1])]

    def get_text_encoder_outputs_caching_strategy(self, args):
        if args.cache_text_encoder_outputs:
            return strategy_sdxl.SdxlTextEncoderOutputsCachingStrategy(
                args.cache_text_encoder_outputs_to_disk, None, args.skip_cache_check, is_weighted=args.weighted_captions
            )
        else:
            return None
        
    def cache_text_encoder_outputs_if_needed(
        self, args, accelerator: Accelerator, unet, vae, text_encoders, dataset: train_util.DatasetGroup, weight_dtype
    ):
        if args.cache_text_encoder_outputs:
            if not args.lowram:
                # メモリ消費を減らす
                logger.info("move vae and unet to cpu to save memory")
                org_vae_device = vae.device
                org_unet_device = unet.device
                vae.to("cpu")
                unet.to("cpu")
                clean_memory_on_device(accelerator.device)

            # When TE is not be trained, it will not be prepared so we need to use explicit autocast
            text_encoders[0].to(accelerator.device, dtype=weight_dtype)
            text_encoders[1].to(accelerator.device, dtype=weight_dtype)
            with accelerator.autocast():
                dataset.new_cache_text_encoder_outputs(text_encoders + [accelerator.unwrap_model(text_encoders[-1])], accelerator)
            accelerator.wait_for_everyone()

            text_encoders[0].to("cpu", dtype=torch.float32)  # Text Encoder doesn't work with fp16 on CPU
            text_encoders[1].to("cpu", dtype=torch.float32)
            clean_memory_on_device(accelerator.device)

            if not args.lowram:
                logger.info("move vae and unet back to original device")
                vae.to(org_vae_device)
                unet.to(org_unet_device)
        else:
            # Text Encoderから毎回出力を取得するので、GPUに乗せておく
            text_encoders[0].to(accelerator.device, dtype=weight_dtype)
            text_encoders[1].to(accelerator.device, dtype=weight_dtype)

    def call_unet(
        self,
        args,
        accelerator,
        unet,
        noisy_latents,
        timesteps,
        text_conds,
        batch,
        weight_dtype,
        indices: Optional[List[int]] = None,
    ):
        noisy_latents = noisy_latents.to(weight_dtype)  # TODO check why noisy_latents is not weight_dtype

        # get size embeddings
        orig_size = batch["original_sizes_hw"]
        crop_size = batch["crop_top_lefts"]
        target_size = batch["target_sizes_hw"]
        embs = sdxl_train_util.get_size_embeddings(orig_size, crop_size, target_size, accelerator.device).to(weight_dtype)

        # concat embeddings
        encoder_hidden_states1, encoder_hidden_states2, pool2 = text_conds
        vector_embedding = torch.cat([pool2, embs], dim=1).to(weight_dtype)
        text_embedding = torch.cat([encoder_hidden_states1, encoder_hidden_states2], dim=2).to(weight_dtype)

        if indices is not None and len(indices) > 0:
            noisy_latents = noisy_latents[indices]
            timesteps = timesteps[indices]
            text_embedding = text_embedding[indices]
            vector_embedding = vector_embedding[indices]

        noise_pred = unet(noisy_latents, timesteps, text_embedding, vector_embedding)
        return noise_pred

    def sample_images(self, accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet):
        sdxl_train_util.sample_images(accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet)

    def save_model_on_epoch_end_or_stepwise(self, args, on_epoch_end, accelerator, save_dtype, epoch, num_train_epochs, global_step, text_encoders, vae, unet):
        src_path = self.src_stable_diffusion_ckpt if self.save_stable_diffusion_format else self.src_diffusers_model_path
        sdxl_train_util.save_sd_model_on_epoch_end_or_stepwise(
            args,
            on_epoch_end,
            accelerator,
            src_path,
            self.save_stable_diffusion_format,
            self.use_safetensors,
            save_dtype,
            epoch,
            num_train_epochs,
            global_step,
            accelerator.unwrap_model(text_encoders[0]), #text_encoder1
            accelerator.unwrap_model(text_encoders[1]), #text_encoder2
            accelerator.unwrap_model(unet),
            vae,
            self.logit_scale,
            self.ckpt_info,
        )

    def save_model_on_train_end(self, args, accelerator, save_dtype, epoch, global_step, text_encoders, vae, unet):
        src_path = self.src_stable_diffusion_ckpt if self.save_stable_diffusion_format else self.src_diffusers_model_path
        sdxl_train_util.save_sd_model_on_train_end(
            args,
            src_path,
            self.save_stable_diffusion_format,
            self.use_safetensors,
            save_dtype,
            epoch,
            global_step,
            accelerator.unwrap_model(text_encoders[0]), #text_encoder1
            accelerator.unwrap_model(text_encoders[1]), #text_encoder2
            accelerator.unwrap_model(unet),
            vae,
            self.logit_scale,
            self.ckpt_info,
        )

def setup_parser() -> argparse.ArgumentParser:
    parser = train_native.setup_parser()
    sdxl_train_util.add_sdxl_training_arguments(parser)

    parser.add_argument(
        "--learning_rate_te1",
        type=float,
        default=None,
        help="learning rate for text encoder 1 (ViT-L) / text encoder 1 (ViT-L)の学習率",
    )
    parser.add_argument(
        "--learning_rate_te2",
        type=float,
        default=None,
        help="learning rate for text encoder 2 (BiG-G) / text encoder 2 (BiG-G)の学習率",
    )
    parser.add_argument(
        "--block_lr",
        type=str,
        default=None,
        help=f"learning rates for each block of U-Net, comma-separated, {sdxl_model_util.UNET_NUM_BLOCKS_FOR_BLOCK_LR} values / "
        + f"U-Netの各ブロックの学習率、カンマ区切り、{sdxl_model_util.UNET_NUM_BLOCKS_FOR_BLOCK_LR}個の値",
    )
    return parser

if __name__ == "__main__":
    parser = setup_parser()

    args = parser.parse_args()
    train_util.verify_command_line_training_args(args)
    args = train_util.read_config_from_file(args, parser)

    trainer = SdxlNativeTrainer()
    trainer.train(args)