import argparse from dataclasses import ( asdict, dataclass, ) import functools import random from textwrap import dedent, indent import json from pathlib import Path # from toolz import curry from typing import ( List, Optional, Sequence, Tuple, Union, ) import toml import voluptuous from voluptuous import ( Any, ExactSequence, MultipleInvalid, Object, Required, Schema, ) from transformers import CLIPTokenizer from . import train_util from .train_util import ( DreamBoothSubset, FineTuningSubset, ControlNetSubset, DreamBoothDataset, FineTuningDataset, ControlNetDataset, DatasetGroup, ) def add_config_arguments(parser: argparse.ArgumentParser): parser.add_argument("--dataset_config", type=Path, default=None, help="config file for detail settings / 詳細な設定用の設定ファイル") # TODO: inherit Params class in Subset, Dataset @dataclass class BaseSubsetParams: image_dir: Optional[str] = None num_repeats: int = 1 shuffle_caption: bool = False keep_tokens: int = 0 color_aug: bool = False flip_aug: bool = False face_crop_aug_range: Optional[Tuple[float, float]] = None random_crop: bool = False caption_dropout_rate: float = 0.0 caption_dropout_every_n_epochs: int = 0 caption_tag_dropout_rate: float = 0.0 token_warmup_min: int = 1 token_warmup_step: float = 0 @dataclass class DreamBoothSubsetParams(BaseSubsetParams): is_reg: bool = False class_tokens: Optional[str] = None caption_extension: str = ".caption" @dataclass class FineTuningSubsetParams(BaseSubsetParams): metadata_file: Optional[str] = None @dataclass class ControlNetSubsetParams(BaseSubsetParams): conditioning_data_dir: str = None caption_extension: str = ".caption" @dataclass class BaseDatasetParams: tokenizer: Union[CLIPTokenizer, List[CLIPTokenizer]] = None max_token_length: int = None resolution: Optional[Tuple[int, int]] = None debug_dataset: bool = False @dataclass class DreamBoothDatasetParams(BaseDatasetParams): batch_size: int = 1 enable_bucket: bool = False min_bucket_reso: int = 256 max_bucket_reso: int = 1024 bucket_reso_steps: int = 64 bucket_no_upscale: bool = False prior_loss_weight: float = 1.0 @dataclass class FineTuningDatasetParams(BaseDatasetParams): batch_size: int = 1 enable_bucket: bool = False min_bucket_reso: int = 256 max_bucket_reso: int = 1024 bucket_reso_steps: int = 64 bucket_no_upscale: bool = False @dataclass class ControlNetDatasetParams(BaseDatasetParams): batch_size: int = 1 enable_bucket: bool = False min_bucket_reso: int = 256 max_bucket_reso: int = 1024 bucket_reso_steps: int = 64 bucket_no_upscale: bool = False @dataclass class SubsetBlueprint: params: Union[DreamBoothSubsetParams, FineTuningSubsetParams] @dataclass class DatasetBlueprint: is_dreambooth: bool is_controlnet: bool params: Union[DreamBoothDatasetParams, FineTuningDatasetParams] subsets: Sequence[SubsetBlueprint] @dataclass class DatasetGroupBlueprint: datasets: Sequence[DatasetBlueprint] @dataclass class Blueprint: dataset_group: DatasetGroupBlueprint class ConfigSanitizer: # @curry @staticmethod def __validate_and_convert_twodim(klass, value: Sequence) -> Tuple: Schema(ExactSequence([klass, klass]))(value) return tuple(value) # @curry @staticmethod def __validate_and_convert_scalar_or_twodim(klass, value: Union[float, Sequence]) -> Tuple: Schema(Any(klass, ExactSequence([klass, klass])))(value) try: Schema(klass)(value) return (value, value) except: return ConfigSanitizer.__validate_and_convert_twodim(klass, value) # subset schema SUBSET_ASCENDABLE_SCHEMA = { "color_aug": bool, "face_crop_aug_range": functools.partial(__validate_and_convert_twodim.__func__, float), "flip_aug": bool, "num_repeats": int, "random_crop": bool, "shuffle_caption": bool, "keep_tokens": int, "token_warmup_min": int, "token_warmup_step": Any(float,int), } # DO means DropOut DO_SUBSET_ASCENDABLE_SCHEMA = { "caption_dropout_every_n_epochs": int, "caption_dropout_rate": Any(float, int), "caption_tag_dropout_rate": Any(float, int), } # DB means DreamBooth DB_SUBSET_ASCENDABLE_SCHEMA = { "caption_extension": str, "class_tokens": str, } DB_SUBSET_DISTINCT_SCHEMA = { Required("image_dir"): str, "is_reg": bool, } # FT means FineTuning FT_SUBSET_DISTINCT_SCHEMA = { Required("metadata_file"): str, "image_dir": str, } CN_SUBSET_ASCENDABLE_SCHEMA = { "caption_extension": str, } CN_SUBSET_DISTINCT_SCHEMA = { Required("image_dir"): str, Required("conditioning_data_dir"): str, } # datasets schema DATASET_ASCENDABLE_SCHEMA = { "batch_size": int, "bucket_no_upscale": bool, "bucket_reso_steps": int, "enable_bucket": bool, "max_bucket_reso": int, "min_bucket_reso": int, "resolution": functools.partial(__validate_and_convert_scalar_or_twodim.__func__, int), } # options handled by argparse but not handled by user config ARGPARSE_SPECIFIC_SCHEMA = { "debug_dataset": bool, "max_token_length": Any(None, int), "prior_loss_weight": Any(float, int), } # for handling default None value of argparse ARGPARSE_NULLABLE_OPTNAMES = [ "face_crop_aug_range", "resolution", ] # prepare map because option name may differ among argparse and user config ARGPARSE_OPTNAME_TO_CONFIG_OPTNAME = { "train_batch_size": "batch_size", "dataset_repeats": "num_repeats", } def __init__(self, support_dreambooth: bool, support_finetuning: bool, support_controlnet: bool, support_dropout: bool) -> None: assert support_dreambooth or support_finetuning or support_controlnet, "Neither DreamBooth mode nor fine tuning mode specified. Please specify one mode or more. / DreamBooth モードか fine tuning モードのどちらも指定されていません。1つ以上指定してください。" self.db_subset_schema = self.__merge_dict( self.SUBSET_ASCENDABLE_SCHEMA, self.DB_SUBSET_DISTINCT_SCHEMA, self.DB_SUBSET_ASCENDABLE_SCHEMA, self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {}, ) self.ft_subset_schema = self.__merge_dict( self.SUBSET_ASCENDABLE_SCHEMA, self.FT_SUBSET_DISTINCT_SCHEMA, self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {}, ) self.cn_subset_schema = self.__merge_dict( self.SUBSET_ASCENDABLE_SCHEMA, self.CN_SUBSET_DISTINCT_SCHEMA, self.CN_SUBSET_ASCENDABLE_SCHEMA, self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {}, ) self.db_dataset_schema = self.__merge_dict( self.DATASET_ASCENDABLE_SCHEMA, self.SUBSET_ASCENDABLE_SCHEMA, self.DB_SUBSET_ASCENDABLE_SCHEMA, self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {}, {"subsets": [self.db_subset_schema]}, ) self.ft_dataset_schema = self.__merge_dict( self.DATASET_ASCENDABLE_SCHEMA, self.SUBSET_ASCENDABLE_SCHEMA, self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {}, {"subsets": [self.ft_subset_schema]}, ) self.cn_dataset_schema = self.__merge_dict( self.DATASET_ASCENDABLE_SCHEMA, self.SUBSET_ASCENDABLE_SCHEMA, self.CN_SUBSET_ASCENDABLE_SCHEMA, self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {}, {"subsets": [self.cn_subset_schema]}, ) if support_dreambooth and support_finetuning: def validate_flex_dataset(dataset_config: dict): subsets_config = dataset_config.get("subsets", []) if support_controlnet and all(["conditioning_data_dir" in subset for subset in subsets_config]): return Schema(self.cn_dataset_schema)(dataset_config) # check dataset meets FT style # NOTE: all FT subsets should have "metadata_file" elif all(["metadata_file" in subset for subset in subsets_config]): return Schema(self.ft_dataset_schema)(dataset_config) # check dataset meets DB style # NOTE: all DB subsets should have no "metadata_file" elif all(["metadata_file" not in subset for subset in subsets_config]): return Schema(self.db_dataset_schema)(dataset_config) else: raise voluptuous.Invalid("DreamBooth subset and fine tuning subset cannot be mixed in the same dataset. Please split them into separate datasets. / DreamBoothのサブセットとfine tuninのサブセットを同一のデータセットに混在させることはできません。別々のデータセットに分割してください。") self.dataset_schema = validate_flex_dataset elif support_dreambooth: self.dataset_schema = self.db_dataset_schema elif support_finetuning: self.dataset_schema = self.ft_dataset_schema elif support_controlnet: self.dataset_schema = self.cn_dataset_schema self.general_schema = self.__merge_dict( self.DATASET_ASCENDABLE_SCHEMA, self.SUBSET_ASCENDABLE_SCHEMA, self.DB_SUBSET_ASCENDABLE_SCHEMA if support_dreambooth else {}, self.CN_SUBSET_ASCENDABLE_SCHEMA if support_controlnet else {}, self.DO_SUBSET_ASCENDABLE_SCHEMA if support_dropout else {}, ) self.user_config_validator = Schema({ "general": self.general_schema, "datasets": [self.dataset_schema], }) self.argparse_schema = self.__merge_dict( self.general_schema, self.ARGPARSE_SPECIFIC_SCHEMA, {optname: Any(None, self.general_schema[optname]) for optname in self.ARGPARSE_NULLABLE_OPTNAMES}, {a_name: self.general_schema[c_name] for a_name, c_name in self.ARGPARSE_OPTNAME_TO_CONFIG_OPTNAME.items()}, ) self.argparse_config_validator = Schema(Object(self.argparse_schema), extra=voluptuous.ALLOW_EXTRA) def sanitize_user_config(self, user_config: dict) -> dict: try: return self.user_config_validator(user_config) except MultipleInvalid: # TODO: エラー発生時のメッセージをわかりやすくする print("Invalid user config / ユーザ設定の形式が正しくないようです") raise # NOTE: In nature, argument parser result is not needed to be sanitize # However this will help us to detect program bug def sanitize_argparse_namespace(self, argparse_namespace: argparse.Namespace) -> argparse.Namespace: try: return self.argparse_config_validator(argparse_namespace) except MultipleInvalid: # XXX: this should be a bug print("Invalid cmdline parsed arguments. This should be a bug. / コマンドラインのパース結果が正しくないようです。プログラムのバグの可能性が高いです。") raise # NOTE: value would be overwritten by latter dict if there is already the same key @staticmethod def __merge_dict(*dict_list: dict) -> dict: merged = {} for schema in dict_list: # merged |= schema for k, v in schema.items(): merged[k] = v return merged class BlueprintGenerator: BLUEPRINT_PARAM_NAME_TO_CONFIG_OPTNAME = { } def __init__(self, sanitizer: ConfigSanitizer): self.sanitizer = sanitizer # runtime_params is for parameters which is only configurable on runtime, such as tokenizer def generate(self, user_config: dict, argparse_namespace: argparse.Namespace, **runtime_params) -> Blueprint: sanitized_user_config = self.sanitizer.sanitize_user_config(user_config) sanitized_argparse_namespace = self.sanitizer.sanitize_argparse_namespace(argparse_namespace) # convert argparse namespace to dict like config # NOTE: it is ok to have extra entries in dict optname_map = self.sanitizer.ARGPARSE_OPTNAME_TO_CONFIG_OPTNAME argparse_config = {optname_map.get(optname, optname): value for optname, value in vars(sanitized_argparse_namespace).items()} general_config = sanitized_user_config.get("general", {}) dataset_blueprints = [] for dataset_config in sanitized_user_config.get("datasets", []): # NOTE: if subsets have no "metadata_file", these are DreamBooth datasets/subsets subsets = dataset_config.get("subsets", []) is_dreambooth = all(["metadata_file" not in subset for subset in subsets]) is_controlnet = all(["conditioning_data_dir" in subset for subset in subsets]) if is_controlnet: subset_params_klass = ControlNetSubsetParams dataset_params_klass = ControlNetDatasetParams elif is_dreambooth: subset_params_klass = DreamBoothSubsetParams dataset_params_klass = DreamBoothDatasetParams else: subset_params_klass = FineTuningSubsetParams dataset_params_klass = FineTuningDatasetParams subset_blueprints = [] for subset_config in subsets: params = self.generate_params_by_fallbacks(subset_params_klass, [subset_config, dataset_config, general_config, argparse_config, runtime_params]) subset_blueprints.append(SubsetBlueprint(params)) params = self.generate_params_by_fallbacks(dataset_params_klass, [dataset_config, general_config, argparse_config, runtime_params]) dataset_blueprints.append(DatasetBlueprint(is_dreambooth, is_controlnet, params, subset_blueprints)) dataset_group_blueprint = DatasetGroupBlueprint(dataset_blueprints) return Blueprint(dataset_group_blueprint) @staticmethod def generate_params_by_fallbacks(param_klass, fallbacks: Sequence[dict]): name_map = BlueprintGenerator.BLUEPRINT_PARAM_NAME_TO_CONFIG_OPTNAME search_value = BlueprintGenerator.search_value default_params = asdict(param_klass()) param_names = default_params.keys() params = {name: search_value(name_map.get(name, name), fallbacks, default_params.get(name)) for name in param_names} return param_klass(**params) @staticmethod def search_value(key: str, fallbacks: Sequence[dict], default_value = None): for cand in fallbacks: value = cand.get(key) if value is not None: return value return default_value def generate_dataset_group_by_blueprint(dataset_group_blueprint: DatasetGroupBlueprint): datasets: List[Union[DreamBoothDataset, FineTuningDataset, ControlNetDataset]] = [] for dataset_blueprint in dataset_group_blueprint.datasets: if dataset_blueprint.is_controlnet: subset_klass = ControlNetSubset dataset_klass = ControlNetDataset elif dataset_blueprint.is_dreambooth: subset_klass = DreamBoothSubset dataset_klass = DreamBoothDataset else: subset_klass = FineTuningSubset dataset_klass = FineTuningDataset subsets = [subset_klass(**asdict(subset_blueprint.params)) for subset_blueprint in dataset_blueprint.subsets] dataset = dataset_klass(subsets=subsets, **asdict(dataset_blueprint.params)) datasets.append(dataset) # print info info = "" for i, dataset in enumerate(datasets): is_dreambooth = isinstance(dataset, DreamBoothDataset) is_controlnet = isinstance(dataset, ControlNetDataset) info += dedent(f"""\ [Dataset {i}] batch_size: {dataset.batch_size} resolution: {(dataset.width, dataset.height)} enable_bucket: {dataset.enable_bucket} """) if dataset.enable_bucket: info += indent(dedent(f"""\ min_bucket_reso: {dataset.min_bucket_reso} max_bucket_reso: {dataset.max_bucket_reso} bucket_reso_steps: {dataset.bucket_reso_steps} bucket_no_upscale: {dataset.bucket_no_upscale} \n"""), " ") else: info += "\n" for j, subset in enumerate(dataset.subsets): info += indent(dedent(f"""\ [Subset {j} of Dataset {i}] image_dir: "{subset.image_dir}" image_count: {subset.img_count} num_repeats: {subset.num_repeats} shuffle_caption: {subset.shuffle_caption} keep_tokens: {subset.keep_tokens} caption_dropout_rate: {subset.caption_dropout_rate} caption_dropout_every_n_epoches: {subset.caption_dropout_every_n_epochs} caption_tag_dropout_rate: {subset.caption_tag_dropout_rate} color_aug: {subset.color_aug} flip_aug: {subset.flip_aug} face_crop_aug_range: {subset.face_crop_aug_range} random_crop: {subset.random_crop} token_warmup_min: {subset.token_warmup_min}, token_warmup_step: {subset.token_warmup_step}, """), " ") if is_dreambooth: info += indent(dedent(f"""\ is_reg: {subset.is_reg} class_tokens: {subset.class_tokens} caption_extension: {subset.caption_extension} \n"""), " ") elif not is_controlnet: info += indent(dedent(f"""\ metadata_file: {subset.metadata_file} \n"""), " ") print(info) # make buckets first because it determines the length of dataset # and set the same seed for all datasets seed = random.randint(0, 2**31) # actual seed is seed + epoch_no for i, dataset in enumerate(datasets): print(f"[Dataset {i}]") dataset.make_buckets() dataset.set_seed(seed) return DatasetGroup(datasets) def generate_dreambooth_subsets_config_by_subdirs(train_data_dir: Optional[str] = None, reg_data_dir: Optional[str] = None): def extract_dreambooth_params(name: str) -> Tuple[int, str]: tokens = name.split('_') try: n_repeats = int(tokens[0]) except ValueError as e: print(f"ignore directory without repeats / 繰り返し回数のないディレクトリを無視します: {name}") return 0, "" caption_by_folder = '_'.join(tokens[1:]) return n_repeats, caption_by_folder def generate(base_dir: Optional[str], is_reg: bool): if base_dir is None: return [] base_dir: Path = Path(base_dir) if not base_dir.is_dir(): return [] subsets_config = [] for subdir in base_dir.iterdir(): if not subdir.is_dir(): continue num_repeats, class_tokens = extract_dreambooth_params(subdir.name) if num_repeats < 1: continue subset_config = {"image_dir": str(subdir), "num_repeats": num_repeats, "is_reg": is_reg, "class_tokens": class_tokens} subsets_config.append(subset_config) return subsets_config subsets_config = [] subsets_config += generate(train_data_dir, False) subsets_config += generate(reg_data_dir, True) return subsets_config def generate_controlnet_subsets_config_by_subdirs(train_data_dir: Optional[str] = None, conditioning_data_dir: Optional[str] = None, caption_extension: str = ".txt"): def generate(base_dir: Optional[str]): if base_dir is None: return [] base_dir: Path = Path(base_dir) if not base_dir.is_dir(): return [] subsets_config = [] for subdir in base_dir.iterdir(): if not subdir.is_dir(): continue subset_config = {"image_dir": str(subdir), "conditioning_data_dir": conditioning_data_dir, "caption_extension": caption_extension, "num_repeats": 1} subsets_config.append(subset_config) return subsets_config subsets_config = [] subsets_config += generate(train_data_dir, False) return subsets_config def load_user_config(file: str) -> dict: file: Path = Path(file) if not file.is_file(): raise ValueError(f"file not found / ファイルが見つかりません: {file}") if file.name.lower().endswith('.json'): try: with open(file, 'r') as f: config = json.load(f) except Exception: print(f"Error on parsing JSON config file. Please check the format. / JSON 形式の設定ファイルの読み込みに失敗しました。文法が正しいか確認してください。: {file}") raise elif file.name.lower().endswith('.toml'): try: config = toml.load(file) except Exception: print(f"Error on parsing TOML config file. Please check the format. / TOML 形式の設定ファイルの読み込みに失敗しました。文法が正しいか確認してください。: {file}") raise else: raise ValueError(f"not supported config file format / 対応していない設定ファイルの形式です: {file}") return config # for config test if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--support_dreambooth", action="store_true") parser.add_argument("--support_finetuning", action="store_true") parser.add_argument("--support_controlnet", action="store_true") parser.add_argument("--support_dropout", action="store_true") parser.add_argument("dataset_config") config_args, remain = parser.parse_known_args() parser = argparse.ArgumentParser() train_util.add_dataset_arguments(parser, config_args.support_dreambooth, config_args.support_finetuning, config_args.support_dropout) train_util.add_training_arguments(parser, config_args.support_dreambooth) argparse_namespace = parser.parse_args(remain) train_util.prepare_dataset_args(argparse_namespace, config_args.support_finetuning) print("[argparse_namespace]") print(vars(argparse_namespace)) user_config = load_user_config(config_args.dataset_config) print("\n[user_config]") print(user_config) sanitizer = ConfigSanitizer(config_args.support_dreambooth, config_args.support_finetuning, config_args.support_controlnet, config_args.support_dropout) sanitized_user_config = sanitizer.sanitize_user_config(user_config) print("\n[sanitized_user_config]") print(sanitized_user_config) blueprint = BlueprintGenerator(sanitizer).generate(user_config, argparse_namespace) print("\n[blueprint]") print(blueprint)