| import gradio as gr | |
| import os | |
| class BasicTraining: | |
| def __init__( | |
| self, | |
| learning_rate_value='1e-6', | |
| lr_scheduler_value='constant', | |
| lr_warmup_value='0', | |
| finetuning: bool = False, | |
| ): | |
| self.learning_rate_value = learning_rate_value | |
| self.lr_scheduler_value = lr_scheduler_value | |
| self.lr_warmup_value = lr_warmup_value | |
| self.finetuning = finetuning | |
| with gr.Row(): | |
| self.train_batch_size = gr.Slider( | |
| minimum=1, | |
| maximum=64, | |
| label='Train batch size', | |
| value=1, | |
| step=1, | |
| ) | |
| self.epoch = gr.Number(label='Epoch', value=1, precision=0) | |
| self.save_every_n_epochs = gr.Number( | |
| label='Save every N epochs', value=1, precision=0 | |
| ) | |
| self.caption_extension = gr.Textbox( | |
| label='Caption Extension', | |
| placeholder='(Optional) Extension for caption files. default: .caption', | |
| ) | |
| with gr.Row(): | |
| self.mixed_precision = gr.Dropdown( | |
| label='Mixed precision', | |
| choices=[ | |
| 'no', | |
| 'fp16', | |
| 'bf16', | |
| ], | |
| value='fp16', | |
| ) | |
| self.save_precision = gr.Dropdown( | |
| label='Save precision', | |
| choices=[ | |
| 'float', | |
| 'fp16', | |
| 'bf16', | |
| ], | |
| value='fp16', | |
| ) | |
| self.num_cpu_threads_per_process = gr.Slider( | |
| minimum=1, | |
| maximum=os.cpu_count(), | |
| step=1, | |
| label='Number of CPU threads per core', | |
| value=2, | |
| ) | |
| self.seed = gr.Textbox( | |
| label='Seed', placeholder='(Optional) eg:1234' | |
| ) | |
| self.cache_latents = gr.Checkbox(label='Cache latents', value=True) | |
| self.cache_latents_to_disk = gr.Checkbox( | |
| label='Cache latents to disk', value=False | |
| ) | |
| with gr.Row(): | |
| self.learning_rate = gr.Number( | |
| label='Learning rate', value=learning_rate_value | |
| ) | |
| self.lr_scheduler = gr.Dropdown( | |
| label='LR Scheduler', | |
| choices=[ | |
| 'adafactor', | |
| 'constant', | |
| 'constant_with_warmup', | |
| 'cosine', | |
| 'cosine_with_restarts', | |
| 'linear', | |
| 'polynomial', | |
| ], | |
| value=lr_scheduler_value, | |
| ) | |
| self.lr_warmup = gr.Slider( | |
| label='LR warmup (% of steps)', | |
| value=lr_warmup_value, | |
| minimum=0, | |
| maximum=100, | |
| step=1, | |
| ) | |
| self.optimizer = gr.Dropdown( | |
| label='Optimizer', | |
| choices=[ | |
| 'AdamW', | |
| 'AdamW8bit', | |
| 'Adafactor', | |
| 'DAdaptation', | |
| 'DAdaptAdaGrad', | |
| 'DAdaptAdam', | |
| 'DAdaptAdan', | |
| 'DAdaptAdanIP', | |
| 'DAdaptAdamPreprint', | |
| 'DAdaptLion', | |
| 'DAdaptSGD', | |
| 'Lion', | |
| 'Lion8bit', | |
| 'Prodigy', | |
| 'SGDNesterov', | |
| 'SGDNesterov8bit', | |
| ], | |
| value='AdamW8bit', | |
| interactive=True, | |
| ) | |
| with gr.Row(): | |
| self.optimizer_args = gr.Textbox( | |
| label='Optimizer extra arguments', | |
| placeholder='(Optional) eg: relative_step=True scale_parameter=True warmup_init=True', | |
| ) | |
| with gr.Row(visible=not finetuning): | |
| self.max_resolution = gr.Textbox( | |
| label='Max resolution', | |
| value='512,512', | |
| placeholder='512,512', | |
| ) | |
| self.stop_text_encoder_training = gr.Slider( | |
| minimum=-1, | |
| maximum=100, | |
| value=0, | |
| step=1, | |
| label='Stop text encoder training', | |
| ) | |
| self.enable_bucket = gr.Checkbox(label='Enable buckets', value=True) |