diff --git a/README.md b/README.md index 06f652a..28f1138 100644 --- a/README.md +++ b/README.md @@ -100,6 +100,14 @@ English (TODO) [google translate](https://github-com.translate.goog/liasece/sd-w 在 stable-diffusion-webui 的运行命令行中可以看到训练的过程及进度。 +### 使用 SD2 或 SDXL 作为基础模型 + +请正确勾选 `Base on Stable Diffusion V2` 或者 `Base on Stable Diffusion XL` ,否则会导致训练失败。 + +### 高级 sd_script 参数追加或覆盖 + +在 `Append or override the sd_script args.` 文本框中输入参数,务必使用 `--` 开头的参数,例如:`--lr_scheduler="constant_with_warmup" --max_grad_norm=0.0` 。插件代码中会使用 `--` 分隔符作为参数间的分隔符。 + ### 一次训练多种参数训练,充分利用你睡觉时的 GPU 有时,一套训练配置可能并不是最优的。等待你的训练结束然后再重新开始训练,这样的效率太低了。因此,你可以一次性配置多种参数,点击一次训练,自动组合不同的参数进行训练。 diff --git a/liasece_sd_webui_train_tools/ArgsList.py b/liasece_sd_webui_train_tools/ArgsList.py index 8455837..17bd883 100644 --- a/liasece_sd_webui_train_tools/ArgsList.py +++ b/liasece_sd_webui_train_tools/ArgsList.py @@ -13,6 +13,7 @@ import gc import torch +from liasece_sd_webui_train_tools.util import * import liasece_sd_webui_train_tools.PythonContextWarper as pc with pc.PythonContextWarper( to_module_path= os.path.abspath(os.path.join(os.path.dirname(__file__), "sd_scripts")), @@ -30,7 +31,6 @@ def __init__(self): self.save_json_folder: Union[str, None] = None self.save_json_name: Union[str, None] = None self.load_json_path: Union[str, None] = None - self.multi_run_folder: Union[str, None] = None self.reg_img_folder: Union[str, None] = None self.sample_prompts: Union[str, None] = None # path to a txt file that has all of the sample prompts in it, # one per line. Only goes to 75 tokens, will cut off the rest. Just place the prompts into the txt file per line @@ -188,6 +188,11 @@ def __init__(self): self.locon_dim: Union[int, None] = None # deprecated self.locon_alpha: Union[int, None] = None # deprecated self.locon: bool = False # deprecated + self.use_sdxl: bool = False # use the sdxl trainer + self.no_half_vae: bool = False # Disable the half-precision (mixed-precision) VAE. VAE for SDXL seems to produce NaNs in some cases. This option is useful to avoid the NaNs. + self.cache_text_encoder_outputs: bool = False # Cache the outputs of the text encoders. This option is useful to reduce the GPU memory usage. This option cannot be used with options for shuffling or dropping the captions. + self.cache_text_encoder_outputs_to_disk: bool = False # Cache the outputs of the text encoders. This option is useful to reduce the GPU memory usage. This option cannot be used with options for shuffling or dropping the captions. + self.ext_sd_script_args: str = "" # Append or override the sd_script args. (e.g. `--lr_scheduler="constant_with_warmup" --max_grad_norm=0.0`) # Creates the dict that is used for the rest of the code, to facilitate easier json saving and loading def convert_args_to_dict(self): @@ -196,32 +201,11 @@ def convert_args_to_dict(self): def create_args(self) -> argparse.Namespace: parser = Parser() args = self.convert_args_to_dict() - multi_path = args['multi_run_folder'] - if multi_path and ensure_path(multi_path, "multi_run_folder"): - for file in os.listdir(multi_path): - if os.path.isdir(file) or file.split(".")[-1] != "json": - continue - args = self.convert_args_to_dict() - args['json_load_skip_list'] = None - try: - ensure_file_paths(args) - except FileNotFoundError: - print("failed to find one or more folders or paths, skipping.") - continue - if args['tag_occurrence_txt_file']: - get_occurrence_of_tags(args) - args = parser.create_args(self.change_dict_to_internal_names(args)) - train_network.train(args) - gc.collect() - torch.cuda.empty_cache() - if not os.path.exists(os.path.join(multi_path, "complete")): - os.makedirs(os.path.join(multi_path, "complete")) - os.rename(os.path.join(multi_path, file), os.path.join(multi_path, "complete", file)) - print("completed all training") - quit() ensure_file_paths(args) if args['tag_occurrence_txt_file']: get_occurrence_of_tags(args) + if self.use_sdxl: + self.no_half_vae = True args = parser.create_args(self.change_dict_to_internal_names(args)) return args @@ -283,7 +267,7 @@ def find_max_steps(args: dict) -> int: def ensure_file_paths(args: dict) -> None: failed_to_find = False - folders_to_check = ['img_folder', 'output_folder', 'save_json_folder', 'multi_run_folder', + folders_to_check = ['img_folder', 'output_folder', 'save_json_folder', 'reg_img_folder', 'log_dir'] for folder in folders_to_check: if folder in args and args[folder] is not None: @@ -360,15 +344,33 @@ def ensure_path(path, name, ext_list=None) -> bool: class Parser: def __init__(self) -> None: - self.parser = train_network.setup_parser() + parser = train_network.setup_parser() + parser.add_argument( + "--cache_text_encoder_outputs", action="store_true", help="cache text encoder outputs / text encoderの出力をキャッシュする" + ) + parser.add_argument( + "--cache_text_encoder_outputs_to_disk", + action="store_true", + help="cache text encoder outputs to disk / text encoderの出力をディスクにキャッシュする", + ) + self.parser = parser def create_args(self, args: dict) -> argparse.Namespace: remove_epochs = False - args_list = [] - skip_list = ["save_json_folder", "load_json_path", "multi_run_folder", "json_load_skip_list", + args_list: list[str] = [] + skip_list = ["save_json_folder", "load_json_path", "json_load_skip_list", "tag_occurrence_txt_file", "sort_tag_occurrence_alphabetically", "save_json_only", "warmup_lr_ratio", "optimizer_args", "locon_dim", "locon_alpha", "locon", "lyco", "network_args", - "resolution", "height_resolution"] + "resolution", "height_resolution", "use_sdxl", "ext_sd_script_args"] + + # decode ext_sd_script_args + if "ext_sd_script_args" in args and args["ext_sd_script_args"]: + ext_sd_script_args = args["ext_sd_script_args"].split("--") + for arg in ext_sd_script_args: + if not arg: + continue + args_list.append(f"--{arg}") + for key, value in args.items(): if not value: continue @@ -376,6 +378,15 @@ def create_args(self, args: dict) -> argparse.Namespace: continue if key == "max_train_steps": remove_epochs = True + # check key is in the parser + already_exists = False + for arg in args_list: + if arg.startswith(f"--{key}"): + already_exists = True + break + if already_exists: + printD(f"Skipping {key} as it already exists") + continue if isinstance(value, bool): args_list.append(f"--{key}") else: diff --git a/liasece_sd_webui_train_tools/config_file.py b/liasece_sd_webui_train_tools/config_file.py index c10942e..1f8b1d9 100644 --- a/liasece_sd_webui_train_tools/config_file.py +++ b/liasece_sd_webui_train_tools/config_file.py @@ -35,12 +35,14 @@ {"train_finish_generate_all_checkpoint_preview": True}, {"train_optimizer_type": ["Lion"]}, {"train_learning_rate": "0.0001"}, + {"sd_script_args": ""}, {"train_net_dim": 128}, {"train_alpha": 64}, {"train_clip_skip": 2}, {"train_mixed_precision": "fp16"}, {"train_xformers": True}, {"train_base_on_sd_v2": False}, + {"use_sdxl": False}, # preview {"preview_include_sub_img": False}, {"preview_txt2img_prompt": "best quality,Amazing,finely detail,extremely detailed CG unity 8k wallpaper"}, diff --git a/liasece_sd_webui_train_tools/sd_scripts b/liasece_sd_webui_train_tools/sd_scripts index 4c6f312..e69d341 160000 --- a/liasece_sd_webui_train_tools/sd_scripts +++ b/liasece_sd_webui_train_tools/sd_scripts @@ -1 +1 @@ -Subproject commit 4c6f3125fce31f86aa949026030e923f3d7d826b +Subproject commit e69d34103bfd588958cf26bf98694c06cb8976aa diff --git a/liasece_sd_webui_train_tools/train.py b/liasece_sd_webui_train_tools/train.py index db258d8..32cecd4 100644 --- a/liasece_sd_webui_train_tools/train.py +++ b/liasece_sd_webui_train_tools/train.py @@ -6,9 +6,11 @@ import pkg_resources from liasece_sd_webui_train_tools.ArgsList import ArgStore +from liasece_sd_webui_train_tools.util import * from modules import script_loading import liasece_sd_webui_train_tools.sd_scripts.train_network as train_network +import liasece_sd_webui_train_tools.sd_scripts.sdxl_train_network as sdxl_train_network import liasece_sd_webui_train_tools.PythonContextWarper as pc import liasece_sd_webui_train_tools.util as util @@ -34,4 +36,9 @@ def train(cfg: ArgStore) -> None: sub_module=["library", "networks"], ): # begin training - train_network.train(args) + if cfg.use_sdxl: + trainer = sdxl_train_network.SdxlNetworkTrainer() + else: + trainer = train_network.NetworkTrainer() + printD("train begin", args) + trainer.train(args) diff --git a/liasece_sd_webui_train_tools/train_ui.py b/liasece_sd_webui_train_tools/train_ui.py index f7a9c97..1ff6d7e 100644 --- a/liasece_sd_webui_train_tools/train_ui.py +++ b/liasece_sd_webui_train_tools/train_ui.py @@ -31,12 +31,14 @@ def on_train_begin_click(id: str, project: str, version: str, train_finish_generate_all_checkpoint_preview: bool, train_optimizer_type: list[str], train_learning_rate: str, + sd_script_args: str, train_net_dim: int, train_alpha: int, train_clip_skip: int, train_mixed_precision: str, train_xformers: bool, train_base_on_sd_v2: bool, + use_sdxl: bool, # use sdxl # preview view config preview_include_sub_img: bool, # txt2txt @@ -61,12 +63,14 @@ def on_train_begin_click(id: str, project: str, version: str, "train_finish_generate_all_checkpoint_preview": train_finish_generate_all_checkpoint_preview, "train_optimizer_type": train_optimizer_type, "train_learning_rate": train_learning_rate, + "sd_script_args": sd_script_args, "train_net_dim": int(train_net_dim), "train_alpha": int(train_alpha), "train_clip_skip": int(train_clip_skip), "train_mixed_precision": train_mixed_precision, "train_xformers": train_xformers, "train_base_on_sd_v2": train_base_on_sd_v2, + "use_sdxl": use_sdxl, }) save_preview_config(project, version, { # preview view config @@ -119,6 +123,8 @@ def on_train_begin_click(id: str, project: str, version: str, cfg.mixed_precision = train_mixed_precision cfg.xformers = train_xformers cfg.v2 = train_base_on_sd_v2 + cfg.use_sdxl = use_sdxl + cfg.ext_sd_script_args = sd_script_args # check if reg path exist if os.path.exists(os.path.join(processed_path, "..", "reg")): cfg.reg_img_folder = os.path.abspath(os.path.join(processed_path, "..", "reg")) diff --git a/liasece_sd_webui_train_tools/ui.py b/liasece_sd_webui_train_tools/ui.py index 0222d2f..a72e130 100644 --- a/liasece_sd_webui_train_tools/ui.py +++ b/liasece_sd_webui_train_tools/ui.py @@ -130,6 +130,7 @@ def new_ui(): train_base_model_refresh_button = ui.ToolButton(value=ui.refresh_symbol, elem_id="train_base_model_refresh_button") with gr.Row(): train_base_on_sd_v2 = gr.Checkbox(label="Base on Stable Diffusion V2", value=False, elem_id="train_base_on_sd_v2", interactive = True) + use_sdxl = gr.Checkbox(label="Base on Stable Diffusion XL", value=False, elem_id="use_sdxl", interactive = True) with gr.Row(): train_xformers = gr.Checkbox(label="Use xformers", value=True, elem_id="train_xformers", interactive = True) with gr.Row(): @@ -139,11 +140,12 @@ def new_ui(): with gr.Column(): train_batch_size = gr.Number(value=1, label="Batch size", elem_id="train_batch_size", interactive = True) train_num_epochs = gr.Number(value=40, label="Number of epochs", elem_id="train_num_epochs", interactive = True) - train_learning_rate = gr.Textbox(value="0.0001", label="Learning rate", elem_id="train_learning_rate", interactive = True) + train_learning_rate = gr.Textbox(value="0.0001", label="Learning rate(Multi-select e.g. 0.0001,0.0002)", elem_id="train_learning_rate", interactive = True) + sd_script_args = gr.Textbox(value="", label="Append or override the sd_script args. (e.g. `--lr_scheduler=\"constant_with_warmup\" --max_grad_norm=0.0`)", elem_id="sd_script_args", interactive = True) with gr.Column(): train_net_dim = gr.Number(value=128, label="Net dim (128 ~ 144MB)", elem_id="train_net_dim", interactive = True) train_alpha = gr.Number(value=64, label="Alpha (default is half of Net dim)", elem_id="train_alpha", interactive = True) - train_optimizer_type = gr.Dropdown(label="Optimizer type",value=["Lion"], choices=["Adam", "AdamW", "AdamW8bit", "Lion", "SGDNesterov", "SGDNesterov8bit", "DAdaptation", "AdaFactor"], multiselect = True, interactive = True, elem_id="train_optimizer_type") + train_optimizer_type = gr.Dropdown(label="Optimizer type(Multi-select)",value=["Lion"], choices=["Adam", "AdamW", "AdamW8bit", "Lion", "SGDNesterov", "SGDNesterov8bit", "DAdaptation", "AdaFactor"], multiselect = True, interactive = True, elem_id="train_optimizer_type") train_mixed_precision = gr.Dropdown(label="Mixed precision (If your graphics card supports bf16 better)",value="fp16", choices=["fp16", "bf16"], interactive = True, elem_id="train_mixed_precision") with gr.Row(): with gr.Column(scale=2): @@ -248,12 +250,14 @@ def train_config_inputs(): train_finish_generate_all_checkpoint_preview, train_optimizer_type, train_learning_rate, + sd_script_args, train_net_dim, train_alpha, train_clip_skip, train_mixed_precision, train_xformers, train_base_on_sd_v2, + use_sdxl, ] def preview_config_inputs(): return [ diff --git a/requirements.txt b/requirements.txt index d33444d..f146bbc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,10 +1,10 @@ -accelerate==0.15.0 -transformers +accelerate==0.25.0 +diffusers[torch]==0.25.0 +transformers==4.36.2 ftfy albumentations opencv-python einops -diffusers[torch]==0.10.2 pytorch-lightning bitsandbytes tensorboard