Method train

train_textual_inversion.py:193–788

(self, args)

        return [emb]

    def train(self, args):
        if args.output_name is None:
            args.output_name = args.token_string
        use_template = args.use_object_template or args.use_style_template

        args_util.verify_training_args(args)
        accelerator_setup.prepare_dataset_args(args, True)
        setup_logging(args, reset=True)

        cache_latents = args.cache_latents

        if args.seed is not None:
            set_seed(args.seed)

        tokenize_strategy = self.get_tokenize_strategy(args)
        strategy_base.TokenizeStrategy.set_strategy(tokenize_strategy)
        tokenizers = self.get_tokenizers(tokenize_strategy)  # will be removed after sample_image is refactored

        # prepare caching strategy: this must be set before preparing dataset. because dataset may use this strategy for initialization.
        latents_caching_strategy = self.get_latents_caching_strategy(args)
        strategy_base.LatentsCachingStrategy.set_strategy(latents_caching_strategy)

        # acceleratorを準備する
        logger.info("prepare accelerator")
        accelerator = accelerator_setup.prepare_accelerator(args)

        # mixed precisionに対応した型を用意しておき適宜castする
        weight_dtype, save_dtype = accelerator_setup.prepare_dtype(args)
        vae_dtype = torch.float32 if args.no_half_vae else weight_dtype

        # モデルを読み込む
        model_version, text_encoders, vae, unet = self.load_target_model(args, weight_dtype, accelerator)

        # Convert the init_word to token_id
        init_token_ids_list = []
        if args.init_word is not None:
            for i, tokenizer in enumerate(tokenizers):
                init_token_ids = tokenizer.encode(args.init_word, add_special_tokens=False)
                if len(init_token_ids) > 1 and len(init_token_ids) != args.num_vectors_per_token:
                    accelerator.print(
                        f"token length for init words is not same to num_vectors_per_token, init words is repeated or truncated / "
                        + f"初期化単語のトークン長がnum_vectors_per_tokenと合わないため、繰り返しまたは切り捨てが発生します:  tokenizer {i+1}, length {len(init_token_ids)}"
                    )
                init_token_ids_list.append(init_token_ids)
        else:
            init_token_ids_list = [None] * len(tokenizers)

        # tokenizerに新しい単語を追加する。追加する単語の数はnum_vectors_per_token
        # token_stringが hoge の場合、"hoge", "hoge1", "hoge2", ... が追加される
        # add new word to tokenizer, count is num_vectors_per_token
        # if token_string is hoge, "hoge", "hoge1", "hoge2", ... are added

        self.assert_token_string(args.token_string, tokenizers)

        token_strings = [args.token_string] + [f"{args.token_string}{i+1}" for i in range(args.num_vectors_per_token - 1)]
        token_ids_list = []
        token_embeds_list = []
        for i, (tokenizer, text_encoder, init_token_ids) in enumerate(zip(tokenizers, text_encoders, init_token_ids_list)):