mirror of https://github.com/bmaltais/kohya_ss
451 lines
18 KiB
Python
451 lines
18 KiB
Python
import gradio as gr
|
|
from typing import Tuple
|
|
from .custom_logging import setup_logging
|
|
|
|
# Set up logging
|
|
log = setup_logging()
|
|
|
|
class BasicTraining:
|
|
"""
|
|
This class configures and initializes the basic training settings for a machine learning model,
|
|
including options for SDXL, learning rate, learning rate scheduler, and training epochs.
|
|
|
|
Attributes:
|
|
sdxl_checkbox (gr.Checkbox): Checkbox to enable SDXL training.
|
|
learning_rate_value (str): Initial learning rate value.
|
|
lr_scheduler_value (str): Initial learning rate scheduler value.
|
|
lr_warmup_value (str): Initial learning rate warmup value.
|
|
finetuning (bool): If True, enables fine-tuning of the model.
|
|
dreambooth (bool): If True, enables Dreambooth training.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
sdxl_checkbox: gr.Checkbox,
|
|
learning_rate_value: float = "1e-6",
|
|
lr_scheduler_value: str = "constant",
|
|
lr_warmup_value: float = "0",
|
|
lr_warmup_steps_value: int = 0,
|
|
finetuning: bool = False,
|
|
dreambooth: bool = False,
|
|
config: dict = {},
|
|
) -> None:
|
|
"""
|
|
Initializes the BasicTraining object with the given parameters.
|
|
|
|
Args:
|
|
sdxl_checkbox (gr.Checkbox): Checkbox to enable SDXL training.
|
|
learning_rate_value (str): Initial learning rate value.
|
|
lr_scheduler_value (str): Initial learning rate scheduler value.
|
|
lr_warmup_value (str): Initial learning rate warmup value.
|
|
finetuning (bool): If True, enables fine-tuning of the model.
|
|
dreambooth (bool): If True, enables Dreambooth training.
|
|
"""
|
|
self.sdxl_checkbox = sdxl_checkbox
|
|
self.learning_rate_value = learning_rate_value
|
|
self.lr_scheduler_value = lr_scheduler_value
|
|
self.lr_warmup_value = lr_warmup_value
|
|
self.lr_warmup_steps_value= lr_warmup_steps_value
|
|
self.finetuning = finetuning
|
|
self.dreambooth = dreambooth
|
|
self.config = config
|
|
|
|
# Initialize old_lr_warmup and old_lr_warmup_steps with default values
|
|
self.old_lr_warmup = 0
|
|
self.old_lr_warmup_steps = 0
|
|
|
|
# Initialize the UI components
|
|
self.initialize_ui_components()
|
|
|
|
def initialize_ui_components(self) -> None:
|
|
"""
|
|
Initializes the UI components for the training settings.
|
|
"""
|
|
# Initialize the training controls
|
|
self.init_training_controls()
|
|
# Initialize the precision and resources controls
|
|
self.init_precision_and_resources_controls()
|
|
# Initialize the learning rate and optimizer controls
|
|
self.init_lr_and_optimizer_controls()
|
|
# Initialize the gradient and learning rate controls
|
|
self.init_grad_and_lr_controls()
|
|
# Initialize the learning rate controls
|
|
self.init_learning_rate_controls()
|
|
# Initialize the scheduler controls
|
|
self.init_scheduler_controls()
|
|
# Initialize the resolution and bucket controls
|
|
self.init_resolution_and_bucket_controls()
|
|
# Setup the behavior of the SDXL checkbox
|
|
self.setup_sdxl_checkbox_behavior()
|
|
|
|
def init_training_controls(self) -> None:
|
|
"""
|
|
Initializes the training controls for the model.
|
|
"""
|
|
# Create a row for the training controls
|
|
with gr.Row():
|
|
# Initialize the train batch size slider
|
|
self.train_batch_size = gr.Slider(
|
|
minimum=1,
|
|
maximum=64,
|
|
label="Train batch size",
|
|
value=1,
|
|
step=self.config.get("basic.train_batch_size", 1),
|
|
)
|
|
# Initialize the epoch number input
|
|
self.epoch = gr.Number(
|
|
label="Epoch", value=self.config.get("basic.epoch", 1), precision=0
|
|
)
|
|
# Initialize the maximum train epochs input
|
|
self.max_train_epochs = gr.Number(
|
|
label="Max train epoch",
|
|
info="training epochs (overrides max_train_steps). 0 = no override",
|
|
step=1,
|
|
# precision=0,
|
|
minimum=0,
|
|
value=self.config.get("basic.max_train_epochs", 0),
|
|
)
|
|
# Initialize the maximum train steps input
|
|
self.max_train_steps = gr.Number(
|
|
label="Max train steps",
|
|
info="Overrides # training steps. 0 = no override",
|
|
step=1,
|
|
# precision=0,
|
|
value=self.config.get("basic.max_train_steps", 1600),
|
|
)
|
|
# Initialize the save every N epochs input
|
|
self.save_every_n_epochs = gr.Number(
|
|
label="Save every N epochs",
|
|
value=self.config.get("basic.save_every_n_epochs", 1),
|
|
precision=0,
|
|
)
|
|
# Initialize the caption extension input
|
|
self.caption_extension = gr.Dropdown(
|
|
label="Caption file extension",
|
|
choices=["", ".cap", ".caption", ".txt"],
|
|
value=".txt",
|
|
interactive=True,
|
|
)
|
|
|
|
def init_precision_and_resources_controls(self) -> None:
|
|
"""
|
|
Initializes the precision and resources controls for the model.
|
|
"""
|
|
with gr.Row():
|
|
# Initialize the seed textbox
|
|
self.seed = gr.Number(
|
|
label="Seed",
|
|
# precision=0,
|
|
step=1,
|
|
minimum=0,
|
|
value=self.config.get("basic.seed", 0),
|
|
info="Set to 0 to make random",
|
|
)
|
|
# Initialize the cache latents checkbox
|
|
self.cache_latents = gr.Checkbox(
|
|
label="Cache latents",
|
|
value=self.config.get("basic.cache_latents", True),
|
|
)
|
|
# Initialize the cache latents to disk checkbox
|
|
self.cache_latents_to_disk = gr.Checkbox(
|
|
label="Cache latents to disk",
|
|
value=self.config.get("basic.cache_latents_to_disk", False),
|
|
)
|
|
|
|
def init_lr_and_optimizer_controls(self) -> None:
|
|
"""
|
|
Initializes the learning rate and optimizer controls for the model.
|
|
"""
|
|
with gr.Row():
|
|
# Initialize the learning rate scheduler dropdown
|
|
self.lr_scheduler = gr.Dropdown(
|
|
label="LR Scheduler",
|
|
choices=[
|
|
"adafactor",
|
|
"constant",
|
|
"constant_with_warmup",
|
|
"cosine",
|
|
"cosine_with_restarts",
|
|
"linear",
|
|
"piecewise_constant",
|
|
"polynomial",
|
|
"cosine_with_min_lr",
|
|
"inverse_sqrt",
|
|
"warmup_stable_decay",
|
|
],
|
|
value=self.config.get("basic.lr_scheduler", self.lr_scheduler_value),
|
|
)
|
|
|
|
# Initialize the learning rate scheduler type dropdown
|
|
self.lr_scheduler_type = gr.Dropdown(
|
|
label="LR Scheduler type",
|
|
info="(Optional) custom scheduler module name",
|
|
choices=[
|
|
"",
|
|
"CosineAnnealingLR",
|
|
],
|
|
value=self.config.get("basic.lr_scheduler_type", ""),
|
|
allow_custom_value=True,
|
|
)
|
|
|
|
# Initialize the optimizer dropdown
|
|
self.optimizer = gr.Dropdown(
|
|
label="Optimizer",
|
|
choices=[
|
|
"AdamW",
|
|
"AdamWScheduleFree",
|
|
"AdamW8bit",
|
|
"Adafactor",
|
|
"bitsandbytes.optim.AdEMAMix8bit",
|
|
"bitsandbytes.optim.PagedAdEMAMix8bit",
|
|
"DAdaptation",
|
|
"DAdaptAdaGrad",
|
|
"DAdaptAdam",
|
|
"DAdaptAdan",
|
|
"DAdaptAdanIP",
|
|
"DAdaptAdamPreprint",
|
|
"DAdaptLion",
|
|
"DAdaptSGD",
|
|
"Lion",
|
|
"Lion8bit",
|
|
"PagedAdamW8bit",
|
|
"PagedAdamW32bit",
|
|
"PagedLion8bit",
|
|
"Prodigy",
|
|
"prodigyplus.ProdigyPlusScheduleFree",
|
|
"pytorch_optimizer.CAME",
|
|
"RAdamScheduleFree",
|
|
"SGDNesterov",
|
|
"SGDNesterov8bit",
|
|
"SGDScheduleFree",
|
|
],
|
|
value=self.config.get("basic.optimizer", "AdamW8bit"),
|
|
interactive=True,
|
|
allow_custom_value=True,
|
|
)
|
|
|
|
def init_grad_and_lr_controls(self) -> None:
|
|
"""
|
|
Initializes the gradient and learning rate controls for the model.
|
|
"""
|
|
with gr.Row():
|
|
# Initialize the maximum gradient norm slider
|
|
self.max_grad_norm = gr.Number(label='Max grad norm', value=1.0, interactive=True)
|
|
# Initialize the learning rate scheduler extra arguments textbox
|
|
self.lr_scheduler_args = gr.Textbox(
|
|
label="LR scheduler extra arguments",
|
|
lines=2,
|
|
placeholder="(Optional) eg: milestones=[1,10,30,50] gamma=0.1",
|
|
value=self.config.get("basic.lr_scheduler_args", ""),
|
|
)
|
|
# Initialize the optimizer extra arguments textbox
|
|
self.optimizer_args = gr.Textbox(
|
|
label="Optimizer extra arguments",
|
|
lines=2,
|
|
placeholder="(Optional) eg: relative_step=True scale_parameter=True warmup_init=True",
|
|
value=self.config.get("basic.optimizer_args", ""),
|
|
)
|
|
|
|
def init_learning_rate_controls(self) -> None:
|
|
"""
|
|
Initializes the learning rate controls for the model.
|
|
"""
|
|
with gr.Row():
|
|
# Adjust visibility based on training modes
|
|
lr_label = (
|
|
"Learning rate Unet"
|
|
if self.finetuning or self.dreambooth
|
|
else "Learning rate"
|
|
)
|
|
# Initialize the learning rate number input
|
|
self.learning_rate = gr.Number(
|
|
label=lr_label,
|
|
value=self.config.get("basic.learning_rate", self.learning_rate_value),
|
|
minimum=-1,
|
|
maximum=1,
|
|
info="Set to 0 to not train the Unet",
|
|
)
|
|
# Initialize the learning rate TE number input
|
|
self.learning_rate_te = gr.Number(
|
|
label="Learning rate TE",
|
|
value=self.config.get(
|
|
"basic.learning_rate_te", self.learning_rate_value
|
|
),
|
|
visible=self.finetuning or self.dreambooth,
|
|
minimum=-1,
|
|
maximum=1,
|
|
info="Set to 0 to not train the Text Encoder",
|
|
)
|
|
# Initialize the learning rate TE1 number input
|
|
self.learning_rate_te1 = gr.Number(
|
|
label="Learning rate TE1",
|
|
value=self.config.get(
|
|
"basic.learning_rate_te1", self.learning_rate_value
|
|
),
|
|
visible=False,
|
|
minimum=-1,
|
|
maximum=1,
|
|
info="Set to 0 to not train the Text Encoder 1",
|
|
)
|
|
# Initialize the learning rate TE2 number input
|
|
self.learning_rate_te2 = gr.Number(
|
|
label="Learning rate TE2",
|
|
value=self.config.get(
|
|
"basic.learning_rate_te2", self.learning_rate_value
|
|
),
|
|
visible=False,
|
|
minimum=-1,
|
|
maximum=1,
|
|
info="Set to 0 to not train the Text Encoder 2",
|
|
)
|
|
# Initialize the learning rate warmup slider
|
|
self.lr_warmup = gr.Slider(
|
|
label="LR warmup (% of total steps)",
|
|
value=self.config.get("basic.lr_warmup", self.lr_warmup_value),
|
|
minimum=0,
|
|
maximum=100,
|
|
step=1,
|
|
)
|
|
# Initialize the learning rate warmup steps override
|
|
self.lr_warmup_steps = gr.Number(
|
|
label="LR warmup steps (override)",
|
|
value=self.config.get("basic.lr_warmup_steps", self.lr_warmup_steps_value),
|
|
minimum=0,
|
|
step=1,
|
|
)
|
|
|
|
def lr_scheduler_changed(scheduler, value, value_lr_warmup_steps):
|
|
if scheduler == "constant":
|
|
self.old_lr_warmup = value
|
|
self.old_lr_warmup_steps = value_lr_warmup_steps
|
|
value = 0
|
|
value_lr_warmup_steps = 0
|
|
interactive=False
|
|
info="Can't use LR warmup with LR Scheduler constant... setting to 0 and disabling field..."
|
|
else:
|
|
if self.old_lr_warmup != 0:
|
|
value = self.old_lr_warmup
|
|
self.old_lr_warmup = 0
|
|
if self.old_lr_warmup_steps != 0:
|
|
value_lr_warmup_steps = self.old_lr_warmup_steps
|
|
self.old_lr_warmup_steps = 0
|
|
interactive=True
|
|
info=""
|
|
return gr.Slider(value=value, interactive=interactive, info=info), gr.Number(value=value_lr_warmup_steps, interactive=interactive, info=info)
|
|
|
|
self.lr_scheduler.change(
|
|
lr_scheduler_changed,
|
|
inputs=[self.lr_scheduler, self.lr_warmup, self.lr_warmup_steps],
|
|
outputs=[self.lr_warmup, self.lr_warmup_steps],
|
|
)
|
|
|
|
def init_scheduler_controls(self) -> None:
|
|
"""
|
|
Initializes the scheduler controls for the model.
|
|
"""
|
|
with gr.Row(visible=not self.finetuning):
|
|
# Initialize the learning rate scheduler number of cycles textbox
|
|
self.lr_scheduler_num_cycles = gr.Number(
|
|
label="LR # cycles",
|
|
minimum=1,
|
|
# precision=0, # round to nearest integer
|
|
step=1, # Increment value by 1
|
|
info="Number of restarts for cosine scheduler with restarts",
|
|
value=self.config.get("basic.lr_scheduler_num_cycles", 1),
|
|
)
|
|
# Initialize the learning rate scheduler power textbox
|
|
self.lr_scheduler_power = gr.Number(
|
|
label="LR power",
|
|
minimum=0.0,
|
|
step=0.01,
|
|
info="Polynomial power for polynomial scheduler",
|
|
value=self.config.get("basic.lr_scheduler_power", 1.0),
|
|
)
|
|
|
|
def init_resolution_and_bucket_controls(self) -> None:
|
|
"""
|
|
Initializes the resolution and bucket controls for the model.
|
|
"""
|
|
with gr.Row(visible=not self.finetuning):
|
|
# Initialize the maximum resolution textbox
|
|
self.max_resolution = gr.Textbox(
|
|
label="Max resolution",
|
|
value=self.config.get("basic.max_resolution", "512,512"),
|
|
placeholder="512,512",
|
|
)
|
|
# Initialize the stop text encoder training slider
|
|
self.stop_text_encoder_training = gr.Slider(
|
|
minimum=-1,
|
|
maximum=100,
|
|
value=self.config.get("basic.stop_text_encoder_training", 0),
|
|
step=1,
|
|
label="Stop TE (% of total steps)",
|
|
)
|
|
# Initialize the enable buckets checkbox
|
|
self.enable_bucket = gr.Checkbox(
|
|
label="Enable buckets",
|
|
value=self.config.get("basic.enable_bucket", True),
|
|
)
|
|
# Initialize the minimum bucket resolution slider
|
|
self.min_bucket_reso = gr.Slider(
|
|
label="Minimum bucket resolution",
|
|
value=self.config.get("basic.min_bucket_reso", 256),
|
|
minimum=64,
|
|
maximum=4096,
|
|
step=64,
|
|
info="Minimum size in pixel a bucket can be (>= 64)",
|
|
)
|
|
# Initialize the maximum bucket resolution slider
|
|
self.max_bucket_reso = gr.Slider(
|
|
label="Maximum bucket resolution",
|
|
value=self.config.get("basic.max_bucket_reso", 2048),
|
|
minimum=64,
|
|
maximum=4096,
|
|
step=64,
|
|
info="Maximum size in pixel a bucket can be (>= 64)",
|
|
)
|
|
|
|
def setup_sdxl_checkbox_behavior(self) -> None:
|
|
"""
|
|
Sets up the behavior of the SDXL checkbox based on the finetuning and dreambooth flags.
|
|
"""
|
|
self.sdxl_checkbox.change(
|
|
self.update_learning_rate_te,
|
|
inputs=[
|
|
self.sdxl_checkbox,
|
|
gr.Checkbox(value=self.finetuning, visible=False),
|
|
gr.Checkbox(value=self.dreambooth, visible=False),
|
|
],
|
|
outputs=[
|
|
self.learning_rate_te,
|
|
self.learning_rate_te1,
|
|
self.learning_rate_te2,
|
|
],
|
|
)
|
|
|
|
def update_learning_rate_te(
|
|
self,
|
|
sdxl_checkbox: gr.Checkbox,
|
|
finetuning: bool,
|
|
dreambooth: bool,
|
|
) -> Tuple[gr.Number, gr.Number, gr.Number]:
|
|
"""
|
|
Updates the visibility of the learning rate TE, TE1, and TE2 based on the SDXL checkbox and finetuning/dreambooth flags.
|
|
|
|
Args:
|
|
sdxl_checkbox (gr.Checkbox): The SDXL checkbox.
|
|
finetuning (bool): Whether finetuning is enabled.
|
|
dreambooth (bool): Whether dreambooth is enabled.
|
|
|
|
Returns:
|
|
Tuple[gr.Number, gr.Number, gr.Number]: A tuple containing the updated visibility for learning rate TE, TE1, and TE2.
|
|
"""
|
|
# Determine the visibility condition based on finetuning and dreambooth flags
|
|
visibility_condition = finetuning or dreambooth
|
|
# Return a tuple of gr.Number instances with updated visibility
|
|
return (
|
|
gr.Number(visible=(not sdxl_checkbox and visibility_condition)),
|
|
gr.Number(visible=(sdxl_checkbox and visibility_condition)),
|
|
gr.Number(visible=(sdxl_checkbox and visibility_condition)),
|
|
)
|