Skip to content

Commit

Permalink
Update lr_warmup_steps code
Browse files Browse the repository at this point in the history
  • Loading branch information
@bmaltais
bmaltais committed Sep 13, 2024
1 parent e655ebf commit 20d7e65
Show file tree
Hide file tree
Showing 4 changed files with 85 additions and 184 deletions.
38 changes: 10 additions & 28 deletions kohya_gui/dreambooth_gui.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -696,22 +696,6 @@ def train_model(
# End of path validation
#

# This function validates files or folder paths. Simply add new variables containing file of folder path
# to validate below
# if not validate_paths(
# dataset_config=dataset_config,
# headless=headless,
# log_tracker_config=log_tracker_config,
# logging_dir=logging_dir,
# output_dir=output_dir,
# pretrained_model_name_or_path=pretrained_model_name_or_path,
# reg_data_dir=reg_data_dir,
# resume=resume,
# train_data_dir=train_data_dir,
# vae=vae,
# ):
# return TRAIN_BUTTON_VISIBLE

if not print_only and check_if_model_exist(
output_name, output_dir, save_model_as, headless=headless
):
Expand All @@ -721,13 +705,6 @@ def train_model(
log.info(
"Dataset config toml file used, skipping total_steps, train_batch_size, gradient_accumulation_steps, epoch, reg_factor, max_train_steps calculations..."
)
if max_train_steps > 0:
if lr_warmup_steps > 0:
lr_warmup_steps = int(lr_warmup_steps)
else:
lr_warmup_steps = float(lr_warmup / 100) if lr_warmup != 0 else 0
else:
lr_warmup_steps = 0

if max_train_steps == 0:
max_train_steps_info = f"Max train steps: 0. sd-scripts will therefore default to 1600. Please specify a different value if required."
Expand Down Expand Up @@ -810,13 +787,18 @@ def train_model(
else:
max_train_steps_info = f"Max train steps: {max_train_steps}"

if lr_warmup_steps > 0:
lr_warmup_steps = int(lr_warmup_steps)
else:
lr_warmup_steps = float(lr_warmup / 100) if lr_warmup != 0 else 0

log.info(f"Total steps: {total_steps}")

# Calculate lr_warmup_steps
if lr_warmup_steps > 0:
lr_warmup_steps = int(lr_warmup_steps)
if lr_warmup > 0:
log.warning("Both lr_warmup and lr_warmup_steps are set. lr_warmup_steps will be used.")
elif lr_warmup != 0:
lr_warmup_steps = lr_warmup / 100
else:
lr_warmup_steps = 0

log.info(f"Train batch size: {train_batch_size}")
log.info(f"Gradient accumulation steps: {gradient_accumulation_steps}")
log.info(f"Epoch: {epoch}")
Expand Down
25 changes: 8 additions & 17 deletions kohya_gui/finetune_gui.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -726,18 +726,6 @@ def train_model(
# End of path validation
#

# if not validate_paths(
# dataset_config=dataset_config,
# finetune_image_folder=image_folder,
# headless=headless,
# log_tracker_config=log_tracker_config,
# logging_dir=logging_dir,
# output_dir=output_dir,
# pretrained_model_name_or_path=pretrained_model_name_or_path,
# resume=resume,
# ):
# return TRAIN_BUTTON_VISIBLE

if not print_only and check_if_model_exist(
output_name, output_dir, save_model_as, headless
):
Expand Down Expand Up @@ -868,13 +856,16 @@ def train_model(

log.info(max_train_steps_info)

if max_train_steps != 0:
if lr_warmup_steps > 0:
lr_warmup_steps = int(lr_warmup_steps)
else:
lr_warmup_steps = float(lr_warmup / 100) if lr_warmup != 0 else 0
# Calculate lr_warmup_steps
if lr_warmup_steps > 0:
lr_warmup_steps = int(lr_warmup_steps)
if lr_warmup > 0:
log.warning("Both lr_warmup and lr_warmup_steps are set. lr_warmup_steps will be used.")
elif lr_warmup != 0:
lr_warmup_steps = lr_warmup / 100
else:
lr_warmup_steps = 0

log.info(f"lr_warmup_steps = {lr_warmup_steps}")

accelerate_path = get_executable_path("accelerate")
Expand Down
172 changes: 57 additions & 115 deletions kohya_gui/lora_gui.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -894,21 +894,6 @@ def train_model(
# End of path validation
#

# if not validate_paths(
# dataset_config=dataset_config,
# headless=headless,
# log_tracker_config=log_tracker_config,
# logging_dir=logging_dir,
# network_weights=network_weights,
# output_dir=output_dir,
# pretrained_model_name_or_path=pretrained_model_name_or_path,
# reg_data_dir=reg_data_dir,
# resume=resume,
# train_data_dir=train_data_dir,
# vae=vae,
# ):
# return TRAIN_BUTTON_VISIBLE

if int(bucket_reso_steps) < 1:
output_message(
msg="Bucket resolution steps need to be greater than 0",
Expand Down Expand Up @@ -949,123 +934,80 @@ def train_model(
# unet_lr = 0

if dataset_config:
log.info(
"Dataset config toml file used, skipping total_steps, train_batch_size, gradient_accumulation_steps, epoch, reg_factor, max_train_steps calculations..."
)
if max_train_steps > 0:
# calculate stop encoder training
if stop_text_encoder_training == 0:
stop_text_encoder_training = 0
else:
stop_text_encoder_training = math.ceil(
float(max_train_steps) / 100 * int(stop_text_encoder_training)
)

if lr_warmup_steps > 0:
lr_warmup_steps = int(lr_warmup_steps)
else:
lr_warmup_steps = float(lr_warmup / 100) if lr_warmup != 0 else 0
else:
stop_text_encoder_training = 0
lr_warmup_steps = 0

if max_train_steps == 0:
max_train_steps_info = f"Max train steps: 0. sd-scripts will therefore default to 1600. Please specify a different value if required."
else:
max_train_steps_info = f"Max train steps: {max_train_steps}"

log.info("Dataset config TOML file used; skipping calculations for total_steps, train_batch_size, gradient_accumulation_steps, epoch, reg_factor, and max_train_steps.")
else:
if train_data_dir == "":
log.error("Train data dir is empty")
if not train_data_dir:
log.error("Train data directory is empty.")
return TRAIN_BUTTON_VISIBLE

# Get a list of all subfolders in train_data_dir
subfolders = [
f
for f in os.listdir(train_data_dir)
f for f in os.listdir(train_data_dir)
if os.path.isdir(os.path.join(train_data_dir, f))
]

total_steps = 0

# Loop through each subfolder and extract the number of repeats
for folder in subfolders:
try:
# Extract the number of repeats from the folder name
repeats = int(folder.split("_")[0])
log.info(f"Folder {folder}: {repeats} repeats found")

# Count the number of images in the folder
num_images = len(
[
f
for f, lower_f in (
(file, file.lower())
for file in os.listdir(os.path.join(train_data_dir, folder))
)
if lower_f.endswith((".jpg", ".jpeg", ".png", ".webp"))
]
)

log.info(f"Folder {folder}: {num_images} images found")

# Calculate the total number of steps for this folder
steps = repeats * num_images

# log.info the result
log.info(f"Folder {folder}: {num_images} * {repeats} = {steps} steps")

total_steps += steps

except ValueError:
# Handle the case where the folder name does not contain an underscore
log.info(
f"Error: '{folder}' does not contain an underscore, skipping..."
)

if reg_data_dir == "":
reg_factor = 1
else:
log.warning(
"Regularisation images are used... Will double the number of steps required..."
)
reg_factor = 2

log.info(f"Regulatization factor: {reg_factor}")
repeats_str = folder.split("_")[0]
repeats = int(repeats_str)
log.info(f"Folder '{folder}': {repeats} repeats found.")
except (ValueError, IndexError):
log.info(f"Skipping folder '{folder}': unable to extract repeat count.")
continue

folder_path = os.path.join(train_data_dir, folder)
image_extensions = (".jpg", ".jpeg", ".png", ".webp")
num_images = len([
file for file in os.listdir(folder_path)
if file.lower().endswith(image_extensions)
])
log.info(f"Folder '{folder}': {num_images} images found.")

steps = repeats * num_images
log.info(f"Folder '{folder}': {num_images} images * {repeats} repeats = {steps} steps.")
total_steps += steps

reg_factor = 2 if reg_data_dir else 1
if reg_factor == 2:
log.warning("Regularization images are used; the number of required steps will be doubled.")

log.info(f"Regularization factor: {reg_factor}")

if max_train_steps == 0:
# calculate max_train_steps
max_train_steps = int(
math.ceil(
float(total_steps)
/ int(train_batch_size)
/ int(gradient_accumulation_steps)
* int(epoch)
* int(reg_factor)
)
if train_batch_size == 0 or gradient_accumulation_steps == 0:
log.error("train_batch_size and gradient_accumulation_steps must be greater than zero.")
return TRAIN_BUTTON_VISIBLE

max_train_steps = int(math.ceil(
total_steps / train_batch_size / gradient_accumulation_steps * epoch * reg_factor
))
max_train_steps_info = (
f"Calculated max_train_steps: ({total_steps} / {train_batch_size} / "
f"{gradient_accumulation_steps} * {epoch} * {reg_factor}) = {max_train_steps}"
)
max_train_steps_info = f"max_train_steps ({total_steps} / {train_batch_size} / {gradient_accumulation_steps} * {epoch} * {reg_factor}) = {max_train_steps}"
else:
if max_train_steps == 0:
max_train_steps_info = f"Max train steps: 0. sd-scripts will therefore default to 1600. Please specify a different value if required."
else:
max_train_steps_info = f"Max train steps: {max_train_steps}"

# calculate stop encoder training
if stop_text_encoder_training == 0:
stop_text_encoder_training = 0
else:
stop_text_encoder_training = math.ceil(
float(max_train_steps) / 100 * int(stop_text_encoder_training)
)

if lr_warmup_steps > 0:
lr_warmup_steps = int(lr_warmup_steps)
else:
lr_warmup_steps = float(lr_warmup / 100) if lr_warmup != 0 else 0
max_train_steps_info = f"Max train steps: {max_train_steps}"

log.info(f"Total steps: {total_steps}")

# Calculate stop_text_encoder_training
if max_train_steps > 0 and stop_text_encoder_training > 0:
stop_text_encoder_training = math.ceil(
max_train_steps * stop_text_encoder_training / 100
)
else:
stop_text_encoder_training = 0

# Calculate lr_warmup_steps
if lr_warmup_steps > 0:
lr_warmup_steps = int(lr_warmup_steps)
if lr_warmup > 0:
log.warning("Both lr_warmup and lr_warmup_steps are set. lr_warmup_steps will be used.")
elif lr_warmup != 0:
lr_warmup_steps = lr_warmup / 100
else:
lr_warmup_steps = 0

log.info(f"Train batch size: {train_batch_size}")
log.info(f"Gradient accumulation steps: {gradient_accumulation_steps}")
log.info(f"Epoch: {epoch}")
Expand Down
34 changes: 10 additions & 24 deletions kohya_gui/textual_inversion_gui.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -561,20 +561,6 @@ def train_model(
# End of path validation
#

# if not validate_paths(
# dataset_config=dataset_config,
# headless=headless,
# log_tracker_config=log_tracker_config,
# logging_dir=logging_dir,
# output_dir=output_dir,
# pretrained_model_name_or_path=pretrained_model_name_or_path,
# reg_data_dir=reg_data_dir,
# resume=resume,
# train_data_dir=train_data_dir,
# vae=vae,
# ):
# return TRAIN_BUTTON_VISIBLE

if token_string == "":
output_message(msg="Token string is missing", headless=headless)
return TRAIN_BUTTON_VISIBLE
Expand All @@ -600,11 +586,6 @@ def train_model(
stop_text_encoder_training = math.ceil(
float(max_train_steps) / 100 * int(stop_text_encoder_training_pct)
)

if lr_warmup_steps > 0:
lr_warmup_steps = int(lr_warmup_steps)
else:
lr_warmup_steps = float(lr_warmup / 100) if lr_warmup != 0 else 0
else:
stop_text_encoder_training = 0
lr_warmup_steps = 0
Expand Down Expand Up @@ -699,13 +680,18 @@ def train_model(
float(max_train_steps) / 100 * int(stop_text_encoder_training_pct)
)

if lr_warmup_steps > 0:
lr_warmup_steps = int(lr_warmup_steps)
else:
lr_warmup_steps = float(lr_warmup / 100) if lr_warmup != 0 else 0

log.info(f"Total steps: {total_steps}")

# Calculate lr_warmup_steps
if lr_warmup_steps > 0:
lr_warmup_steps = int(lr_warmup_steps)
if lr_warmup > 0:
log.warning("Both lr_warmup and lr_warmup_steps are set. lr_warmup_steps will be used.")
elif lr_warmup != 0:
lr_warmup_steps = lr_warmup / 100
else:
lr_warmup_steps = 0

log.info(f"Train batch size: {train_batch_size}")
log.info(f"Gradient accumulation steps: {gradient_accumulation_steps}")
log.info(f"Epoch: {epoch}")
Expand Down

0 comments on commit 20d7e65

Please sign in to comment.