StableDiffusion/KohyaSS
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #34 from bmaltais/dev

Browse Source 
Add new features
This commit is contained in:
bmaltais 2023-01-08 19:32:21 -05:00 committed by GitHub
commit c5c298201c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 99 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
README.md
View File

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

28
finetune_gui.py
View File

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

51
library/basic_caption_gui.py
View File

@ -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
],
)

38
library/common_gui.py
View File

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