Merge pull request #34 from bmaltais/dev

Add new features

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):

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)

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
+        if not prefix=='' or not postfix=='':
-        add_pre_postfix(
+            # Add prefix and postfix
-            folder=images_dir_input,
+            add_pre_postfix(
-            caption_file_ext=caption_file_ext,
+                folder=images_dir_input,
-            prefix=prefix,
+                caption_file_ext=caption_file_ext,
-            postfix=postfix,
+                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(
+            find = gr.Textbox(
-                label='Overwrite existing captions in folder',
+                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(
+            replace = gr.Textbox(
-                label='Caption file extension',
+                label='Replacement text',
-                placeholder='(Optional) Default: .caption',
+                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
             ],
         )

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: