From 19ca9782c4e0b9d12e8822a76ccc47904125c231 Mon Sep 17 00:00:00 2001 From: bmaltais Date: Tue, 9 Apr 2024 20:38:19 -0400 Subject: [PATCH 01/10] Relocate preset dropdown --- kohya_gui/finetune_gui.py | 5 +++-- kohya_gui/lora_gui.py | 18 +++++++++--------- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/kohya_gui/finetune_gui.py b/kohya_gui/finetune_gui.py index d9e99b529..2191486f2 100644 --- a/kohya_gui/finetune_gui.py +++ b/kohya_gui/finetune_gui.py @@ -810,8 +810,9 @@ def list_presets(path): training_preset = gr.Dropdown( label="Presets", - choices=[""] + list_presets(f"{presets_dir}/finetune"), - elem_id="myDropdown", + choices=["none"] + list_presets(f"{presets_dir}/finetune"), + # elem_id="myDropdown", + value="none", ) with gr.Accordion("Basic", open="True"): diff --git a/kohya_gui/lora_gui.py b/kohya_gui/lora_gui.py index 29a78e5a2..fb4dee1b7 100644 --- a/kohya_gui/lora_gui.py +++ b/kohya_gui/lora_gui.py @@ -1126,7 +1126,7 @@ def list_presets(path): json_files = [] # Insert an empty string at the beginning - json_files.insert(0, "none") + #json_files.insert(0, "none") for file in os.listdir(path): if file.endswith(".json"): @@ -1141,14 +1141,14 @@ def list_presets(path): return json_files - with gr.Accordion("Basic", open="True"): - training_preset = gr.Dropdown( - label="Presets", - choices=[""] + list_presets(rf"{presets_dir}/lora"), - elem_id="myDropdown", - value="none", - ) + training_preset = gr.Dropdown( + label="Presets", + choices=["none"] + list_presets(rf"{presets_dir}/lora"), + # elem_id="myDropdown", + value="none", + ) + with gr.Accordion("Basic", open="True"): with gr.Group(elem_id="basic_tab"): with gr.Row(): LoRA_type = gr.Dropdown( @@ -1399,7 +1399,7 @@ def update_LoRA_settings( conv_dim, network_dim, ): - log.info("LoRA type changed...") + log.debug("LoRA type changed...") lora_settings_config = { "network_row": { From 49c34a0ac2b15a6a82061187295fd18618494ef6 Mon Sep 17 00:00:00 2001 From: bmaltais Date: Wed, 10 Apr 2024 18:56:57 -0400 Subject: [PATCH 02/10] Fix issue with WD14 prefix --- kohya_gui/wd14_caption_gui.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/kohya_gui/wd14_caption_gui.py b/kohya_gui/wd14_caption_gui.py index 50e3daaba..8f97f02f8 100644 --- a/kohya_gui/wd14_caption_gui.py +++ b/kohya_gui/wd14_caption_gui.py @@ -1,7 +1,7 @@ import gradio as gr from easygui import msgbox import subprocess -from .common_gui import get_folder_path, scriptdir, list_dirs +from .common_gui import get_folder_path, add_pre_postfix, scriptdir, list_dirs from .class_gui_config import KohyaSSGUIConfig import os @@ -46,8 +46,8 @@ def caption_images( log.info(f"Captioning files in {train_data_dir}...") run_cmd = rf'accelerate launch "{scriptdir}/sd-scripts/finetune/tag_images_by_wd14_tagger.py"' - if always_first_tags: - run_cmd += f' --always_first_tags="{always_first_tags}"' + # if always_first_tags: + # run_cmd += f' --always_first_tags="{always_first_tags}"' if append_tags: run_cmd += f" --append_tags" run_cmd += f" --batch_size={int(batch_size)}" @@ -73,8 +73,8 @@ def caption_images( if remove_underscore: run_cmd += f" --remove_underscore" run_cmd += f' --repo_id="{repo_id}"' - if tag_replacement: - run_cmd += f" --tag_replacement" + if not tag_replacement == "": + run_cmd += f" --tag_replacement={tag_replacement}" if not thresh == 0.35: run_cmd += f" --thresh={thresh}" if not undesired_tags == "": @@ -95,6 +95,13 @@ def caption_images( # Run the command subprocess.run(run_cmd, shell=True, env=env) + + # Add prefix and postfix + add_pre_postfix( + folder=train_data_dir, + caption_file_ext=caption_extension, + prefix=always_first_tags, + ) log.info("...captioning done") @@ -220,7 +227,7 @@ def list_train_dirs(path): with gr.Row(): always_first_tags = gr.Textbox( label="Prefix to add to WD14 caption", - info="comma-separated list of tags to always put at the beginning, e.g. 1girl,1boy", + info="comma-separated list of tags to always put at the beginning, e.g. 1girl, 1boy, ", placeholder="(Optional)", interactive=True, value=config.get("wd14_caption.always_first_tags", ""), @@ -306,7 +313,7 @@ def list_train_dirs(path): # Advanced Settings with gr.Row(): batch_size = gr.Number( - value=config.get("wd14_caption.batch_size", 8), + value=config.get("wd14_caption.batch_size", 1), label="Batch size", interactive=True, ) From fd8bfa8a39b1c08a2f39d04f93a8288285d0c8f7 Mon Sep 17 00:00:00 2001 From: bmaltais Date: Wed, 10 Apr 2024 20:48:12 -0400 Subject: [PATCH 03/10] Upgrade Gradio release --- .release | 2 +- README.md | 6 ++++ kohya_gui/blip2_caption_gui.py | 2 +- kohya_gui/class_basic_training.py | 4 +-- kohya_gui/class_source_model.py | 2 +- kohya_gui/dreambooth_gui.py | 14 +++----- kohya_gui/finetune_gui.py | 15 +++------ kohya_gui/group_images_gui.py | 2 +- kohya_gui/lora_gui.py | 25 +++++---------- kohya_gui/manual_caption_gui.py | 8 ++--- kohya_gui/textual_inversion_gui.py | 20 ++++-------- requirements.txt | 2 +- test/config/dreambooth-AdamW.json | 44 +++++++++++++++++++++++--- test/config/finetune-AdamW.json | 17 ++++++++-- test/config/locon-AdamW.json | 51 ++++++++++++++++++++++++++---- 15 files changed, 138 insertions(+), 76 deletions(-) diff --git a/.release b/.release index 76a0279d3..ab61f4fb4 100644 --- a/.release +++ b/.release @@ -1 +1 @@ -v23.1.4 \ No newline at end of file +v23.1.5 \ No newline at end of file diff --git a/README.md b/README.md index 6e2cb4f41..7e6a3c3f6 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,7 @@ The GUI allows you to set the training parameters and generate and run the requi - [SDXL training](#sdxl-training) - [Masked loss](#masked-loss) - [Change History](#change-history) + - [2024/04/10 (v23.1.5)](#20240410-v2315) - [2024/04/08 (v23.1.4)](#20240408-v2314) - [2024/04/08 (v23.1.3)](#20240408-v2313) - [2024/04/08 (v23.1.2)](#20240408-v2312) @@ -407,6 +408,11 @@ ControlNet dataset is used to specify the mask. The mask images should be the RG ## Change History +### 2024/04/10 (v23.1.5) + +- Fix issue with Textual Inversion configuration file selection. +- Upgrade to gradio 4.19.2 to fix several high security risks associated to earlier versions. Hoping this will not introduce undorseen issues. + ### 2024/04/08 (v23.1.4) - Relocate config accordion to the top of the GUI. diff --git a/kohya_gui/blip2_caption_gui.py b/kohya_gui/blip2_caption_gui.py index 7fd866b7e..6336a58ff 100644 --- a/kohya_gui/blip2_caption_gui.py +++ b/kohya_gui/blip2_caption_gui.py @@ -332,7 +332,7 @@ def list_train_dirs(path): ) top_p = gr.Slider( - minimum=-0, + minimum=0, maximum=1, value=0.9, step=0.1, diff --git a/kohya_gui/class_basic_training.py b/kohya_gui/class_basic_training.py index 0df6ddf0a..8b1a040f6 100644 --- a/kohya_gui/class_basic_training.py +++ b/kohya_gui/class_basic_training.py @@ -19,9 +19,9 @@ class BasicTraining: def __init__( self, sdxl_checkbox: gr.Checkbox, - learning_rate_value: str = "1e-6", + learning_rate_value: float = "1e-6", lr_scheduler_value: str = "constant", - lr_warmup_value: str = "0", + lr_warmup_value: float = "0", finetuning: bool = False, dreambooth: bool = False, config: dict = {}, diff --git a/kohya_gui/class_source_model.py b/kohya_gui/class_source_model.py index a390e42fd..ddab2ddae 100644 --- a/kohya_gui/class_source_model.py +++ b/kohya_gui/class_source_model.py @@ -257,7 +257,7 @@ def list_dataset_config_dirs(path: str) -> list: min_width=60, ) with gr.Column(): - gr.Box(visible=False) + gr.Group(visible=False) with gr.Row(): self.training_comment = gr.Textbox( diff --git a/kohya_gui/dreambooth_gui.py b/kohya_gui/dreambooth_gui.py index f578604fb..927f68d0a 100644 --- a/kohya_gui/dreambooth_gui.py +++ b/kohya_gui/dreambooth_gui.py @@ -47,7 +47,7 @@ def save_configuration( - save_as, + save_as_bool, file_path, pretrained_model_name_or_path, v2, @@ -160,8 +160,6 @@ def save_configuration( original_file_path = file_path - save_as_bool = True if save_as.get("label") == "True" else False - if save_as_bool: log.info("Save as...") file_path = get_saveasfile_path(file_path) @@ -301,8 +299,6 @@ def open_configuration( # Get list of function parameters and values parameters = list(locals().items()) - ask_for_file = True if ask_for_file.get("label") == "True" else False - original_file_path = file_path if ask_for_file: @@ -740,8 +736,8 @@ def dreambooth_tab( headless=False, config: KohyaSSGUIConfig = {}, ): - dummy_db_true = gr.Label(value=True, visible=False) - dummy_db_false = gr.Label(value=False, visible=False) + dummy_db_true = gr.Checkbox(value=True, visible=False) + dummy_db_false = gr.Checkbox(value=False, visible=False) dummy_headless = gr.Label(value=headless, visible=False) with gr.Tab("Training"), gr.Column(variant="compact"): @@ -764,9 +760,9 @@ def dreambooth_tab( with gr.Accordion("Basic", open="True"): with gr.Group(elem_id="basic_tab"): basic_training = BasicTraining( - learning_rate_value="1e-5", + learning_rate_value=1e-5, lr_scheduler_value="cosine", - lr_warmup_value="10", + lr_warmup_value=10, dreambooth=True, sdxl_checkbox=source_model.sdxl_checkbox, config=config, diff --git a/kohya_gui/finetune_gui.py b/kohya_gui/finetune_gui.py index 2191486f2..1cb8a7b31 100644 --- a/kohya_gui/finetune_gui.py +++ b/kohya_gui/finetune_gui.py @@ -53,7 +53,7 @@ def save_configuration( - save_as, + save_as_bool, file_path, pretrained_model_name_or_path, v2, @@ -170,8 +170,6 @@ def save_configuration( original_file_path = file_path - save_as_bool = True if save_as.get("label") == "True" else False - if save_as_bool: log.info("Save as...") file_path = get_saveasfile_path(file_path) @@ -319,9 +317,6 @@ def open_configuration( # Get list of function parameters and values parameters = list(locals().items()) - ask_for_file = True if ask_for_file.get("label") == "True" else False - apply_preset = True if apply_preset.get("label") == "True" else False - # Check if we are "applying" a preset or a config if apply_preset: log.info(f"Applying preset {training_preset}...") @@ -764,8 +759,8 @@ def train_model( def finetune_tab(headless=False, config: dict = {}): - dummy_db_true = gr.Label(value=True, visible=False) - dummy_db_false = gr.Label(value=False, visible=False) + dummy_db_true = gr.Checkbox(value=True, visible=False) + dummy_db_false = gr.Checkbox(value=False, visible=False) dummy_headless = gr.Label(value=headless, visible=False) with gr.Tab("Training"), gr.Column(variant="compact"): gr.Markdown("Train a custom model using kohya finetune python code...") @@ -818,7 +813,7 @@ def list_presets(path): with gr.Accordion("Basic", open="True"): with gr.Group(elem_id="basic_tab"): basic_training = BasicTraining( - learning_rate_value="1e-5", + learning_rate_value=1e-5, finetuning=True, sdxl_checkbox=source_model.sdxl_checkbox, config=config, @@ -837,7 +832,7 @@ def list_presets(path): with gr.Row(): gradient_accumulation_steps = gr.Number( label="Gradient accumulate steps", - value="1", + value=1, ) block_lr = gr.Textbox( label="Block LR (SDXL)", diff --git a/kohya_gui/group_images_gui.py b/kohya_gui/group_images_gui.py index bbd78d852..55552e300 100644 --- a/kohya_gui/group_images_gui.py +++ b/kohya_gui/group_images_gui.py @@ -146,7 +146,7 @@ def list_output_dirs(path): group_size = gr.Slider( label="Group size", info="Number of images to group together", - value="4", + value=4, minimum=1, maximum=64, step=1, diff --git a/kohya_gui/lora_gui.py b/kohya_gui/lora_gui.py index fb4dee1b7..bfc540356 100644 --- a/kohya_gui/lora_gui.py +++ b/kohya_gui/lora_gui.py @@ -83,7 +83,7 @@ def update_network_args_with_kohya_lora_vars( def save_configuration( - save_as, + save_as_bool, file_path, pretrained_model_name_or_path, v2, @@ -233,9 +233,6 @@ def save_configuration( original_file_path = file_path - # Determine whether to save as a new file or overwrite the existing file - save_as_bool = True if save_as.get("label") == "True" else False - # If saving as a new file, get the file path for saving if save_as_bool: log.info("Save as...") @@ -423,12 +420,6 @@ def open_configuration( # Get list of function parameters and values parameters = list(locals().items()) - # Convert 'ask_for_file' and 'apply_preset' from string to boolean based on their 'label' value - # This corrects a critical oversight in the original code, where `.get("label")` method calls were - # made on boolean variables instead of dictionaries - ask_for_file = True if ask_for_file.get("label") == "True" else False - apply_preset = True if apply_preset.get("label") == "True" else False - # Determines if a preset configuration is being applied if apply_preset: if training_preset != "none": @@ -1091,8 +1082,8 @@ def lora_tab( headless=False, config: dict = {}, ): - dummy_db_true = gr.Label(value=True, visible=False) - dummy_db_false = gr.Label(value=False, visible=False) + dummy_db_true = gr.Checkbox(value=True, visible=False) + dummy_db_false = gr.Checkbox(value=False, visible=False) dummy_headless = gr.Label(value=headless, visible=False) with gr.Tab("Training"), gr.Column(variant="compact") as tab: @@ -1210,9 +1201,9 @@ def list_presets(path): info="Automatically determine the dim(rank) from the weight file.", ) basic_training = BasicTraining( - learning_rate_value="0.0001", + learning_rate_value=0.0001, lr_scheduler_value="cosine", - lr_warmup_value="10", + lr_warmup_value=10, sdxl_checkbox=source_model.sdxl_checkbox, config=config, ) @@ -1220,7 +1211,7 @@ def list_presets(path): with gr.Row(): text_encoder_lr = gr.Number( label="Text Encoder learning rate", - value="0.0001", + value=0.0001, info="(Optional)", minimum=0, maximum=1, @@ -1228,7 +1219,7 @@ def list_presets(path): unet_lr = gr.Number( label="Unet learning rate", - value="0.0001", + value=0.0001, info="(Optional)", minimum=0, maximum=1, @@ -1286,7 +1277,7 @@ def list_presets(path): visible=False, ) constrain = gr.Number( - value="0.0", + value=0.0, label="Constrain OFT", info="Limits the norm of the oft_blocks, ensuring that their magnitude does not exceed a specified threshold, thus controlling the extent of the transformation applied.", visible=False, diff --git a/kohya_gui/manual_caption_gui.py b/kohya_gui/manual_caption_gui.py index 2a80bfc06..bc90d2158 100644 --- a/kohya_gui/manual_caption_gui.py +++ b/kohya_gui/manual_caption_gui.py @@ -269,8 +269,8 @@ def list_images_dirs(path): with gr.Tab("Manual Captioning"): gr.Markdown("This utility allows quick captioning and tagging of images.") - page = gr.Number(-1, visible=False) - max_page = gr.Number(1, visible=False) + page = gr.Number(value=-1, visible=False) + max_page = gr.Number(value=1, visible=False) loaded_images_dir = gr.Text(visible=False) with gr.Group(), gr.Row(): images_dir = gr.Dropdown( @@ -336,7 +336,7 @@ def list_images_dirs(path): def render_pagination(): gr.Button("< Prev", elem_id="open_folder").click( paginate, - inputs=[page, max_page, gr.Number(-1, visible=False)], + inputs=[page, max_page, gr.Number(value=-1, visible=False)], outputs=[page], ) page_count = gr.Label("Page 1", label="Page") @@ -352,7 +352,7 @@ def render_pagination(): ) gr.Button("Next >", elem_id="open_folder").click( paginate, - inputs=[page, max_page, gr.Number(1, visible=False)], + inputs=[page, max_page, gr.Number(value=1, visible=False)], outputs=[page], ) return page_count diff --git a/kohya_gui/textual_inversion_gui.py b/kohya_gui/textual_inversion_gui.py index 49fc1761b..1505797a2 100644 --- a/kohya_gui/textual_inversion_gui.py +++ b/kohya_gui/textual_inversion_gui.py @@ -47,7 +47,7 @@ def save_configuration( - save_as, + save_as_bool, file_path, pretrained_model_name_or_path, v2, @@ -159,8 +159,6 @@ def save_configuration( original_file_path = file_path - save_as_bool = True if save_as.get("label") == "True" else False - if save_as_bool: log.info("Save as...") file_path = get_saveasfile_path(file_path) @@ -301,8 +299,6 @@ def open_configuration( # Get list of function parameters and values parameters = list(locals().items()) - ask_for_file = True if ask_for_file.get("label") == "True" else False - original_file_path = file_path if ask_for_file: @@ -710,8 +706,8 @@ def train_model( def ti_tab(headless=False, default_output_dir=None, config: dict = {}): - dummy_db_true = gr.Label(value=True, visible=False) - dummy_db_false = gr.Label(value=False, visible=False) + dummy_db_true = gr.Checkbox(value=True, visible=False) + dummy_db_false = gr.Checkbox(value=False, visible=False) dummy_headless = gr.Label(value=headless, visible=False) current_embedding_dir = ( @@ -725,7 +721,7 @@ def ti_tab(headless=False, default_output_dir=None, config: dict = {}): # Setup Configuration Files Gradio with gr.Accordion("Configuration", open=False): - configuration = ConfigurationFile(headless=headless) + configuration = ConfigurationFile(headless=headless, config=config) with gr.Accordion("Accelerate launch", open=False), gr.Column(): accelerate_launch = AccelerateLaunch(config=config) @@ -824,9 +820,9 @@ def list_embedding_files(path): value="caption", ) basic_training = BasicTraining( - learning_rate_value="1e-5", + learning_rate_value=1e-5, lr_scheduler_value="cosine", - lr_warmup_value="10", + lr_warmup_value=10, sdxl_checkbox=source_model.sdxl_checkbox, config=config, ) @@ -863,10 +859,6 @@ def list_embedding_files(path): ) gradio_dataset_balancing_tab(headless=headless) - # Setup Configuration Files Gradio - with gr.Accordion("Configuration", open=False): - configuration = ConfigurationFile(headless=headless) - with gr.Column(), gr.Group(): with gr.Row(): button_run = gr.Button("Start training", variant="primary") diff --git a/requirements.txt b/requirements.txt index ce7404eb3..375ec4687 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,7 +7,7 @@ easygui==0.98.3 einops==0.7.0 fairscale==0.4.13 ftfy==6.1.1 -gradio==3.50.2 +gradio==4.19.2 huggingface-hub==0.20.1 imagesize==1.4.1 invisible-watermark==0.2.0 diff --git a/test/config/dreambooth-AdamW.json b/test/config/dreambooth-AdamW.json index 1f41f7d36..6ca69acf0 100644 --- a/test/config/dreambooth-AdamW.json +++ b/test/config/dreambooth-AdamW.json @@ -5,43 +5,72 @@ "bucket_reso_steps": 64, "cache_latents": true, "cache_latents_to_disk": false, - "caption_dropout_every_n_epochs": 0.0, + "caption_dropout_every_n_epochs": 0, "caption_dropout_rate": 0.05, "caption_extension": "", "clip_skip": 2, "color_aug": false, + "dataset_config": "", + "debiased_estimation_loss": false, "enable_bucket": true, "epoch": 1, + "extra_accelerate_launch_args": "", "flip_aug": false, + "full_bf16": false, "full_fp16": false, - "gradient_accumulation_steps": 1.0, + "gpu_ids": "", + "gradient_accumulation_steps": 1, "gradient_checkpointing": false, + "huber_c": 0.1, + "huber_schedule": "snr", + "ip_noise_gamma": 0, + "ip_noise_gamma_random_strength": false, "keep_tokens": "0", "learning_rate": 5e-05, + "learning_rate_te": 1e-05, + "learning_rate_te1": 1e-05, + "learning_rate_te2": 1e-05, + "log_tracker_config": "", + "log_tracker_name": "", "logging_dir": "./test/logs", + "loss_type": "l2", "lr_scheduler": "constant", + "lr_scheduler_args": "", + "lr_scheduler_num_cycles": "", + "lr_scheduler_power": "", "lr_warmup": 0, + "main_process_port": 0, + "masked_loss": false, + "max_bucket_reso": 2048, "max_data_loader_n_workers": "0", "max_resolution": "512,512", + "max_timestep": 1000, "max_token_length": "75", "max_train_epochs": "", + "max_train_steps": "", "mem_eff_attn": false, + "min_bucket_reso": 256, "min_snr_gamma": 0, + "min_timestep": 0, "mixed_precision": "bf16", "model_list": "runwayml/stable-diffusion-v1-5", + "multi_gpu": false, "multires_noise_discount": 0, "multires_noise_iterations": 0, "no_token_padding": false, - "noise_offset": "0.05", + "noise_offset": 0.05, + "noise_offset_random_strength": false, "noise_offset_type": "Original", "num_cpu_threads_per_process": 2, + "num_machines": 1, + "num_processes": 1, "optimizer": "AdamW", "optimizer_args": "", "output_dir": "./test/output", "output_name": "db-AdamW", "persistent_data_loader_workers": false, "pretrained_model_name_or_path": "runwayml/stable-diffusion-v1-5", - "prior_loss_weight": 1.0, + "prior_loss_weight": 1, "random_crop": false, "reg_data_dir": "", "resume": "", @@ -49,6 +78,7 @@ "sample_every_n_steps": 25, "sample_prompts": "a painting of a gas mask , by darius kawasaki", "sample_sampler": "euler_a", + "save_as_bool": false, "save_every_n_epochs": 1, "save_every_n_steps": 0, "save_last_n_steps": 0, @@ -56,7 +86,9 @@ "save_model_as": "safetensors", "save_precision": "fp16", "save_state": false, + "save_state_on_train_end": false, "scale_v_pred_loss_like_noise_pred": false, + "sdxl": false, "seed": "1234", "shuffle_caption": false, "stop_text_encoder_training": 0, @@ -65,9 +97,11 @@ "use_wandb": false, "v2": false, "v_parameterization": false, + "v_pred_like_loss": 0, "vae": "", "vae_batch_size": 0, "wandb_api_key": "", + "wandb_run_name": "", "weighted_captions": false, - "xformers": true + "xformers": "xformers" } \ No newline at end of file diff --git a/test/config/finetune-AdamW.json b/test/config/finetune-AdamW.json index 36a5fff31..6c6cdf10c 100644 --- a/test/config/finetune-AdamW.json +++ b/test/config/finetune-AdamW.json @@ -7,7 +7,7 @@ "bucket_reso_steps": 1, "cache_latents": true, "cache_latents_to_disk": false, - "caption_dropout_every_n_epochs": 0.0, + "caption_dropout_every_n_epochs": 0, "caption_dropout_rate": 0, "caption_extension": ".txt", "caption_metadata_filename": "meta-1_cap.json", @@ -18,26 +18,34 @@ "dataset_config": "", "dataset_repeats": "50", "epoch": 2, + "extra_accelerate_launch_args": "", "flip_aug": false, "full_bf16": false, "full_fp16": false, "full_path": true, "gpu_ids": "", - "gradient_accumulation_steps": 1.0, + "gradient_accumulation_steps": 1, "gradient_checkpointing": false, + "huber_c": 0.1, + "huber_schedule": "snr", "image_folder": ".\\test\\img\\10_darius kawasaki person", + "ip_noise_gamma": 0, + "ip_noise_gamma_random_strength": false, "keep_tokens": 0, "latent_metadata_filename": "meta-1_lat.json", "learning_rate": 1e-05, "learning_rate_te": 5e-06, "learning_rate_te1": 5e-06, - "learning_rate_te2": 0.0, + "learning_rate_te2": 0, "log_tracker_config": "", "log_tracker_name": "", "logging_dir": "./test/ft", + "loss_type": "l2", "lr_scheduler": "cosine_with_restarts", "lr_scheduler_args": "", "lr_warmup": 10, + "main_process_port": 0, + "masked_loss": false, "max_bucket_reso": "1024", "max_data_loader_n_workers": "0", "max_resolution": "512,512", @@ -55,6 +63,7 @@ "multires_noise_discount": 0, "multires_noise_iterations": 0, "noise_offset": 0, + "noise_offset_random_strength": false, "noise_offset_type": "Original", "num_cpu_threads_per_process": 2, "num_machines": 1, @@ -71,6 +80,7 @@ "sample_every_n_steps": 0, "sample_prompts": "", "sample_sampler": "euler_a", + "save_as_bool": false, "save_every_n_epochs": 1, "save_every_n_steps": 0, "save_last_n_steps": 0, @@ -78,6 +88,7 @@ "save_model_as": "safetensors", "save_precision": "bf16", "save_state": false, + "save_state_on_train_end": false, "scale_v_pred_loss_like_noise_pred": false, "sdxl_cache_text_encoder_outputs": false, "sdxl_checkbox": false, diff --git a/test/config/locon-AdamW.json b/test/config/locon-AdamW.json index 6c38f8e16..97b4d29bd 100644 --- a/test/config/locon-AdamW.json +++ b/test/config/locon-AdamW.json @@ -1,5 +1,6 @@ { "LoRA_type": "Kohya LoCon", + "LyCORIS_preset": "full", "adaptive_noise_scale": 0, "additional_parameters": "", "block_alphas": "", @@ -7,71 +8,100 @@ "block_lr_zero_threshold": "", "bucket_no_upscale": true, "bucket_reso_steps": 64, + "bypass_mode": false, "cache_latents": true, "cache_latents_to_disk": false, - "caption_dropout_every_n_epochs": 0.0, + "caption_dropout_every_n_epochs": 0, "caption_dropout_rate": 0.05, "caption_extension": "", "clip_skip": 2, "color_aug": false, + "constrain": 0, "conv_alpha": 8, - "conv_alphas": "", + "conv_block_alphas": "", + "conv_block_dims": "", "conv_dim": 16, - "conv_dims": "", + "dataset_config": "", + "debiased_estimation_loss": false, "decompose_both": false, "dim_from_weights": false, + "dora_wd": false, "down_lr_weight": "", "enable_bucket": true, "epoch": 1, + "extra_accelerate_launch_args": "", "factor": -1, "flip_aug": false, + "fp8_base": false, + "full_bf16": false, "full_fp16": false, + "gpu_ids": "", "gradient_accumulation_steps": 4, "gradient_checkpointing": false, + "huber_c": 0.1, + "huber_schedule": "snr", + "ip_noise_gamma": 0, + "ip_noise_gamma_random_strength": false, "keep_tokens": "0", "learning_rate": 0.0005, + "log_tracker_config": "", + "log_tracker_name": "", "logging_dir": "./test/logs", "lora_network_weights": "", + "loss_type": "l2", "lr_scheduler": "constant", + "lr_scheduler_args": "", "lr_scheduler_num_cycles": "", "lr_scheduler_power": "", "lr_warmup": 0, + "main_process_port": 0, + "masked_loss": false, + "max_bucket_reso": 2048, "max_data_loader_n_workers": "0", + "max_grad_norm": 1, "max_resolution": "512,512", "max_timestep": 1000, "max_token_length": "75", "max_train_epochs": "", + "max_train_steps": "", "mem_eff_attn": false, "mid_lr_weight": "", + "min_bucket_reso": 256, "min_snr_gamma": 0, "min_timestep": 0, "mixed_precision": "bf16", "model_list": "runwayml/stable-diffusion-v1-5", "module_dropout": 0.1, + "multi_gpu": false, "multires_noise_discount": 0, "multires_noise_iterations": 0, "network_alpha": 8, "network_dim": 16, "network_dropout": 0.1, - "no_token_padding": false, - "noise_offset": "0.05", + "noise_offset": 0.05, + "noise_offset_random_strength": false, "noise_offset_type": "Original", "num_cpu_threads_per_process": 2, + "num_machines": 1, + "num_processes": 1, "optimizer": "AdamW", "optimizer_args": "", "output_dir": "./test/output", "output_name": "locon-AdamW", "persistent_data_loader_workers": false, "pretrained_model_name_or_path": "runwayml/stable-diffusion-v1-5", - "prior_loss_weight": 1.0, + "prior_loss_weight": 1, "random_crop": false, "rank_dropout": 0.1, + "rank_dropout_scale": false, "reg_data_dir": "", + "rescaled": false, "resume": "", "sample_every_n_epochs": 0, "sample_every_n_steps": 25, "sample_prompts": "a painting of a gas mask , by darius kawasaki", "sample_sampler": "euler_a", + "save_as_bool": false, "save_every_n_epochs": 1, "save_every_n_steps": 0, "save_last_n_steps": 0, @@ -79,6 +109,7 @@ "save_model_as": "safetensors", "save_precision": "fp16", "save_state": false, + "save_state_on_train_end": false, "scale_v_pred_loss_like_noise_pred": false, "scale_weight_norms": 1, "sdxl": false, @@ -90,17 +121,23 @@ "text_encoder_lr": 0.0001, "train_batch_size": 1, "train_data_dir": "./test/img", + "train_norm": false, "train_on_input": false, "training_comment": "", "unet_lr": 0.0001, "unit": 1, "up_lr_weight": "", "use_cp": false, + "use_scalar": false, + "use_tucker": false, "use_wandb": false, "v2": false, "v_parameterization": false, + "v_pred_like_loss": 0, + "vae": "", "vae_batch_size": 0, "wandb_api_key": "", + "wandb_run_name": "", "weighted_captions": false, - "xformers": true + "xformers": "xformers" } \ No newline at end of file From f7853fb0753d8da24a666ad7f93f3c5f57ce500a Mon Sep 17 00:00:00 2001 From: bmaltais Date: Wed, 10 Apr 2024 21:31:38 -0400 Subject: [PATCH 04/10] Update requirements --- README.md | 2 ++ kohya_gui.py | 5 ++++- requirements.txt | 2 +- 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 7e6a3c3f6..c45a2ce28 100644 --- a/README.md +++ b/README.md @@ -412,6 +412,8 @@ ControlNet dataset is used to specify the mask. The mask images should be the RG - Fix issue with Textual Inversion configuration file selection. - Upgrade to gradio 4.19.2 to fix several high security risks associated to earlier versions. Hoping this will not introduce undorseen issues. +- Upgrade transformers to 4.38.0 to fix a low severity security issue. +- Add explicit --do_not_share parameter to kohya_gui.py to avoid sharing the GUI on platforms like Kaggle. ### 2024/04/08 (v23.1.4) diff --git a/kohya_gui.py b/kohya_gui.py index 1b77b46b9..7704be80c 100644 --- a/kohya_gui.py +++ b/kohya_gui.py @@ -97,10 +97,11 @@ def UI(**kwargs): launch_kwargs["server_port"] = server_port if inbrowser: launch_kwargs["inbrowser"] = inbrowser - if share: + if share and not kwargs.get("do_not_share", False): launch_kwargs["share"] = share launch_kwargs["debug"] = True interface.launch(**launch_kwargs) + interface.launch() if __name__ == "__main__": @@ -141,6 +142,8 @@ def UI(**kwargs): parser.add_argument("--use-ipex", action="store_true", help="Use IPEX environment") parser.add_argument("--use-rocm", action="store_true", help="Use ROCm environment") + + parser.add_argument("--do_not_share", action="store_true", help="Do not share the gradio UI") args = parser.parse_args() diff --git a/requirements.txt b/requirements.txt index 375ec4687..4bb22c2eb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -27,7 +27,7 @@ scipy==1.11.4 timm==0.6.12 tk==0.1.0 toml==0.10.2 -transformers==4.36.2 +transformers==4.38.0 voluptuous==0.13.1 wandb==0.15.11 scipy==1.11.4 From 9c24a385739d7e9dd806e005332449119677f9d0 Mon Sep 17 00:00:00 2001 From: bmaltais Date: Wed, 10 Apr 2024 21:35:11 -0400 Subject: [PATCH 05/10] Update do_not_share logic --- kohya_gui.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/kohya_gui.py b/kohya_gui.py index 7704be80c..7702b105e 100644 --- a/kohya_gui.py +++ b/kohya_gui.py @@ -88,6 +88,7 @@ def UI(**kwargs): server_port = kwargs.get("server_port", 0) inbrowser = kwargs.get("inbrowser", False) share = kwargs.get("share", False) + do_not_share = kwargs.get("do_not_share", False) server_name = kwargs.get("listen") launch_kwargs["server_name"] = server_name @@ -97,8 +98,11 @@ def UI(**kwargs): launch_kwargs["server_port"] = server_port if inbrowser: launch_kwargs["inbrowser"] = inbrowser - if share and not kwargs.get("do_not_share", False): - launch_kwargs["share"] = share + if do_not_share: + launch_kwargs["share"] = False + else: + if share: + launch_kwargs["share"] = share launch_kwargs["debug"] = True interface.launch(**launch_kwargs) interface.launch() From 48a3242e5411417b19ca0505cf0ce06c8e747d84 Mon Sep 17 00:00:00 2001 From: bmaltais Date: Wed, 10 Apr 2024 22:01:09 -0400 Subject: [PATCH 06/10] Update logic --- kohya_gui.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/kohya_gui.py b/kohya_gui.py index 7702b105e..0e11d4b17 100644 --- a/kohya_gui.py +++ b/kohya_gui.py @@ -105,8 +105,6 @@ def UI(**kwargs): launch_kwargs["share"] = share launch_kwargs["debug"] = True interface.launch(**launch_kwargs) - interface.launch() - if __name__ == "__main__": # torch.cuda.set_per_process_memory_fraction(0.48) @@ -158,6 +156,7 @@ def UI(**kwargs): inbrowser=args.inbrowser, server_port=args.server_port, share=args.share, + do_not_share=args.do_not_share, listen=args.listen, headless=args.headless, language=args.language, From 8633484a5a5aebe17c805bb3b46760873ed1f09b Mon Sep 17 00:00:00 2001 From: bmaltais Date: Thu, 11 Apr 2024 07:33:23 -0400 Subject: [PATCH 07/10] Fix bugs, improve layout logic --- kohya_gui.py | 13 +--- kohya_gui/dreambooth_gui.py | 42 ++++++------ kohya_gui/finetune_gui.py | 101 ++++++++++++++--------------- kohya_gui/lora_gui.py | 49 +++++++------- kohya_gui/tensorboard_gui.py | 4 +- kohya_gui/textual_inversion_gui.py | 46 +++++++------ 6 files changed, 118 insertions(+), 137 deletions(-) diff --git a/kohya_gui.py b/kohya_gui.py index 0e11d4b17..ed9d195df 100644 --- a/kohya_gui.py +++ b/kohya_gui.py @@ -149,15 +149,4 @@ def UI(**kwargs): args = parser.parse_args() - UI( - config_file_path=args.config, - username=args.username, - password=args.password, - inbrowser=args.inbrowser, - server_port=args.server_port, - share=args.share, - do_not_share=args.do_not_share, - listen=args.listen, - headless=args.headless, - language=args.language, - ) + UI(**vars(args)) diff --git a/kohya_gui/dreambooth_gui.py b/kohya_gui/dreambooth_gui.py index 927f68d0a..53c3011a5 100644 --- a/kohya_gui/dreambooth_gui.py +++ b/kohya_gui/dreambooth_gui.py @@ -435,11 +435,8 @@ def train_model( # Get list of function parameters and values parameters = list(locals().items()) - print_only_bool = True if print_only.get("label") == "True" else False log.info(f"Start training Dreambooth...") - headless_bool = True if headless.get("label") == "True" else False - # This function validates files or folder paths. Simply add new variables containing file of folder path # to validate below if not validate_paths( @@ -447,7 +444,7 @@ def train_model( pretrained_model_name_or_path=pretrained_model_name_or_path, train_data_dir=train_data_dir, reg_data_dir=reg_data_dir, - headless=headless_bool, + headless=headless, logging_dir=logging_dir, log_tracker_config=log_tracker_config, resume=resume, @@ -456,8 +453,8 @@ def train_model( ): return - if not print_only_bool and check_if_model_exist( - output_name, output_dir, save_model_as, headless=headless_bool + if not print_only and check_if_model_exist( + output_name, output_dir, save_model_as, headless=headless ): return @@ -693,7 +690,7 @@ def train_model( output_dir, ) - if print_only_bool: + if print_only: log.warning( "Here is the trainer command as a reference. It will not be executed:\n" ) @@ -738,7 +735,7 @@ def dreambooth_tab( ): dummy_db_true = gr.Checkbox(value=True, visible=False) dummy_db_false = gr.Checkbox(value=False, visible=False) - dummy_headless = gr.Label(value=headless, visible=False) + dummy_headless = gr.Checkbox(value=headless, visible=False) with gr.Tab("Training"), gr.Column(variant="compact"): gr.Markdown("Train a custom model using kohya dreambooth python code...") @@ -756,6 +753,21 @@ def dreambooth_tab( with gr.Accordion("Folders", open=False), gr.Group(): folders = Folders(headless=headless, config=config) + with gr.Accordion("Dataset Preparation", open=False): + gr.Markdown( + "This section provide Dreambooth tools to help setup your dataset..." + ) + gradio_dreambooth_folder_creation_tab( + train_data_dir_input=source_model.train_data_dir, + reg_data_dir_input=folders.reg_data_dir, + output_dir_input=folders.output_dir, + logging_dir_input=folders.logging_dir, + headless=headless, + config=config, + ) + + gradio_dataset_balancing_tab(headless=headless) + with gr.Accordion("Parameters", open=False), gr.Column(): with gr.Accordion("Basic", open="True"): with gr.Group(elem_id="basic_tab"): @@ -779,20 +791,6 @@ def dreambooth_tab( with gr.Accordion("Samples", open=False, elem_id="samples_tab"): sample = SampleImages(config=config) - with gr.Accordion("Dataset Preparation", open=False): - gr.Markdown( - "This section provide Dreambooth tools to help setup your dataset..." - ) - gradio_dreambooth_folder_creation_tab( - train_data_dir_input=source_model.train_data_dir, - reg_data_dir_input=folders.reg_data_dir, - output_dir_input=folders.output_dir, - logging_dir_input=folders.logging_dir, - headless=headless, - config=config, - ) - gradio_dataset_balancing_tab(headless=headless) - with gr.Column(), gr.Group(): with gr.Row(): button_run = gr.Button("Start training", variant="primary") diff --git a/kohya_gui/finetune_gui.py b/kohya_gui/finetune_gui.py index 1cb8a7b31..cc4b28505 100644 --- a/kohya_gui/finetune_gui.py +++ b/kohya_gui/finetune_gui.py @@ -469,12 +469,11 @@ def train_model( ): # Get list of function parameters and values parameters = list(locals().items()) - - print_only_bool = True if print_only.get("label") == "True" else False + + log.debug(f"headless = {headless} ; print_only = {print_only}") + log.info(f"Start Finetuning...") - headless_bool = True if headless.get("label") == "True" else False - if train_dir != "" and not os.path.exists(train_dir): os.mkdir(train_dir) @@ -482,7 +481,7 @@ def train_model( output_dir=output_dir, pretrained_model_name_or_path=pretrained_model_name_or_path, finetune_image_folder=image_folder, - headless=headless_bool, + headless=headless, logging_dir=logging_dir, log_tracker_config=log_tracker_config, resume=resume, @@ -490,8 +489,8 @@ def train_model( ): return - if not print_only_bool and check_if_model_exist( - output_name, output_dir, save_model_as, headless_bool + if not print_only and check_if_model_exist( + output_name, output_dir, save_model_as, headless ): return @@ -519,7 +518,7 @@ def train_model( rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}" ) - if not print_only_bool: + if not print_only: # Run the command subprocess.run(run_cmd, shell=True, env=env) @@ -552,7 +551,7 @@ def train_model( rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}" ) - if not print_only_bool: + if not print_only: # Run the command subprocess.run(run_cmd, shell=True, env=env) @@ -724,7 +723,7 @@ def train_model( output_dir, ) - if print_only_bool: + if print_only: log.warning( "Here is the trainer command as a reference. It will not be executed:\n" ) @@ -761,7 +760,7 @@ def train_model( def finetune_tab(headless=False, config: dict = {}): dummy_db_true = gr.Checkbox(value=True, visible=False) dummy_db_false = gr.Checkbox(value=False, visible=False) - dummy_headless = gr.Label(value=headless, visible=False) + dummy_headless = gr.Checkbox(value=headless, visible=False) with gr.Tab("Training"), gr.Column(variant="compact"): gr.Markdown("Train a custom model using kohya finetune python code...") @@ -785,6 +784,46 @@ def finetune_tab(headless=False, config: dict = {}): logging_dir = folders.logging_dir train_dir = folders.reg_data_dir + with gr.Accordion("Dataset Preparation", open=False): + with gr.Row(): + max_resolution = gr.Textbox( + label="Resolution (width,height)", value="512,512" + ) + min_bucket_reso = gr.Textbox(label="Min bucket resolution", value="256") + max_bucket_reso = gr.Textbox( + label="Max bucket resolution", value="1024" + ) + batch_size = gr.Textbox(label="Batch size", value="1") + with gr.Row(): + create_caption = gr.Checkbox( + label="Generate caption metadata", value=True + ) + create_buckets = gr.Checkbox( + label="Generate image buckets metadata", value=True + ) + use_latent_files = gr.Dropdown( + label="Use latent files", + choices=[ + "No", + "Yes", + ], + value="Yes", + ) + with gr.Accordion("Advanced parameters", open=False): + with gr.Row(): + caption_metadata_filename = gr.Textbox( + label="Caption metadata filename", + value="meta_cap.json", + ) + latent_metadata_filename = gr.Textbox( + label="Latent metadata filename", value="meta_lat.json" + ) + with gr.Row(): + full_path = gr.Checkbox(label="Use full path", value=True) + weighted_captions = gr.Checkbox( + label="Weighted captions", value=False + ) + with gr.Accordion("Parameters", open=False), gr.Column(): def list_presets(path): @@ -853,46 +892,6 @@ def list_presets(path): with gr.Accordion("Samples", open=False, elem_id="samples_tab"): sample = SampleImages(config=config) - with gr.Accordion("Dataset Preparation", open=False): - with gr.Row(): - max_resolution = gr.Textbox( - label="Resolution (width,height)", value="512,512" - ) - min_bucket_reso = gr.Textbox(label="Min bucket resolution", value="256") - max_bucket_reso = gr.Textbox( - label="Max bucket resolution", value="1024" - ) - batch_size = gr.Textbox(label="Batch size", value="1") - with gr.Row(): - create_caption = gr.Checkbox( - label="Generate caption metadata", value=True - ) - create_buckets = gr.Checkbox( - label="Generate image buckets metadata", value=True - ) - use_latent_files = gr.Dropdown( - label="Use latent files", - choices=[ - "No", - "Yes", - ], - value="Yes", - ) - with gr.Accordion("Advanced parameters", open=False): - with gr.Row(): - caption_metadata_filename = gr.Textbox( - label="Caption metadata filename", - value="meta_cap.json", - ) - latent_metadata_filename = gr.Textbox( - label="Latent metadata filename", value="meta_lat.json" - ) - with gr.Row(): - full_path = gr.Checkbox(label="Use full path", value=True) - weighted_captions = gr.Checkbox( - label="Weighted captions", value=False - ) - with gr.Column(), gr.Group(): with gr.Row(): button_run = gr.Button("Start training", variant="primary") diff --git a/kohya_gui/lora_gui.py b/kohya_gui/lora_gui.py index bfc540356..56986ab8e 100644 --- a/kohya_gui/lora_gui.py +++ b/kohya_gui/lora_gui.py @@ -633,16 +633,14 @@ def train_model( parameters = list(locals().items()) global command_running - print_only_bool = True if print_only.get("label") == "True" else False log.info(f"Start training LoRA {LoRA_type} ...") - headless_bool = True if headless.get("label") == "True" else False if not validate_paths( output_dir=output_dir, pretrained_model_name_or_path=pretrained_model_name_or_path, train_data_dir=train_data_dir, reg_data_dir=reg_data_dir, - headless=headless_bool, + headless=headless, logging_dir=logging_dir, log_tracker_config=log_tracker_config, resume=resume, @@ -655,7 +653,7 @@ def train_model( if int(bucket_reso_steps) < 1: output_message( msg="Bucket resolution steps need to be greater than 0", - headless=headless_bool, + headless=headless, ) return @@ -665,7 +663,7 @@ def train_model( if float(noise_offset) > 1 or float(noise_offset) < 0: output_message( msg="Noise offset need to be a value between 0 and 1", - headless=headless_bool, + headless=headless, ) return @@ -676,12 +674,12 @@ def train_model( if stop_text_encoder_training_pct > 0: output_message( msg='Output "stop text encoder training" is not yet supported. Ignoring', - headless=headless_bool, + headless=headless, ) stop_text_encoder_training_pct = 0 - if not print_only_bool and check_if_model_exist( - output_name, output_dir, save_model_as, headless=headless_bool + if not print_only and check_if_model_exist( + output_name, output_dir, save_model_as, headless=headless ): return @@ -913,7 +911,7 @@ def train_model( # Determine the training configuration based on learning rate values # Sets flags for training specific components based on the provided learning rates. if float(learning_rate) == unet_lr_float == text_encoder_lr_float == 0: - output_message(msg="Please input learning rate values.", headless=headless_bool) + output_message(msg="Please input learning rate values.", headless=headless) return # Flag to train text encoder only if its learning rate is non-zero and unet's is zero. network_train_text_encoder_only = text_encoder_lr_float != 0 and unet_lr_float == 0 @@ -1041,7 +1039,7 @@ def train_model( output_dir, ) - if print_only_bool: + if print_only: log.warning( "Here is the trainer command as a reference. It will not be executed:\n" ) @@ -1084,7 +1082,7 @@ def lora_tab( ): dummy_db_true = gr.Checkbox(value=True, visible=False) dummy_db_false = gr.Checkbox(value=False, visible=False) - dummy_headless = gr.Label(value=headless, visible=False) + dummy_headless = gr.Checkbox(value=headless, visible=False) with gr.Tab("Training"), gr.Column(variant="compact") as tab: gr.Markdown( @@ -1111,6 +1109,21 @@ def lora_tab( with gr.Accordion("Folders", open=False), gr.Group(): folders = Folders(headless=headless, config=config) + with gr.Accordion("Dataset Preparation", open=False): + gr.Markdown( + "This section provide Dreambooth tools to help setup your dataset..." + ) + gradio_dreambooth_folder_creation_tab( + train_data_dir_input=source_model.train_data_dir, + reg_data_dir_input=folders.reg_data_dir, + output_dir_input=folders.output_dir, + logging_dir_input=folders.logging_dir, + headless=headless, + config=config, + ) + + gradio_dataset_balancing_tab(headless=headless) + with gr.Accordion("Parameters", open=False), gr.Column(): def list_presets(path): @@ -1900,20 +1913,6 @@ def update_LoRA_settings( ], ) - with gr.Accordion("Dataset Preparation", open=False): - gr.Markdown( - "This section provide Dreambooth tools to help setup your dataset..." - ) - gradio_dreambooth_folder_creation_tab( - train_data_dir_input=source_model.train_data_dir, - reg_data_dir_input=folders.reg_data_dir, - output_dir_input=folders.output_dir, - logging_dir_input=folders.logging_dir, - headless=headless, - config=config, - ) - gradio_dataset_balancing_tab(headless=headless) - with gr.Column(), gr.Group(): with gr.Row(): button_run = gr.Button("Start training", variant="primary") diff --git a/kohya_gui/tensorboard_gui.py b/kohya_gui/tensorboard_gui.py index edd1d42ae..53b7116f5 100644 --- a/kohya_gui/tensorboard_gui.py +++ b/kohya_gui/tensorboard_gui.py @@ -20,8 +20,6 @@ def start_tensorboard(headless, logging_dir, wait_time=5): os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0" global tensorboard_proc - headless_bool = True if headless.get("label") == "True" else False - # Read the TENSORBOARD_PORT from the environment, or use the default tensorboard_port = os.environ.get("TENSORBOARD_PORT", DEFAULT_TENSORBOARD_PORT) @@ -62,7 +60,7 @@ def start_tensorboard(headless, logging_dir, wait_time=5): log.error("Failed to start Tensorboard:", e) return - if not headless_bool: + if not headless: # Wait for some time to allow TensorBoard to start up time.sleep(wait_time) diff --git a/kohya_gui/textual_inversion_gui.py b/kohya_gui/textual_inversion_gui.py index 1505797a2..cbbec180c 100644 --- a/kohya_gui/textual_inversion_gui.py +++ b/kohya_gui/textual_inversion_gui.py @@ -434,17 +434,14 @@ def train_model( # Get list of function parameters and values parameters = list(locals().items()) - print_only_bool = True if print_only.get("label") == "True" else False log.info(f"Start training TI...") - headless_bool = True if headless.get("label") == "True" else False - if not validate_paths( output_dir=output_dir, pretrained_model_name_or_path=pretrained_model_name_or_path, train_data_dir=train_data_dir, reg_data_dir=reg_data_dir, - headless=headless_bool, + headless=headless, logging_dir=logging_dir, log_tracker_config=log_tracker_config, resume=resume, @@ -454,15 +451,15 @@ def train_model( return if token_string == "": - output_message(msg="Token string is missing", headless=headless_bool) + output_message(msg="Token string is missing", headless=headless) return if init_word == "": - output_message(msg="Init word is missing", headless=headless_bool) + output_message(msg="Init word is missing", headless=headless) return - if not print_only_bool and check_if_model_exist( - output_name, output_dir, save_model_as, headless_bool + if not print_only and check_if_model_exist( + output_name, output_dir, save_model_as, headless ): return @@ -670,7 +667,7 @@ def train_model( output_dir, ) - if print_only_bool: + if print_only: log.warning( "Here is the trainer command as a reference. It will not be executed:\n" ) @@ -708,7 +705,7 @@ def train_model( def ti_tab(headless=False, default_output_dir=None, config: dict = {}): dummy_db_true = gr.Checkbox(value=True, visible=False) dummy_db_false = gr.Checkbox(value=False, visible=False) - dummy_headless = gr.Label(value=headless, visible=False) + dummy_headless = gr.Checkbox(value=headless, visible=False) current_embedding_dir = ( default_output_dir @@ -739,6 +736,21 @@ def ti_tab(headless=False, default_output_dir=None, config: dict = {}): with gr.Accordion("Folders", open=False), gr.Group(): folders = Folders(headless=headless, config=config) + with gr.Accordion("Dataset Preparation", open=False): + gr.Markdown( + "This section provide Dreambooth tools to help setup your dataset..." + ) + gradio_dreambooth_folder_creation_tab( + train_data_dir_input=source_model.train_data_dir, + reg_data_dir_input=folders.reg_data_dir, + output_dir_input=folders.output_dir, + logging_dir_input=folders.logging_dir, + headless=headless, + config=config, + ) + + gradio_dataset_balancing_tab(headless=headless) + with gr.Accordion("Parameters", open=False), gr.Column(): with gr.Accordion("Basic", open="True"): with gr.Group(elem_id="basic_tab"): @@ -845,20 +857,6 @@ def list_embedding_files(path): with gr.Accordion("Samples", open=False, elem_id="samples_tab"): sample = SampleImages(config=config) - with gr.Accordion("Dataset Preparation", open=False): - gr.Markdown( - "This section provide Dreambooth tools to help setup your dataset..." - ) - gradio_dreambooth_folder_creation_tab( - train_data_dir_input=source_model.train_data_dir, - reg_data_dir_input=folders.reg_data_dir, - output_dir_input=folders.output_dir, - logging_dir_input=folders.logging_dir, - headless=headless, - config=config, - ) - gradio_dataset_balancing_tab(headless=headless) - with gr.Column(), gr.Group(): with gr.Row(): button_run = gr.Button("Start training", variant="primary") From 831af8babeb75faff62bcc6a8c6a4f80354f1ff1 Mon Sep 17 00:00:00 2001 From: bmaltais Date: Thu, 11 Apr 2024 08:00:31 -0400 Subject: [PATCH 08/10] Remove shell=True from scripts (#2257) * Remove shell=True from scripts --- README.md | 3 ++- _typos.toml | 1 + kohya_gui/basic_caption_gui.py | 2 +- kohya_gui/blip_caption_gui.py | 2 +- kohya_gui/class_command_executor.py | 2 +- kohya_gui/convert_lcm_gui.py | 2 +- kohya_gui/convert_model_gui.py | 2 +- kohya_gui/extract_lora_from_dylora_gui.py | 2 +- kohya_gui/extract_lora_gui.py | 2 +- kohya_gui/extract_lycoris_locon_gui.py | 2 +- kohya_gui/finetune_gui.py | 6 ++++-- kohya_gui/git_caption_gui.py | 2 +- kohya_gui/group_images_gui.py | 2 +- kohya_gui/merge_lora_gui.py | 2 +- kohya_gui/merge_lycoris_gui.py | 2 +- kohya_gui/resize_lora_gui.py | 2 +- kohya_gui/svd_merge_lora_gui.py | 2 +- kohya_gui/wd14_caption_gui.py | 2 +- 18 files changed, 22 insertions(+), 18 deletions(-) diff --git a/README.md b/README.md index c45a2ce28..55359a2d8 100644 --- a/README.md +++ b/README.md @@ -411,9 +411,10 @@ ControlNet dataset is used to specify the mask. The mask images should be the RG ### 2024/04/10 (v23.1.5) - Fix issue with Textual Inversion configuration file selection. -- Upgrade to gradio 4.19.2 to fix several high security risks associated to earlier versions. Hoping this will not introduce undorseen issues. +- Upgrade to gradio 4.19.2 to fix several high security risks associated to earlier versions. This is a major upgrade, moving from 3.x to 4.x. Hoping this will not introduce undorseen issues. - Upgrade transformers to 4.38.0 to fix a low severity security issue. - Add explicit --do_not_share parameter to kohya_gui.py to avoid sharing the GUI on platforms like Kaggle. +- Remove shell=True from subprocess calls to avoid security issues when using the GUI. ### 2024/04/08 (v23.1.4) diff --git a/_typos.toml b/_typos.toml index ae9e06b18..d73875a92 100644 --- a/_typos.toml +++ b/_typos.toml @@ -27,6 +27,7 @@ rik="rik" koo="koo" yos="yos" wn="wn" +parm = "parm" [files] diff --git a/kohya_gui/basic_caption_gui.py b/kohya_gui/basic_caption_gui.py index ed8d46d21..4866a3f59 100644 --- a/kohya_gui/basic_caption_gui.py +++ b/kohya_gui/basic_caption_gui.py @@ -83,7 +83,7 @@ def caption_images( ) # Run the command based on the operating system - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) # Check if overwrite option is enabled if overwrite: diff --git a/kohya_gui/blip_caption_gui.py b/kohya_gui/blip_caption_gui.py index dab98454c..cda0b2881 100644 --- a/kohya_gui/blip_caption_gui.py +++ b/kohya_gui/blip_caption_gui.py @@ -79,7 +79,7 @@ def caption_images( ) # Run the command in the sd-scripts folder context - subprocess.run(run_cmd, shell=True, env=env, cwd=f"{scriptdir}/sd-scripts") + subprocess.run(run_cmd, env=env, cwd=f"{scriptdir}/sd-scripts") # Add prefix and postfix add_pre_postfix( diff --git a/kohya_gui/class_command_executor.py b/kohya_gui/class_command_executor.py index ccd15a713..c9530cd62 100644 --- a/kohya_gui/class_command_executor.py +++ b/kohya_gui/class_command_executor.py @@ -28,7 +28,7 @@ def execute_command(self, run_cmd: str, **kwargs): if self.process and self.process.poll() is None: log.info("The command is already running. Please wait for it to finish.") else: - self.process = subprocess.Popen(run_cmd, shell=True, **kwargs) + self.process = subprocess.Popen(run_cmd, **kwargs) def kill_command(self): """ diff --git a/kohya_gui/convert_lcm_gui.py b/kohya_gui/convert_lcm_gui.py index dbc928e45..c087eece3 100644 --- a/kohya_gui/convert_lcm_gui.py +++ b/kohya_gui/convert_lcm_gui.py @@ -59,7 +59,7 @@ def convert_lcm(name, model_path, lora_scale, model_type): ) # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) # Return a success message log.info("Done extracting...") diff --git a/kohya_gui/convert_model_gui.py b/kohya_gui/convert_model_gui.py index f8fec7473..aedda3c55 100644 --- a/kohya_gui/convert_model_gui.py +++ b/kohya_gui/convert_model_gui.py @@ -104,7 +104,7 @@ def convert_model( ) # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) ### diff --git a/kohya_gui/extract_lora_from_dylora_gui.py b/kohya_gui/extract_lora_from_dylora_gui.py index d99a15235..9f55b6673 100644 --- a/kohya_gui/extract_lora_from_dylora_gui.py +++ b/kohya_gui/extract_lora_from_dylora_gui.py @@ -64,7 +64,7 @@ def extract_dylora( ) # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) log.info("Done extracting DyLoRA...") diff --git a/kohya_gui/extract_lora_gui.py b/kohya_gui/extract_lora_gui.py index 66c1e6123..4ef019043 100644 --- a/kohya_gui/extract_lora_gui.py +++ b/kohya_gui/extract_lora_gui.py @@ -102,7 +102,7 @@ def extract_lora( ) # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) ### diff --git a/kohya_gui/extract_lycoris_locon_gui.py b/kohya_gui/extract_lycoris_locon_gui.py index 4ae331579..1da7010b8 100644 --- a/kohya_gui/extract_lycoris_locon_gui.py +++ b/kohya_gui/extract_lycoris_locon_gui.py @@ -110,7 +110,7 @@ def extract_lycoris_locon( ) # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) log.info("Done extracting...") diff --git a/kohya_gui/finetune_gui.py b/kohya_gui/finetune_gui.py index cc4b28505..69cf30a5e 100644 --- a/kohya_gui/finetune_gui.py +++ b/kohya_gui/finetune_gui.py @@ -517,10 +517,11 @@ def train_model( env["PYTHONPATH"] = ( rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}" ) + env["TF_ENABLE_ONEDNN_OPTS"] = "0" if not print_only: # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) # create images buckets if generate_image_buckets: @@ -550,10 +551,11 @@ def train_model( env["PYTHONPATH"] = ( rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}" ) + env["TF_ENABLE_ONEDNN_OPTS"] = "0" if not print_only: # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) image_num = len( [ diff --git a/kohya_gui/git_caption_gui.py b/kohya_gui/git_caption_gui.py index a98449749..f0b07defe 100644 --- a/kohya_gui/git_caption_gui.py +++ b/kohya_gui/git_caption_gui.py @@ -51,7 +51,7 @@ def caption_images( ) # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) # Add prefix and postfix add_pre_postfix( diff --git a/kohya_gui/group_images_gui.py b/kohya_gui/group_images_gui.py index 55552e300..914aad255 100644 --- a/kohya_gui/group_images_gui.py +++ b/kohya_gui/group_images_gui.py @@ -53,7 +53,7 @@ def group_images( ) # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) log.info("...grouping done") diff --git a/kohya_gui/merge_lora_gui.py b/kohya_gui/merge_lora_gui.py index 662cd55f0..75f638d3f 100644 --- a/kohya_gui/merge_lora_gui.py +++ b/kohya_gui/merge_lora_gui.py @@ -452,6 +452,6 @@ def merge_lora( ) # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) log.info("Done merging...") diff --git a/kohya_gui/merge_lycoris_gui.py b/kohya_gui/merge_lycoris_gui.py index bb6ae9ee6..3ddbc7112 100644 --- a/kohya_gui/merge_lycoris_gui.py +++ b/kohya_gui/merge_lycoris_gui.py @@ -56,7 +56,7 @@ def merge_lycoris( ) # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) log.info("Done merging...") diff --git a/kohya_gui/resize_lora_gui.py b/kohya_gui/resize_lora_gui.py index 21b3ea533..9253c0a97 100644 --- a/kohya_gui/resize_lora_gui.py +++ b/kohya_gui/resize_lora_gui.py @@ -83,7 +83,7 @@ def resize_lora( ) # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) log.info("Done resizing...") diff --git a/kohya_gui/svd_merge_lora_gui.py b/kohya_gui/svd_merge_lora_gui.py index c14c9c6ad..ebb4f0852 100644 --- a/kohya_gui/svd_merge_lora_gui.py +++ b/kohya_gui/svd_merge_lora_gui.py @@ -99,7 +99,7 @@ def svd_merge_lora( ) # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) ### diff --git a/kohya_gui/wd14_caption_gui.py b/kohya_gui/wd14_caption_gui.py index 8f97f02f8..9768d3153 100644 --- a/kohya_gui/wd14_caption_gui.py +++ b/kohya_gui/wd14_caption_gui.py @@ -94,7 +94,7 @@ def caption_images( env["TF_ENABLE_ONEDNN_OPTS"] = "0" # Run the command - subprocess.run(run_cmd, shell=True, env=env) + subprocess.run(run_cmd, env=env) # Add prefix and postfix add_pre_postfix( From 8bc67a7467f8366db1a4b9b3b14525ec763f1650 Mon Sep 17 00:00:00 2001 From: bmaltais Date: Thu, 11 Apr 2024 08:32:17 -0400 Subject: [PATCH 09/10] Find replace (#2258) * Improve security to limit finding and replacing contents of any files on the system, that the attacker knows the contents of by specifying the list of acceptable extensions for captions. * Update caption extension fields to Dropdown with restricted choices for security reasons --- kohya_gui/basic_caption_gui.py | 4 ++-- kohya_gui/blip_caption_gui.py | 4 ++-- kohya_gui/class_basic_training.py | 9 ++++---- kohya_gui/common_gui.py | 35 ++++++++++++++++++++++++------- kohya_gui/git_caption_gui.py | 4 ++-- kohya_gui/group_images_gui.py | 6 +++--- kohya_gui/manual_caption_gui.py | 4 ++-- kohya_gui/wd14_caption_gui.py | 6 +++--- 8 files changed, 46 insertions(+), 26 deletions(-) diff --git a/kohya_gui/basic_caption_gui.py b/kohya_gui/basic_caption_gui.py index 4866a3f59..c47442ab5 100644 --- a/kohya_gui/basic_caption_gui.py +++ b/kohya_gui/basic_caption_gui.py @@ -190,9 +190,9 @@ def list_images_dirs(path): show_progress=False, ) # Textbox for caption file extension - caption_ext = gr.Textbox( + caption_ext = gr.Dropdown( label="Caption file extension", - placeholder="Extension for caption file (e.g., .caption, .txt)", + choices=[".cap", ".caption", ".txt"], value=".txt", interactive=True, ) diff --git a/kohya_gui/blip_caption_gui.py b/kohya_gui/blip_caption_gui.py index cda0b2881..fc4ca1669 100644 --- a/kohya_gui/blip_caption_gui.py +++ b/kohya_gui/blip_caption_gui.py @@ -142,9 +142,9 @@ def list_train_dirs(path): show_progress=False, ) with gr.Row(): - caption_file_ext = gr.Textbox( + caption_file_ext = gr.Dropdown( label="Caption file extension", - placeholder="Extension for caption file (e.g., .caption, .txt)", + choices=[".cap", ".caption", ".txt"], value=".txt", interactive=True, ) diff --git a/kohya_gui/class_basic_training.py b/kohya_gui/class_basic_training.py index 8b1a040f6..57c8e7356 100644 --- a/kohya_gui/class_basic_training.py +++ b/kohya_gui/class_basic_training.py @@ -98,10 +98,11 @@ def init_training_controls(self) -> None: label="Save every N epochs", value=self.config.get("basic.save_every_n_epochs", 1), precision=0 ) # Initialize the caption extension input - self.caption_extension = gr.Textbox( - label="Caption Extension", - placeholder="(Optional) default: .caption", - value=self.config.get("basic.caption_extension", ""), + self.caption_extension = gr.Dropdown( + label="Caption file extension", + choices=[".cap", ".caption", ".txt"], + value=".txt", + interactive=True, ) def init_precision_and_resources_controls(self) -> None: diff --git a/kohya_gui/common_gui.py b/kohya_gui/common_gui.py index fc5e4944a..97addc159 100644 --- a/kohya_gui/common_gui.py +++ b/kohya_gui/common_gui.py @@ -715,7 +715,6 @@ def has_ext_files(folder_path: str, file_extension: str) -> bool: # If no file with the specified extension is found, return False return False - def find_replace( folder_path: str = "", caption_file_ext: str = ".caption", @@ -747,21 +746,41 @@ def find_replace( ) # Exit the function early return + + # Check if the caption file extension is one of the supported extensions + if caption_file_ext not in [".caption", ".txt", ".txt2", ".cap"]: + log.error( + f"Unsupported file extension {caption_file_ext} for caption files. Please use .caption, .txt, .txt2, or .cap." + ) + # Exit the function early + return + + # Check if the folder path exists + if not os.path.exists(folder_path): + log.error(f"The provided path '{folder_path}' is not a valid folder.") + return # List all caption files in the folder - caption_files = [f for f in os.listdir(folder_path) if f.endswith(caption_file_ext)] + try: + caption_files = [f for f in os.listdir(folder_path) if f.endswith(caption_file_ext)] + except Exception as e: + log.error(f"Error accessing folder {folder_path}: {e}") + return # Iterate over the list of caption files for caption_file in caption_files: # Construct the full path for each caption file file_path = os.path.join(folder_path, caption_file) # Read and replace text - with open(file_path, "r", errors="ignore") as f: - content = f.read().replace(search_text, replace_text) - - # Write the updated content back to the file - with open(file_path, "w") as f: - f.write(content) + try: + with open(file_path, "r", errors="ignore") as f: + content = f.read().replace(search_text, replace_text) + + # Write the updated content back to the file + with open(file_path, "w") as f: + f.write(content) + except Exception as e: + log.error(f"Error processing file {file_path}: {e}") def color_aug_changed(color_aug): diff --git a/kohya_gui/git_caption_gui.py b/kohya_gui/git_caption_gui.py index f0b07defe..e2538b259 100644 --- a/kohya_gui/git_caption_gui.py +++ b/kohya_gui/git_caption_gui.py @@ -114,9 +114,9 @@ def list_train_dirs(path): show_progress=False, ) with gr.Row(): - caption_ext = gr.Textbox( + caption_ext = gr.Dropdown( label="Caption file extension", - placeholder="Extension for caption file (e.g., .caption, .txt)", + choices=[".cap", ".caption", ".txt"], value=".txt", interactive=True, ) diff --git a/kohya_gui/group_images_gui.py b/kohya_gui/group_images_gui.py index 914aad255..5219fb83f 100644 --- a/kohya_gui/group_images_gui.py +++ b/kohya_gui/group_images_gui.py @@ -171,9 +171,9 @@ def list_output_dirs(path): info="Generate caption files for the grouped images based on their folder name", ) - caption_ext = gr.Textbox( - label="Caption Extension", - placeholder="Caption file extension (e.g., .txt)", + caption_ext = gr.Dropdown( + label="Caption file extension", + choices=[".cap", ".caption", ".txt"], value=".txt", interactive=True, ) diff --git a/kohya_gui/manual_caption_gui.py b/kohya_gui/manual_caption_gui.py index bc90d2158..cca885a92 100644 --- a/kohya_gui/manual_caption_gui.py +++ b/kohya_gui/manual_caption_gui.py @@ -298,9 +298,9 @@ def list_images_dirs(path): show_progress=False, ) load_images_button = gr.Button("Load", elem_id="open_folder") - caption_ext = gr.Textbox( + caption_ext = gr.Dropdown( label="Caption file extension", - placeholder="Extension for caption file (e.g., .caption, .txt)", + choices=[".cap", ".caption", ".txt"], value=".txt", interactive=True, ) diff --git a/kohya_gui/wd14_caption_gui.py b/kohya_gui/wd14_caption_gui.py index 9768d3153..c5b8269db 100644 --- a/kohya_gui/wd14_caption_gui.py +++ b/kohya_gui/wd14_caption_gui.py @@ -188,10 +188,10 @@ def list_train_dirs(path): with gr.Row(): - caption_extension = gr.Textbox( + caption_extension = gr.Dropdown( label="Caption file extension", - placeholder="Extension for caption file (e.g., .caption, .txt)", - value=config.get("wd14_caption.caption_extension", ".txt"), + choices=[".cap", ".caption", ".txt"], + value=".txt", interactive=True, ) From 25bb1303fff21cb5bae17236d53504e85c1866df Mon Sep 17 00:00:00 2001 From: bmaltais Date: Thu, 11 Apr 2024 08:55:03 -0400 Subject: [PATCH 10/10] Restrict extensions pre_postfix can use to improve security (#2259) --- README.md | 63 ++++------------------------------------- kohya_gui/common_gui.py | 4 +++ 2 files changed, 9 insertions(+), 58 deletions(-) diff --git a/README.md b/README.md index 55359a2d8..e7bc3dd18 100644 --- a/README.md +++ b/README.md @@ -43,19 +43,12 @@ The GUI allows you to set the training parameters and generate and run the requi - [Masked loss](#masked-loss) - [Change History](#change-history) - [2024/04/10 (v23.1.5)](#20240410-v2315) + - [Security Improvements](#security-improvements) - [2024/04/08 (v23.1.4)](#20240408-v2314) - [2024/04/08 (v23.1.3)](#20240408-v2313) - [2024/04/08 (v23.1.2)](#20240408-v2312) - [2024/04/07 (v23.1.1)](#20240407-v2311) - [2024/04/07 (v23.1.0)](#20240407-v2310) - - [2024/03/21 (v23.0.15)](#20240321-v23015) - - [2024/03/19 (v23.0.14)](#20240319-v23014) - - [2024/03/19 (v23.0.13)](#20240319-v23013) - - [2024/03/16 (v23.0.12)](#20240316-v23012) - - [New Features \& Improvements](#new-features--improvements) - - [Software Updates](#software-updates) - - [Recommendations for Users](#recommendations-for-users) - - [2024/03/13 (v23.0.11)](#20240313-v23011) ## 🦒 Colab @@ -413,8 +406,12 @@ ControlNet dataset is used to specify the mask. The mask images should be the RG - Fix issue with Textual Inversion configuration file selection. - Upgrade to gradio 4.19.2 to fix several high security risks associated to earlier versions. This is a major upgrade, moving from 3.x to 4.x. Hoping this will not introduce undorseen issues. - Upgrade transformers to 4.38.0 to fix a low severity security issue. + +#### Security Improvements + - Add explicit --do_not_share parameter to kohya_gui.py to avoid sharing the GUI on platforms like Kaggle. - Remove shell=True from subprocess calls to avoid security issues when using the GUI. +- Limit caption extensions to a fixed set of extensions to limit the risk of finding and replacing text content in unexpected files. ### 2024/04/08 (v23.1.4) @@ -519,53 +516,3 @@ ControlNet dataset is used to specify the mask. The mask images should be the RG - Added support for "Dataset Preparation" defaults via the config.toml file. - Added a field to allow for the input of extra accelerate launch arguments. - Added new caption tool from https://github.com/kainatquaderee - -### 2024/03/21 (v23.0.15) - -- Add support for toml dataset configuration fole to all trainers -- Add new setup menu option to install Triton 2.1.0 for Windows -- Add support for LyCORIS BOFT and DoRA and QLyCORIS options for LoHA, LoKr and LoCon -- Fix issue with vae path validation -- Other fixes - -### 2024/03/19 (v23.0.14) - -- Fix blip caption issue - -### 2024/03/19 (v23.0.13) - -- Fix issue with image samples. - -### 2024/03/16 (v23.0.12) - -#### New Features & Improvements - -- **Enhanced Logging and Tracking Capabilities** - - Added support for configuring advanced logging and tracking: - - `wandb_run_name`: Set a custom name for your Weights & Biases runs to easily identify and organize your experiments. - - `log_tracker_name` and `log_tracker_config`: Integrate custom logging trackers with your projects. Specify the tracker name and provide its configuration to enable detailed monitoring and logging of your runs. - -- **Custom Path Defaults** - - You can now specify custom paths more easily: - - Simply copy the `config example.toml` file located in the root directory of the repository to `config.toml`. - - Edit the `config.toml` file to adjust paths and settings according to your preferences. - -#### Software Updates - -- **sd-scripts updated to v0.8.5** - - **Bug Fixes:** - - Corrected an issue where the value of timestep embedding was incorrect during SDXL training. This fix ensures accurate training progress and results. - - Addressed a related inference issue with the generation script, improving the reliability of SDXL model outputs. - - **Note:** The exact impact of this bug is currently unknown, but it's recommended to update to v0.8.5 for anyone engaged in SDXL training to ensure optimal performance and results. - -- **Upgrade of `lycoris_lora` Python Module** - - Updated the `lycoris_lora` module to version 2.2.0.post3. This update may include bug fixes, performance improvements, and new features. - -#### Recommendations for Users - -- To benefit from the latest features and improvements, users are encouraged to update their installations and configurations accordingly. - -### 2024/03/13 (v23.0.11) - -- Increase icon size. -- More setup fixes. diff --git a/kohya_gui/common_gui.py b/kohya_gui/common_gui.py index 97addc159..1b3e0fcee 100644 --- a/kohya_gui/common_gui.py +++ b/kohya_gui/common_gui.py @@ -646,6 +646,10 @@ def add_pre_postfix( postfix (str, optional): Postfix to add to the content of the caption files. caption_file_ext (str, optional): Extension of the caption files. """ + # Enforce that the provided extension is one of .caption, .cap, .txt + if caption_file_ext not in (".caption", ".cap", ".txt"): + log.error("Invalid caption file extension. Must be on of .caption, .cap, .txt") + return # If neither prefix nor postfix is provided, return early if prefix == "" and postfix == "":