|
import gradio as gr |
|
from .common_gui import get_folder_path, get_any_file_path |
|
|
|
class AdvancedTraining: |
|
def __init__( |
|
self, |
|
headless=False, |
|
finetuning: bool = False |
|
): |
|
self.headless = headless |
|
self.finetuning = finetuning |
|
def noise_offset_type_change(noise_offset_type): |
|
if noise_offset_type == 'Original': |
|
return (gr.Group.update(visible=True), gr.Group.update(visible=False)) |
|
else: |
|
return (gr.Group.update(visible=False), gr.Group.update(visible=True)) |
|
|
|
with gr.Row(visible=not finetuning): |
|
self.no_token_padding = gr.Checkbox( |
|
label='No token padding', value=False |
|
) |
|
self.gradient_accumulation_steps = gr.Number( |
|
label='Gradient accumulate steps', value='1' |
|
) |
|
self.weighted_captions = gr.Checkbox( |
|
label='Weighted captions', value=False |
|
) |
|
with gr.Row(visible=not finetuning): |
|
self.prior_loss_weight = gr.Number( |
|
label='Prior loss weight', value=1.0 |
|
) |
|
self.vae = gr.Textbox( |
|
label='VAE', |
|
placeholder='(Optiona) path to checkpoint of vae to replace for training', |
|
) |
|
self.vae_button = gr.Button( |
|
'📂', elem_id='open_folder_small', visible=(not headless) |
|
) |
|
self.vae_button.click( |
|
get_any_file_path, |
|
outputs=self.vae, |
|
show_progress=False, |
|
) |
|
with gr.Row(visible=not finetuning): |
|
self.lr_scheduler_num_cycles = gr.Textbox( |
|
label='LR number of cycles', |
|
placeholder='(Optional) For Cosine with restart and polynomial only', |
|
) |
|
|
|
self.lr_scheduler_power = gr.Textbox( |
|
label='LR power', |
|
placeholder='(Optional) For Cosine with restart and polynomial only', |
|
) |
|
|
|
with gr.Row(): |
|
self.additional_parameters = gr.Textbox( |
|
label='Additional parameters', |
|
placeholder='(Optional) Use to provide additional parameters not handled by the GUI. Eg: --some_parameters "value"', |
|
) |
|
with gr.Row(): |
|
self.save_every_n_steps = gr.Number( |
|
label='Save every N steps', |
|
value=0, |
|
precision=0, |
|
info='(Optional) The model is saved every specified steps', |
|
) |
|
self.save_last_n_steps = gr.Number( |
|
label='Save last N steps', |
|
value=0, |
|
precision=0, |
|
info='(Optional) Save only the specified number of models (old models will be deleted)', |
|
) |
|
self.save_last_n_steps_state = gr.Number( |
|
label='Save last N states', |
|
value=0, |
|
precision=0, |
|
info='(Optional) Save only the specified number of states (old models will be deleted)', |
|
) |
|
with gr.Row(): |
|
self.keep_tokens = gr.Slider( |
|
label='Keep n tokens', value='0', minimum=0, maximum=32, step=1 |
|
) |
|
self.clip_skip = gr.Slider( |
|
label='Clip skip', value='1', minimum=1, maximum=12, step=1 |
|
) |
|
self.max_token_length = gr.Dropdown( |
|
label='Max Token Length', |
|
choices=[ |
|
'75', |
|
'150', |
|
'225', |
|
], |
|
value='75', |
|
) |
|
self.full_fp16 = gr.Checkbox( |
|
label='Full fp16 training (experimental)', value=False |
|
) |
|
with gr.Row(): |
|
self.gradient_checkpointing = gr.Checkbox( |
|
label='Gradient checkpointing', value=False |
|
) |
|
self.shuffle_caption = gr.Checkbox(label='Shuffle caption', value=False) |
|
self.persistent_data_loader_workers = gr.Checkbox( |
|
label='Persistent data loader', value=False |
|
) |
|
self.mem_eff_attn = gr.Checkbox( |
|
label='Memory efficient attention', value=False |
|
) |
|
with gr.Row(): |
|
|
|
|
|
|
|
|
|
self.xformers = gr.Checkbox(label='Use xformers', value=True) |
|
self.color_aug = gr.Checkbox(label='Color augmentation', value=False) |
|
self.flip_aug = gr.Checkbox(label='Flip augmentation', value=False) |
|
self.min_snr_gamma = gr.Slider( |
|
label='Min SNR gamma', value=0, minimum=0, maximum=20, step=1 |
|
) |
|
with gr.Row(): |
|
self.bucket_no_upscale = gr.Checkbox( |
|
label="Don't upscale bucket resolution", value=True |
|
) |
|
self.bucket_reso_steps = gr.Slider( |
|
label='Bucket resolution steps', value=64, minimum=1, maximum=128 |
|
) |
|
self.random_crop = gr.Checkbox( |
|
label='Random crop instead of center crop', value=False |
|
) |
|
|
|
with gr.Row(): |
|
self.min_timestep = gr.Slider( |
|
label='Min Timestep', |
|
value=0, |
|
step=1, |
|
minimum=0, |
|
maximum=1000, |
|
info='Values greater than 0 will make the model more img2img focussed. 0 = image only' |
|
) |
|
self.max_timestep = gr.Slider( |
|
label='Max Timestep', |
|
value=1000, |
|
step=1, |
|
minimum=0, |
|
maximum=1000, |
|
info='Values lower than 1000 will make the model more img2img focussed. 1000 = noise only', |
|
) |
|
|
|
with gr.Row(): |
|
self.noise_offset_type = gr.Dropdown( |
|
label='Noise offset type', |
|
choices=[ |
|
'Original', |
|
'Multires', |
|
], |
|
value='Original', |
|
) |
|
with gr.Row(visible=True) as self.noise_offset_original: |
|
self.noise_offset = gr.Slider( |
|
label='Noise offset', |
|
value=0, |
|
minimum=0, |
|
maximum=1, |
|
step=0.01, |
|
info='recommended values are 0.05 - 0.15', |
|
) |
|
self.adaptive_noise_scale = gr.Slider( |
|
label='Adaptive noise scale', |
|
value=0, |
|
minimum=-1, |
|
maximum=1, |
|
step=0.001, |
|
info='(Experimental, Optional) Since the latent is close to a normal distribution, it may be a good idea to specify a value around 1/10 the noise offset.', |
|
) |
|
with gr.Row(visible=False) as self.noise_offset_multires: |
|
self.multires_noise_iterations = gr.Slider( |
|
label='Multires noise iterations', |
|
value=0, |
|
minimum=0, |
|
maximum=64, |
|
step=1, |
|
info='enable multires noise (recommended values are 6-10)', |
|
) |
|
self.multires_noise_discount = gr.Slider( |
|
label='Multires noise discount', |
|
value=0, |
|
minimum=0, |
|
maximum=1, |
|
step=0.01, |
|
info='recommended values are 0.8. For LoRAs with small datasets, 0.1-0.3', |
|
) |
|
self.noise_offset_type.change( |
|
noise_offset_type_change, |
|
inputs=[self.noise_offset_type], |
|
outputs=[self.noise_offset_original, self.noise_offset_multires] |
|
) |
|
with gr.Row(): |
|
self.caption_dropout_every_n_epochs = gr.Number( |
|
label='Dropout caption every n epochs', value=0 |
|
) |
|
self.caption_dropout_rate = gr.Slider( |
|
label='Rate of caption dropout', value=0, minimum=0, maximum=1 |
|
) |
|
self.vae_batch_size = gr.Slider( |
|
label='VAE batch size', minimum=0, maximum=32, value=0, step=1 |
|
) |
|
with gr.Row(): |
|
self.save_state = gr.Checkbox(label='Save training state', value=False) |
|
self.resume = gr.Textbox( |
|
label='Resume from saved training state', |
|
placeholder='path to "last-state" state folder to resume from', |
|
) |
|
self.resume_button = gr.Button( |
|
'📂', elem_id='open_folder_small', visible=(not headless) |
|
) |
|
self.resume_button.click( |
|
get_folder_path, |
|
outputs=self.resume, |
|
show_progress=False, |
|
) |
|
self.max_train_epochs = gr.Textbox( |
|
label='Max train epoch', |
|
placeholder='(Optional) Override number of epoch', |
|
) |
|
self.max_data_loader_n_workers = gr.Textbox( |
|
label='Max num workers for DataLoader', |
|
placeholder='(Optional) Override number of epoch. Default: 8', |
|
value='0', |
|
) |
|
with gr.Row(): |
|
self.wandb_api_key = gr.Textbox( |
|
label='WANDB API Key', |
|
value='', |
|
placeholder='(Optional)', |
|
info='Users can obtain and/or generate an api key in the their user settings on the website: https://wandb.ai/login', |
|
) |
|
self.use_wandb = gr.Checkbox( |
|
label='WANDB Logging', |
|
value=False, |
|
info='If unchecked, tensorboard will be used as the default for logging.', |
|
) |
|
self.scale_v_pred_loss_like_noise_pred = gr.Checkbox( |
|
label='Scale v prediction loss', |
|
value=False, |
|
info='Only for SD v2 models. By scaling the loss according to the time step, the weights of global noise prediction and local noise prediction become the same, and the improvement of details may be expected.', |
|
) |
|
|
