diff --git a/README.md b/README.md index d5df129..5ed7ea9 100644 --- a/README.md +++ b/README.md @@ -101,6 +101,9 @@ Once you have created the LoRA network you can generate images via auto1111 by i ## Change history +* 2023/01/08 (v19.4.2): + - Add find/replace option to Basic Caption utility + - Add resume training and save_state option to finetune UI * 2023/01/06 (v19.4.1): - Emergency fix for new version of gradio causing issues with drop down menus. Please run `pip install -U -r requirements.txt` to fix the issue after pulling this repo. * 2023/01/06 (v19.4): diff --git a/finetune_gui.py b/finetune_gui.py index 9c8e8a5..76b7298 100644 --- a/finetune_gui.py +++ b/finetune_gui.py @@ -57,6 +57,8 @@ def save_configuration( use_8bit_adam, xformers, clip_skip, + save_state, + resume, ): original_file_path = file_path @@ -111,6 +113,8 @@ def save_configuration( 'use_8bit_adam': use_8bit_adam, 'xformers': xformers, 'clip_skip': clip_skip, + 'save_state': save_state, + 'resume': resume, } # Save the data to the selected file @@ -156,12 +160,14 @@ def open_config_file( use_8bit_adam, xformers, clip_skip, + save_state, + resume, ): original_file_path = file_path file_path = get_file_path(file_path) if file_path != '' and file_path != None: - print(file_path) + print(f'Loading config file {file_path}') # load variables from JSON file with open(file_path, 'r') as f: my_data = json.load(f) @@ -210,6 +216,8 @@ def open_config_file( my_data.get('use_8bit_adam', use_8bit_adam), my_data.get('xformers', xformers), my_data.get('clip_skip', clip_skip), + my_data.get('save_state', save_state), + my_data.get('resume', resume), ) @@ -248,6 +256,8 @@ def train_model( use_8bit_adam, xformers, clip_skip, + save_state, + resume, ): def save_inference_file(output_dir, v2, v_parameterization): # Copy inference model for v2 if required @@ -365,6 +375,10 @@ def train_model( run_cmd += f' --save_model_as={save_model_as}' if int(clip_skip) > 1: run_cmd += f' --clip_skip={str(clip_skip)}' + if save_state: + run_cmd += ' --save_state' + if not resume == '': + run_cmd += f' --resume={resume}' print(run_cmd) # Run the command @@ -698,6 +712,16 @@ def finetune_tab(): clip_skip = gr.Slider( label='Clip skip', value='1', minimum=1, maximum=12, step=1 ) + with gr.Row(): + save_state = gr.Checkbox( + label='Save training state', value=False + ) + resume = gr.Textbox( + label='Resume from saved training state', + placeholder='path to "last-state" state folder to resume from', + ) + resume_button = gr.Button('📂', elem_id='open_folder_small') + resume_button.click(get_folder_path, outputs=resume) with gr.Box(): with gr.Row(): create_caption = gr.Checkbox( @@ -744,6 +768,8 @@ def finetune_tab(): use_8bit_adam, xformers, clip_skip, + save_state, + resume, ] button_run.click(train_model, inputs=settings_list) diff --git a/library/basic_caption_gui.py b/library/basic_caption_gui.py index 3c1d719..550d757 100644 --- a/library/basic_caption_gui.py +++ b/library/basic_caption_gui.py @@ -1,7 +1,7 @@ import gradio as gr from easygui import msgbox import subprocess -from .common_gui import get_folder_path, add_pre_postfix +from .common_gui import get_folder_path, add_pre_postfix, find_replace def caption_images( @@ -10,7 +10,7 @@ def caption_images( overwrite_input, caption_file_ext, prefix, - postfix, + postfix, find, replace ): # Check for images_dir_input if images_dir_input == '': @@ -35,13 +35,21 @@ def caption_images( subprocess.run(run_cmd) if overwrite_input: - # Add prefix and postfix - add_pre_postfix( - folder=images_dir_input, - caption_file_ext=caption_file_ext, - prefix=prefix, - postfix=postfix, - ) + if not prefix=='' or not postfix=='': + # Add prefix and postfix + add_pre_postfix( + folder=images_dir_input, + caption_file_ext=caption_file_ext, + prefix=prefix, + postfix=postfix, + ) + if not find=='': + find_replace( + folder=images_dir_input, + caption_file_ext=caption_file_ext, + find=find, + replace=replace, + ) else: if not prefix == '' or not postfix == '': msgbox( @@ -73,6 +81,16 @@ def gradio_basic_caption_gui_tab(): button_images_dir_input.click( get_folder_path, outputs=images_dir_input ) + caption_file_ext = gr.Textbox( + label='Caption file extension', + placeholder='(Optional) Default: .caption', + interactive=True, + ) + overwrite_input = gr.Checkbox( + label='Overwrite existing captions in folder', + interactive=True, + value=False, + ) with gr.Row(): prefix = gr.Textbox( label='Prefix to add to caption', @@ -81,7 +99,7 @@ def gradio_basic_caption_gui_tab(): ) caption_text_input = gr.Textbox( label='Caption text', - placeholder='Eg: , by some artist. Leave empti if you just want to add pre or postfix', + placeholder='Eg: , by some artist. Leave empty if you just want to add pre or postfix', interactive=True, ) postfix = gr.Textbox( @@ -90,14 +108,14 @@ def gradio_basic_caption_gui_tab(): interactive=True, ) with gr.Row(): - overwrite_input = gr.Checkbox( - label='Overwrite existing captions in folder', + find = gr.Textbox( + label='Find text', + placeholder='Eg: , by some artist. Leave empty if you just want to add pre or postfix', interactive=True, - value=False, ) - caption_file_ext = gr.Textbox( - label='Caption file extension', - placeholder='(Optional) Default: .caption', + replace = gr.Textbox( + label='Replacement text', + placeholder='Eg: , by some artist. Leave empty if you just want to replace with nothing', interactive=True, ) caption_button = gr.Button('Caption images') @@ -111,5 +129,6 @@ def gradio_basic_caption_gui_tab(): caption_file_ext, prefix, postfix, + find, replace ], ) diff --git a/library/common_gui.py b/library/common_gui.py index f54267b..ff7b2c5 100644 --- a/library/common_gui.py +++ b/library/common_gui.py @@ -7,6 +7,15 @@ def get_dir_and_file(file_path): dir_path, file_name = os.path.split(file_path) return (dir_path, file_name) +def has_ext_files(directory, extension): + # Iterate through all the files in the directory + for file in os.listdir(directory): + # If the file name ends with extension, return True + if file.endswith(extension): + return True + # If no extension files were found, return False + return False + def get_file_path(file_path='', defaultextension='.json', extension_name='Config files'): current_file_path = file_path # print(f'current file path: {current_file_path}') @@ -127,13 +136,13 @@ def get_saveasfilename_path(file_path='', extensions='*', extension_name='Config def add_pre_postfix( folder='', prefix='', postfix='', caption_file_ext='.caption' ): + if not has_ext_files(folder, caption_file_ext): + msgbox(f'No files with extension {caption_file_ext} were found in {folder}...') + return + if prefix == '' and postfix == '': return - # set caption extention to default in case it was not provided - if caption_file_ext == '': - caption_file_ext = '.caption' - files = [f for f in os.listdir(folder) if f.endswith(caption_file_ext)] if not prefix == '': prefix = f'{prefix} ' @@ -147,6 +156,27 @@ def add_pre_postfix( f.seek(0, 0) f.write(f'{prefix}{content}{postfix}') f.close() + +def find_replace( + folder='', caption_file_ext='.caption', find='', replace='' +): + print('Running caption find/replace') + if not has_ext_files(folder, caption_file_ext): + msgbox(f'No files with extension {caption_file_ext} were found in {folder}...') + return + + if find == '': + return + + files = [f for f in os.listdir(folder) if f.endswith(caption_file_ext)] + for file in files: + with open(os.path.join(folder, file), 'r') as f: + content = f.read() + f.close + content = content.replace(find, replace) + with open(os.path.join(folder, file), 'w') as f: + f.write(content) + f.close() def color_aug_changed(color_aug): if color_aug: