File size: 10,381 Bytes
fe6327d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
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():
            # This use_8bit_adam element should be removed in a future release as it is no longer used
            # use_8bit_adam = gr.Checkbox(
            #     label='Use 8bit adam', value=False, visible=False
            # )
            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.',
            )