cyberharem/train/train.py

import glob
import json
import logging
import math
import os.path
from typing import Optional, Union

from gchar.games.base import Character
from hbutils.string import plural_word
from hbutils.system import TemporaryDirectory
from hcpdiff.train_ac import Trainer
from hcpdiff.train_ac_single import TrainerSingleCard
from hcpdiff.utils import load_config_with_cli

from .embedding import create_embedding, _DEFAULT_TRAIN_MODEL
from ..dataset import load_dataset_for_character, save_recommended_tags
from ..utils import data_to_cli_args, get_ch_name

_DEFAULT_TRAIN_CFG = 'cfgs/train/examples/lora_anime_character.yaml'


def _min_training_steps(dataset_size: int, unit: int = 20):
 steps = 4000.9 + (720.9319 - 4000.9) / (1 + (dataset_size / 297.2281) ** 0.6543184)
 return int(round(steps / unit)) * unit


def train_plora(
 source: Union[str, Character], name: Optional[str] = None,
 epochs: int = 13, min_steps: Optional[int] = None,
 save_for_times: int = 15, no_min_steps: bool = False,
 batch_size: int = 4, pretrained_model: str = _DEFAULT_TRAIN_MODEL,
 workdir: str = None, emb_n_words: int = 4, emb_init_text: str = '*[0.017, 1]',
 unet_rank: float = 8, text_encoder_rank: float = 4,
 cfg_file: str = _DEFAULT_TRAIN_CFG, single_card: bool = True,
 dataset_type: str = 'stage3-1200', use_ratio: bool = True,
):
 with load_dataset_for_character(source, dataset_type) as (ch, ds_dir):
 if ch is None:
 if name is None:
 raise ValueError(f'Name should be specified when using custom source - {source!r}.')
 else:
 name = name or get_ch_name(ch)

 dataset_size = len(glob.glob(os.path.join(ds_dir, '*.png')))
 logging.info(f'{plural_word(dataset_size, "image")} found in dataset.')

 actual_steps = epochs * dataset_size
 if not no_min_steps:
 actual_steps = max(actual_steps, _min_training_steps(dataset_size, 20))
 if min_steps is not None:
 actual_steps = max(actual_steps, min_steps)
 save_per_steps = max(int(math.ceil(actual_steps / save_for_times / 20) * 20), 20)
 steps = int(math.ceil(actual_steps / save_per_steps) * save_per_steps)
 epochs = int(math.ceil(steps / dataset_size))
 logging.info(f'Training for {plural_word(steps, "step")}, {plural_word(epochs, "epoch")}, '
 f'save per {plural_word(save_per_steps, "step")} ...')

 workdir = workdir or os.path.join('runs', name)
 os.makedirs(workdir, exist_ok=True)
 # os.makedirs(workdir)
 save_recommended_tags(ds_dir, name, workdir)
 with open(os.path.join(workdir, 'meta.json'), 'w', encoding='utf-8') as f:
 json.dump({
 'dataset': {
 'size': dataset_size,
 'type': dataset_type,
 },
 }, f, indent=4, sort_keys=True, ensure_ascii=False)

 with TemporaryDirectory() as embs_dir:
 logging.info(f'Creating embeddings {name!r} at {embs_dir!r}, '
 f'n_words: {emb_n_words!r}, init_text: {emb_init_text!r}, '
 f'pretrained_model: {pretrained_model!r}.')
 create_embedding(
 name, emb_n_words, emb_init_text,
 replace=True,
 pretrained_model=pretrained_model,
 embs_dir=embs_dir,
 )

 cli_args = data_to_cli_args({
 'train': {
 'train_steps': steps,
 'save_step': save_per_steps,
 'scheduler': {
 'num_training_steps': steps,
 }
 },
 'model': {
 'pretrained_model_name_or_path': pretrained_model,
 },
 'character_name': name,
 'dataset_dir': ds_dir,
 'exp_dir': workdir,
 'unet_rank': unet_rank,
 'text_encoder_rank': text_encoder_rank,
 'tokenizer_pt': {
 'emb_dir': embs_dir,
 },
 'data': {
 'dataset1': {
 'batch_size': batch_size,

 'bucket': {
 '_target_': 'hcpdiff.data.bucket.RatioBucket.from_files',
 'target_area': '${times:512,512}',
 'num_bucket': 5,
 } if use_ratio else {
 '_target_': 'hcpdiff.data.bucket.SizeBucket.from_files',
 'target_area': '---',
 'num_bucket': 1,
 }
 },
 },
 })
 conf = load_config_with_cli(cfg_file, args_list=cli_args) # skip --cfg

 logging.info(f'Training with {cfg_file!r}, args: {cli_args!r} ...')
 if single_card:
 logging.info('Training with single card ...')
 trainer = TrainerSingleCard(conf)
 else:
 logging.info('Training with non-single cards ...')
 trainer = Trainer(conf)

 trainer.train()