common_gui.py

from tkinter import filedialog, Tk
from easygui import msgbox
import os
import re
import gradio as gr
import easygui
import shutil
import sys
import json

from library.custom_logging import setup_logging
from datetime import datetime

# Set up logging
log = setup_logging()

folder_symbol = "\U0001f4c2"  # 📂
refresh_symbol = "\U0001f504"  # 🔄
save_style_symbol = "\U0001f4be"  # 💾
document_symbol = "\U0001F4C4"  # 📄

# define a list of substrings to search for v2 base models
V2_BASE_MODELS = [
    "stabilityai/stable-diffusion-2-1-base/blob/main/v2-1_512-ema-pruned",
    "stabilityai/stable-diffusion-2-1-base",
    "stabilityai/stable-diffusion-2-base",
]

# define a list of substrings to search for v_parameterization models
V_PARAMETERIZATION_MODELS = [
    "stabilityai/stable-diffusion-2-1/blob/main/v2-1_768-ema-pruned",
    "stabilityai/stable-diffusion-2-1",
    "stabilityai/stable-diffusion-2",
]

# define a list of substrings to v1.x models
V1_MODELS = [
    "CompVis/stable-diffusion-v1-4",
    "runwayml/stable-diffusion-v1-5",
]

# define a list of substrings to search for SDXL base models
SDXL_MODELS = [
    "stabilityai/stable-diffusion-xl-base-1.0",
    "stabilityai/stable-diffusion-xl-refiner-1.0",
]

# define a list of substrings to search for
ALL_PRESET_MODELS = V2_BASE_MODELS + V_PARAMETERIZATION_MODELS + V1_MODELS + SDXL_MODELS

ENV_EXCLUSION = ["COLAB_GPU", "RUNPOD_POD_ID"]


def check_if_model_exist(output_name, output_dir, save_model_as, headless=False):
    if headless:
        log.info(
            "Headless mode, skipping verification if model already exist... if model already exist it will be overwritten..."
        )
        return False

    if save_model_as in ["diffusers", "diffusers_safetendors"]:
        ckpt_folder = os.path.join(output_dir, output_name)
        if os.path.isdir(ckpt_folder):
            msg = f"A diffuser model with the same name {ckpt_folder} already exists. Do you want to overwrite it?"
            if not easygui.ynbox(msg, "Overwrite Existing Model?"):
                log.info("Aborting training due to existing model with same name...")
                return True
    elif save_model_as in ["ckpt", "safetensors"]:
        ckpt_file = os.path.join(output_dir, output_name + "." + save_model_as)
        if os.path.isfile(ckpt_file):
            msg = f"A model with the same file name {ckpt_file} already exists. Do you want to overwrite it?"
            if not easygui.ynbox(msg, "Overwrite Existing Model?"):
                log.info("Aborting training due to existing model with same name...")
                return True
    else:
        log.info(
            'Can\'t verify if existing model exist when save model is set a "same as source model", continuing to train model...'
        )
        return False

    return False


def output_message(msg="", title="", headless=False):
    if headless:
        log.info(msg)
    else:
        msgbox(msg=msg, title=title)


def update_my_data(my_data):
    # Update the optimizer based on the use_8bit_adam flag
    use_8bit_adam = my_data.get("use_8bit_adam", False)
    my_data.setdefault("optimizer", "AdamW8bit" if use_8bit_adam else "AdamW")

    # Update model_list to custom if empty or pretrained_model_name_or_path is not a preset model
    model_list = my_data.get("model_list", [])
    pretrained_model_name_or_path = my_data.get("pretrained_model_name_or_path", "")
    if not model_list or pretrained_model_name_or_path not in ALL_PRESET_MODELS:
        my_data["model_list"] = "custom"

    # Convert values to int if they are strings
    for key in ["epoch", "save_every_n_epochs", "lr_warmup"]:
        value = my_data.get(key, 0)
        if isinstance(value, str) and value.strip().isdigit():
            my_data[key] = int(value)
        elif not value:
            my_data[key] = 0

    # Convert values to float if they are strings
    for key in ["noise_offset", "learning_rate", "text_encoder_lr", "unet_lr"]:
        value = my_data.get(key, 0)
        if isinstance(value, str) and value.strip().isdigit():
            my_data[key] = float(value)
        elif not value:
            my_data[key] = 0

    # Update LoRA_type if it is set to LoCon
    if my_data.get("LoRA_type", "Standard") == "LoCon":
        my_data["LoRA_type"] = "LyCORIS/LoCon"

    # Update model save choices due to changes for LoRA and TI training
    if "save_model_as" in my_data:
        if (
            my_data.get("LoRA_type") or my_data.get("num_vectors_per_token")
        ) and my_data.get("save_model_as") not in ["safetensors", "ckpt"]:
            message = "Updating save_model_as to safetensors because the current value in the config file is no longer applicable to {}"
            if my_data.get("LoRA_type"):
                log.info(message.format("LoRA"))
            if my_data.get("num_vectors_per_token"):
                log.info(message.format("TI"))
            my_data["save_model_as"] = "safetensors"

    # Update xformers if it is set to True and is a boolean
    xformers_value = my_data.get("xformers", None)
    if isinstance(xformers_value, bool):
        if xformers_value:
            my_data["xformers"] = "xformers"
        else:
            my_data["xformers"] = "none"

    return my_data


def get_dir_and_file(file_path):
    dir_path, file_name = os.path.split(file_path)
    return (dir_path, file_name)


def get_file_path(
    file_path="", default_extension=".json", extension_name="Config files"
):
    if not any(var in os.environ for var in ENV_EXCLUSION) and sys.platform != "darwin":
        current_file_path = file_path
        # log.info(f'current file path: {current_file_path}')

        initial_dir, initial_file = get_dir_and_file(file_path)

        # Create a hidden Tkinter root window
        root = Tk()
        root.wm_attributes("-topmost", 1)
        root.withdraw()

        # Show the open file dialog and get the selected file path
        file_path = filedialog.askopenfilename(
            filetypes=(
                (extension_name, f"*{default_extension}"),
                ("All files", "*.*"),
            ),
            defaultextension=default_extension,
            initialfile=initial_file,
            initialdir=initial_dir,
        )

        # Destroy the hidden root window
        root.destroy()

        # If no file is selected, use the current file path
        if not file_path:
            file_path = current_file_path
        current_file_path = file_path
        # log.info(f'current file path: {current_file_path}')

    return file_path


def get_any_file_path(file_path=""):
    if not any(var in os.environ for var in ENV_EXCLUSION) and sys.platform != "darwin":
        current_file_path = file_path
        # log.info(f'current file path: {current_file_path}')

        initial_dir, initial_file = get_dir_and_file(file_path)

        root = Tk()
        root.wm_attributes("-topmost", 1)
        root.withdraw()
        file_path = filedialog.askopenfilename(
            initialdir=initial_dir,
            initialfile=initial_file,
        )
        root.destroy()

        if file_path == "":
            file_path = current_file_path

    return file_path


def remove_doublequote(file_path):
    if file_path != None:
        file_path = file_path.replace('"', "")

    return file_path


def get_folder_path(folder_path=""):
    if not any(var in os.environ for var in ENV_EXCLUSION) and sys.platform != "darwin":
        current_folder_path = folder_path

        initial_dir, initial_file = get_dir_and_file(folder_path)

        root = Tk()
        root.wm_attributes("-topmost", 1)
        root.withdraw()
        folder_path = filedialog.askdirectory(initialdir=initial_dir)
        root.destroy()

        if folder_path == "":
            folder_path = current_folder_path

    return folder_path


def get_saveasfile_path(
    file_path="", defaultextension=".json", extension_name="Config files"
):
    if not any(var in os.environ for var in ENV_EXCLUSION) and sys.platform != "darwin":
        current_file_path = file_path
        # log.info(f'current file path: {current_file_path}')

        initial_dir, initial_file = get_dir_and_file(file_path)

        root = Tk()
        root.wm_attributes("-topmost", 1)
        root.withdraw()
        save_file_path = filedialog.asksaveasfile(
            filetypes=(
                (f"{extension_name}", f"{defaultextension}"),
                ("All files", "*"),
            ),
            defaultextension=defaultextension,
            initialdir=initial_dir,
            initialfile=initial_file,
        )
        root.destroy()

        # log.info(save_file_path)

        if save_file_path == None:
            file_path = current_file_path
        else:
            log.info(save_file_path.name)
            file_path = save_file_path.name

        # log.info(file_path)

    return file_path


def get_saveasfilename_path(
    file_path="", extensions="*", extension_name="Config files"
):
    if not any(var in os.environ for var in ENV_EXCLUSION) and sys.platform != "darwin":
        current_file_path = file_path
        # log.info(f'current file path: {current_file_path}')

        initial_dir, initial_file = get_dir_and_file(file_path)

        root = Tk()
        root.wm_attributes("-topmost", 1)
        root.withdraw()
        save_file_path = filedialog.asksaveasfilename(
            filetypes=(
                (f"{extension_name}", f"{extensions}"),
                ("All files", "*"),
            ),
            defaultextension=extensions,
            initialdir=initial_dir,
            initialfile=initial_file,
        )
        root.destroy()

        if save_file_path == "":
            file_path = current_file_path
        else:
            # log.info(save_file_path)
            file_path = save_file_path

    return file_path


def add_pre_postfix(
    folder: str = "",
    prefix: str = "",
    postfix: str = "",
    caption_file_ext: str = ".caption",
) -> None:
    """
    Add prefix and/or postfix to the content of caption files within a folder.
    If no caption files are found, create one with the requested prefix and/or postfix.

    Args:
        folder (str): Path to the folder containing caption files.
        prefix (str, optional): Prefix to add to the content of the caption files.
        postfix (str, optional): Postfix to add to the content of the caption files.
        caption_file_ext (str, optional): Extension of the caption files.
    """

    if prefix == "" and postfix == "":
        return

    image_extensions = (".jpg", ".jpeg", ".png", ".webp")
    image_files = [
        f for f in os.listdir(folder) if f.lower().endswith(image_extensions)
    ]

    for image_file in image_files:
        caption_file_name = os.path.splitext(image_file)[0] + caption_file_ext
        caption_file_path = os.path.join(folder, caption_file_name)

        if not os.path.exists(caption_file_path):
            with open(caption_file_path, "w", encoding="utf8") as f:
                separator = " " if prefix and postfix else ""
                f.write(f"{prefix}{separator}{postfix}")
        else:
            with open(caption_file_path, "r+", encoding="utf8") as f:
                content = f.read()
                content = content.rstrip()
                f.seek(0, 0)

                prefix_separator = " " if prefix else ""
                postfix_separator = " " if postfix else ""
                f.write(
                    f"{prefix}{prefix_separator}{content}{postfix_separator}{postfix}"
                )


def has_ext_files(folder_path: str, file_extension: str) -> bool:
    """
    Check if there are any files with the specified extension in the given folder.

    Args:
        folder_path (str): Path to the folder containing files.
        file_extension (str): Extension of the files to look for.

    Returns:
        bool: True if files with the specified extension are found, False otherwise.
    """
    for file in os.listdir(folder_path):
        if file.endswith(file_extension):
            return True
    return False


def find_replace(
    folder_path: str = "",
    caption_file_ext: str = ".caption",
    search_text: str = "",
    replace_text: str = "",
) -> None:
    """
    Find and replace text in caption files within a folder.

    Args:
        folder_path (str, optional): Path to the folder containing caption files.
        caption_file_ext (str, optional): Extension of the caption files.
        search_text (str, optional): Text to search for in the caption files.
        replace_text (str, optional): Text to replace the search text with.
    """
    log.info("Running caption find/replace")

    if not has_ext_files(folder_path, caption_file_ext):
        msgbox(
            f"No files with extension {caption_file_ext} were found in {folder_path}..."
        )
        return

    if search_text == "":
        return

    caption_files = [f for f in os.listdir(folder_path) if f.endswith(caption_file_ext)]

    for caption_file in caption_files:
        with open(os.path.join(folder_path, caption_file), "r", errors="ignore") as f:
            content = f.read()

        content = content.replace(search_text, replace_text)

        with open(os.path.join(folder_path, caption_file), "w") as f:
            f.write(content)


def color_aug_changed(color_aug):
    if color_aug:
        msgbox(
            'Disabling "Cache latent" because "Color augmentation" has been selected...'
        )
        return gr.Checkbox(value=False, interactive=False)
    else:
        return gr.Checkbox(value=True, interactive=True)


def save_inference_file(output_dir, v2, v_parameterization, output_name):
    # List all files in the directory
    files = os.listdir(output_dir)

    # Iterate over the list of files
    for file in files:
        # Check if the file starts with the value of output_name
        if file.startswith(output_name):
            # Check if it is a file or a directory
            if os.path.isfile(os.path.join(output_dir, file)):
                # Split the file name and extension
                file_name, ext = os.path.splitext(file)

                # Copy the v2-inference-v.yaml file to the current file, with a .yaml extension
                if v2 and v_parameterization:
                    log.info(
                        f"Saving v2-inference-v.yaml as {output_dir}/{file_name}.yaml"
                    )
                    shutil.copy(
                        f"./v2_inference/v2-inference-v.yaml",
                        f"{output_dir}/{file_name}.yaml",
                    )
                elif v2:
                    log.info(
                        f"Saving v2-inference.yaml as {output_dir}/{file_name}.yaml"
                    )
                    shutil.copy(
                        f"./v2_inference/v2-inference.yaml",
                        f"{output_dir}/{file_name}.yaml",
                    )


def set_pretrained_model_name_or_path_input(
    model_list,
    pretrained_model_name_or_path,
    pretrained_model_name_or_path_file,
    pretrained_model_name_or_path_folder,
    v2,
    v_parameterization,
    sdxl,
):
    # Check if the given model_list is in the list of SDXL models
    if str(model_list) in SDXL_MODELS:
        log.info("SDXL model selected. Setting sdxl parameters")
        v2 = gr.Checkbox(value=False, visible=False)
        v_parameterization = gr.Checkbox(value=False, visible=False)
        sdxl = gr.Checkbox(value=True, visible=False)
        pretrained_model_name_or_path = gr.Textbox(
            value=str(model_list), visible=False
        )
        pretrained_model_name_or_path_file = gr.Button(visible=False)
        pretrained_model_name_or_path_folder = gr.Button(visible=False)
        return (
            model_list,
            pretrained_model_name_or_path,
            pretrained_model_name_or_path_file,
            pretrained_model_name_or_path_folder,
            v2,
            v_parameterization,
            sdxl,
        )

    # Check if the given model_list is in the list of V2 base models
    if str(model_list) in V2_BASE_MODELS:
        log.info("SD v2 base model selected. Setting --v2 parameter")
        v2 = gr.Checkbox(value=True, visible=False)
        v_parameterization = gr.Checkbox(value=False, visible=False)
        sdxl = gr.Checkbox(value=False, visible=False)
        pretrained_model_name_or_path = gr.Textbox(
            value=str(model_list), visible=False
        )
        pretrained_model_name_or_path_file = gr.Button(visible=False)
        pretrained_model_name_or_path_folder = gr.Button(visible=False)
        return (
            model_list,
            pretrained_model_name_or_path,
            pretrained_model_name_or_path_file,
            pretrained_model_name_or_path_folder,
            v2,
            v_parameterization,
            sdxl,
        )

    # Check if the given model_list is in the list of V parameterization models
    if str(model_list) in V_PARAMETERIZATION_MODELS:
        log.info(
            "SD v2 model selected. Setting --v2 and --v_parameterization parameters"
        )
        v2 = gr.Checkbox(value=True, visible=False)
        v_parameterization = gr.Checkbox(value=True, visible=False)
        sdxl = gr.Checkbox(value=False, visible=False)
        pretrained_model_name_or_path = gr.Textbox(
            value=str(model_list), visible=False
        )
        pretrained_model_name_or_path_file = gr.Button(visible=False)
        pretrained_model_name_or_path_folder = gr.Button(visible=False)
        return (
            model_list,
            pretrained_model_name_or_path,
            pretrained_model_name_or_path_file,
            pretrained_model_name_or_path_folder,
            v2,
            v_parameterization,
            sdxl,
        )

    # Check if the given model_list is in the list of V1 models
    if str(model_list) in V1_MODELS:
        log.info(f"{model_list} model selected.")
        v2 = gr.Checkbox(value=False, visible=False)
        v_parameterization = gr.Checkbox(value=False, visible=False)
        sdxl = gr.Checkbox(value=False, visible=False)
        pretrained_model_name_or_path = gr.Textbox(
            value=str(model_list), visible=False
        )
        pretrained_model_name_or_path_file = gr.Button(visible=False)
        pretrained_model_name_or_path_folder = gr.Button(visible=False)
        return (
            model_list,
            pretrained_model_name_or_path,
            pretrained_model_name_or_path_file,
            pretrained_model_name_or_path_folder,
            v2,
            v_parameterization,
            sdxl,
        )

    # Check if the model_list is set to 'custom'
    if model_list == "custom":
        v2 = gr.Checkbox(visible=True)
        v_parameterization = gr.Checkbox(visible=True)
        sdxl = gr.Checkbox(visible=True)
        pretrained_model_name_or_path = gr.Textbox(visible=True)
        pretrained_model_name_or_path_file = gr.Button(visible=True)
        pretrained_model_name_or_path_folder = gr.Button(visible=True)
        return (
            model_list,
            pretrained_model_name_or_path,
            pretrained_model_name_or_path_file,
            pretrained_model_name_or_path_folder,
            v2,
            v_parameterization,
            sdxl,
        )


###
### Gradio common GUI section
###


def get_pretrained_model_name_or_path_file(model_list, pretrained_model_name_or_path):
    pretrained_model_name_or_path = get_any_file_path(pretrained_model_name_or_path)
    # set_model_list(model_list, pretrained_model_name_or_path)


def get_int_or_default(kwargs, key, default_value=0):
    value = kwargs.get(key, default_value)
    if isinstance(value, int):
        return value
    elif isinstance(value, str):
        return int(value)
    elif isinstance(value, float):
        return int(value)
    else:
        log.info(
            f"{key} is not an int, float or a string, setting value to {default_value}"
        )
        return default_value


def get_float_or_default(kwargs, key, default_value=0.0):
    value = kwargs.get(key, default_value)
    if isinstance(value, float):
        return value
    elif isinstance(value, int):
        return float(value)
    elif isinstance(value, str):
        return float(value)
    else:
        log.info(
            f"{key} is not an int, float or a string, setting value to {default_value}"
        )
        return default_value


def get_str_or_default(kwargs, key, default_value=""):
    value = kwargs.get(key, default_value)
    if isinstance(value, str):
        return value
    elif isinstance(value, int):
        return str(value)
    elif isinstance(value, str):
        return str(value)
    else:
        return default_value


# def run_cmd_training(**kwargs):
#     run_cmd = ""

#     lr_scheduler = kwargs.get("lr_scheduler", "")
#     if lr_scheduler:
#         run_cmd += f' --lr_scheduler="{lr_scheduler}"'

#     lr_warmup_steps = kwargs.get("lr_warmup_steps", "")
#     if lr_warmup_steps:
#         if lr_scheduler == "constant":
#             log.info("Can't use LR warmup with LR Scheduler constant... ignoring...")
#         else:
#             run_cmd += f' --lr_warmup_steps="{lr_warmup_steps}"'

#     train_batch_size = kwargs.get("train_batch_size", "")
#     if train_batch_size:
#         run_cmd += f' --train_batch_size="{train_batch_size}"'

#     max_train_steps = kwargs.get("max_train_steps", "")
#     if max_train_steps:
#         run_cmd += f' --max_train_steps="{max_train_steps}"'

#     save_every_n_epochs = kwargs.get("save_every_n_epochs")
#     if save_every_n_epochs:
#         run_cmd += f' --save_every_n_epochs="{int(save_every_n_epochs)}"'

#     mixed_precision = kwargs.get("mixed_precision", "")
#     if mixed_precision:
#         run_cmd += f' --mixed_precision="{mixed_precision}"'

#     save_precision = kwargs.get("save_precision", "")
#     if save_precision:
#         run_cmd += f' --save_precision="{save_precision}"'

#     seed = kwargs.get("seed", "")
#     if seed != "":
#         run_cmd += f' --seed="{seed}"'

#     caption_extension = kwargs.get("caption_extension", "")
#     if caption_extension:
#         run_cmd += f' --caption_extension="{caption_extension}"'

#     cache_latents = kwargs.get("cache_latents")
#     if cache_latents:
#         run_cmd += " --cache_latents"

#     cache_latents_to_disk = kwargs.get("cache_latents_to_disk")
#     if cache_latents_to_disk:
#         run_cmd += " --cache_latents_to_disk"

#     optimizer_type = kwargs.get("optimizer", "AdamW")
#     run_cmd += f' --optimizer_type="{optimizer_type}"'

#     optimizer_args = kwargs.get("optimizer_args", "")
#     if optimizer_args != "":
#         run_cmd += f" --optimizer_args {optimizer_args}"

#     lr_scheduler_args = kwargs.get("lr_scheduler_args", "")
#     if lr_scheduler_args != "":
#         run_cmd += f" --lr_scheduler_args {lr_scheduler_args}"

#     max_grad_norm = kwargs.get("max_grad_norm", "")
#     if max_grad_norm != "":
#         run_cmd += f' --max_grad_norm="{max_grad_norm}"'

#     return run_cmd


def run_cmd_advanced_training(**kwargs):
    run_cmd = ""

    additional_parameters = kwargs.get("additional_parameters")
    if additional_parameters:
        run_cmd += f" {additional_parameters}"

    block_lr = kwargs.get("block_lr")
    if block_lr:
        run_cmd += f' --block_lr="(block_lr)"'

    bucket_no_upscale = kwargs.get("bucket_no_upscale")
    if bucket_no_upscale:
        run_cmd += " --bucket_no_upscale"

    bucket_reso_steps = kwargs.get("bucket_reso_steps")
    if bucket_reso_steps:
        run_cmd += f" --bucket_reso_steps={int(bucket_reso_steps)}"

    cache_latents = kwargs.get("cache_latents")
    if cache_latents:
        run_cmd += " --cache_latents"

    cache_latents_to_disk = kwargs.get("cache_latents_to_disk")
    if cache_latents_to_disk:
        run_cmd += " --cache_latents_to_disk"

    cache_text_encoder_outputs = kwargs.get("cache_text_encoder_outputs")
    if cache_text_encoder_outputs:
        run_cmd += " --cache_text_encoder_outputs"

    caption_dropout_every_n_epochs = kwargs.get("caption_dropout_every_n_epochs")
    if caption_dropout_every_n_epochs and int(caption_dropout_every_n_epochs) > 0:
        run_cmd += (
            f' --caption_dropout_every_n_epochs="{int(caption_dropout_every_n_epochs)}"'
        )

    caption_dropout_rate = kwargs.get("caption_dropout_rate")
    if caption_dropout_rate and float(caption_dropout_rate) > 0:
        run_cmd += f' --caption_dropout_rate="{caption_dropout_rate}"'

    caption_extension = kwargs.get("caption_extension")
    if caption_extension:
        run_cmd += f' --caption_extension="{caption_extension}"'

    clip_skip = kwargs.get("clip_skip")
    if clip_skip and int(clip_skip) > 1:
        run_cmd += f" --clip_skip={int(clip_skip)}"

    color_aug = kwargs.get("color_aug")
    if color_aug:
        run_cmd += " --color_aug"

    dataset_repeats = kwargs.get("dataset_repeats")
    if dataset_repeats:
        run_cmd += f' --dataset_repeats="{dataset_repeats}"'

    debiased_estimation_loss = kwargs.get("debiased_estimation_loss")
    if debiased_estimation_loss:
        run_cmd += " --debiased_estimation_loss"

    dim_from_weights = kwargs.get("dim_from_weights")
    if dim_from_weights and kwargs.get(
        "lora_network_weights"
    ):  # Only if lora_network_weights is true
        run_cmd += f" --dim_from_weights"

    enable_bucket = kwargs.get("enable_bucket")
    if enable_bucket:
        min_bucket_reso = kwargs.get("min_bucket_reso")
        max_bucket_reso = kwargs.get("max_bucket_reso")
        if min_bucket_reso and max_bucket_reso:
            run_cmd += f" --enable_bucket --min_bucket_reso={min_bucket_reso} --max_bucket_reso={max_bucket_reso}"

    in_json = kwargs.get("in_json")
    if in_json:
        run_cmd += f' --in_json="{in_json}"'

    flip_aug = kwargs.get("flip_aug")
    if flip_aug:
        run_cmd += " --flip_aug"

    fp8_base = kwargs.get("fp8_base")
    if fp8_base:
        run_cmd += " --fp8_base"

    full_bf16 = kwargs.get("full_bf16")
    if full_bf16:
        run_cmd += " --full_bf16"

    full_fp16 = kwargs.get("full_fp16")
    if full_fp16:
        run_cmd += " --full_fp16"

    gradient_accumulation_steps = kwargs.get("gradient_accumulation_steps")
    if gradient_accumulation_steps and int(gradient_accumulation_steps) > 1:
        run_cmd += f" --gradient_accumulation_steps={int(gradient_accumulation_steps)}"

    gradient_checkpointing = kwargs.get("gradient_checkpointing")
    if gradient_checkpointing:
        run_cmd += " --gradient_checkpointing"

    keep_tokens = kwargs.get("keep_tokens")
    if keep_tokens and int(keep_tokens) > 0:
        run_cmd += f' --keep_tokens="{int(keep_tokens)}"'

    learning_rate = kwargs.get("learning_rate")
    if learning_rate:
        run_cmd += f' --learning_rate="{learning_rate}"'

    learning_rate_te = kwargs.get("learning_rate_te")
    if learning_rate_te:
        run_cmd += f' --learning_rate_te="{learning_rate_te}"'

    learning_rate_te1 = kwargs.get("learning_rate_te1")
    if learning_rate_te1:
        run_cmd += f' --learning_rate_te1="{learning_rate_te1}"'

    learning_rate_te2 = kwargs.get("learning_rate_te2")
    if learning_rate_te2:
        run_cmd += f' --learning_rate_te2="{learning_rate_te2}"'

    logging_dir = kwargs.get("logging_dir")
    if logging_dir:
        run_cmd += f' --logging_dir="{logging_dir}"'

    lora_network_weights = kwargs.get("lora_network_weights")
    if lora_network_weights:
        run_cmd += f' --lora_network_weights="{lora_network_weights}"'

    lr_scheduler = kwargs.get("lr_scheduler")
    if lr_scheduler:
        run_cmd += f' --lr_scheduler="{lr_scheduler}"'

    lr_scheduler_args = kwargs.get("lr_scheduler_args")
    if lr_scheduler_args and lr_scheduler_args != "":
        run_cmd += f" --lr_scheduler_args {lr_scheduler_args}"

    lr_scheduler_num_cycles = kwargs.get("lr_scheduler_num_cycles")
    if lr_scheduler_num_cycles and not lr_scheduler_num_cycles == "":
        run_cmd += f' --lr_scheduler_num_cycles="{lr_scheduler_num_cycles}"'
    else:
        epoch = kwargs.get("epoch")
        if epoch:
            run_cmd += f' --lr_scheduler_num_cycles="{epoch}"'

    lr_scheduler_power = kwargs.get("lr_scheduler_power")
    if lr_scheduler_power and not lr_scheduler_power == "":
        run_cmd += f' --lr_scheduler_power="{lr_scheduler_power}"'

    lr_warmup_steps = kwargs.get("lr_warmup_steps")
    if lr_warmup_steps:
        if lr_scheduler == "constant":
            log.info("Can't use LR warmup with LR Scheduler constant... ignoring...")
        else:
            run_cmd += f' --lr_warmup_steps="{lr_warmup_steps}"'

    gpu_ids = kwargs.get("gpu_ids")
    if gpu_ids:
        run_cmd += f' --gpu_ids="{gpu_ids}"'

    max_data_loader_n_workers = kwargs.get("max_data_loader_n_workers")
    if max_data_loader_n_workers and not max_data_loader_n_workers == "":
        run_cmd += f' --max_data_loader_n_workers="{max_data_loader_n_workers}"'

    max_grad_norm = kwargs.get("max_grad_norm")
    if max_grad_norm and max_grad_norm != "":
        run_cmd += f' --max_grad_norm="{max_grad_norm}"'

    max_resolution = kwargs.get("max_resolution")
    if max_resolution:
        run_cmd += f' --resolution="{max_resolution}"'

    max_timestep = kwargs.get("max_timestep")
    if max_timestep and int(max_timestep) < 1000:
        run_cmd += f" --max_timestep={int(max_timestep)}"

    max_token_length = kwargs.get("max_token_length")
    if max_token_length and int(max_token_length) > 75:
        run_cmd += f" --max_token_length={int(max_token_length)}"

    max_train_epochs = kwargs.get("max_train_epochs")
    if max_train_epochs and not max_train_epochs == "":
        run_cmd += f" --max_train_epochs={max_train_epochs}"

    max_train_steps = kwargs.get("max_train_steps")
    if max_train_steps:
        run_cmd += f' --max_train_steps="{max_train_steps}"'

    mem_eff_attn = kwargs.get("mem_eff_attn")
    if mem_eff_attn:
        run_cmd += " --mem_eff_attn"

    min_snr_gamma = kwargs.get("min_snr_gamma")
    if min_snr_gamma and int(min_snr_gamma) >= 1:
        run_cmd += f" --min_snr_gamma={int(min_snr_gamma)}"

    min_timestep = kwargs.get("min_timestep")
    if min_timestep and int(min_timestep) > 0:
        run_cmd += f" --min_timestep={int(min_timestep)}"

    mixed_precision = kwargs.get("mixed_precision")
    if mixed_precision:
        run_cmd += f' --mixed_precision="{mixed_precision}"'

    multi_gpu = kwargs.get("multi_gpu")
    if multi_gpu:
        run_cmd += " --multi_gpu"

    network_alpha = kwargs.get("network_alpha")
    if network_alpha:
        run_cmd += f' --network_alpha="{network_alpha}"'

    network_args = kwargs.get("network_args")
    if network_args and len(network_args):
        run_cmd += f" --network_args{network_args}"

    network_dim = kwargs.get("network_dim")
    if network_dim:
        run_cmd += f" --network_dim={network_dim}"

    network_dropout = kwargs.get("network_dropout")
    if network_dropout and network_dropout > 0.0:
        run_cmd += f" --network_dropout={network_dropout}"

    network_module = kwargs.get("network_module")
    if network_module:
        run_cmd += f" --network_module={network_module}"

    network_train_text_encoder_only = kwargs.get("network_train_text_encoder_only")
    if network_train_text_encoder_only:
        run_cmd += " --network_train_text_encoder_only"

    network_train_unet_only = kwargs.get("network_train_unet_only")
    if network_train_unet_only:
        run_cmd += " --network_train_unet_only"

    no_half_vae = kwargs.get("no_half_vae")
    if no_half_vae:
        run_cmd += " --no_half_vae"

    no_token_padding = kwargs.get("no_token_padding")
    if no_token_padding:
        run_cmd += " --no_token_padding"

    noise_offset_type = kwargs.get("noise_offset_type")
    if noise_offset_type and noise_offset_type == "Original":
        noise_offset = kwargs.get("noise_offset")
        if noise_offset and float(noise_offset) > 0:
            run_cmd += f" --noise_offset={float(noise_offset)}"

        adaptive_noise_scale = kwargs.get("adaptive_noise_scale")
        if (
            adaptive_noise_scale
            and float(adaptive_noise_scale) != 0
            and float(noise_offset) > 0
        ):
            run_cmd += f" --adaptive_noise_scale={float(adaptive_noise_scale)}"
    elif noise_offset_type and noise_offset_type == "Multires":
        multires_noise_iterations = kwargs.get("multires_noise_iterations")
        if int(multires_noise_iterations) > 0:
            run_cmd += (
                f' --multires_noise_iterations="{int(multires_noise_iterations)}"'
            )

        multires_noise_discount = kwargs.get("multires_noise_discount")
        if multires_noise_discount and float(multires_noise_discount) > 0:
            run_cmd += f' --multires_noise_discount="{float(multires_noise_discount)}"'

    num_machines = kwargs.get("num_machines")
    if num_machines and int(num_machines) > 1:
        run_cmd += f" --num_machines={int(num_machines)}"

    num_processes = kwargs.get("num_processes")
    if num_processes and int(num_processes) > 1:
        run_cmd += f" --num_processes={int(num_processes)}"

    num_cpu_threads_per_process = kwargs.get("num_cpu_threads_per_process")
    if num_cpu_threads_per_process and int(num_cpu_threads_per_process) > 1:
        run_cmd += f" --num_cpu_threads_per_process={int(num_cpu_threads_per_process)}"

    optimizer_args = kwargs.get("optimizer_args")
    if optimizer_args and optimizer_args != "":
        run_cmd += f" --optimizer_args {optimizer_args}"

    optimizer_type = kwargs.get("optimizer")
    if optimizer_type:
        run_cmd += f' --optimizer_type="{optimizer_type}"'

    output_dir = kwargs.get("output_dir")
    if output_dir:
        run_cmd += f' --output_dir="{output_dir}"'

    output_name = kwargs.get("output_name")
    if output_name and not output_name == "":
        run_cmd += f' --output_name="{output_name}"'

    persistent_data_loader_workers = kwargs.get("persistent_data_loader_workers")
    if persistent_data_loader_workers:
        run_cmd += " --persistent_data_loader_workers"

    pretrained_model_name_or_path = kwargs.get("pretrained_model_name_or_path")
    if pretrained_model_name_or_path:
        run_cmd += f' --pretrained_model_name_or_path="{pretrained_model_name_or_path}"'

    prior_loss_weight = kwargs.get("prior_loss_weight")
    if prior_loss_weight and not float(prior_loss_weight) == 1.0:
        run_cmd += f" --prior_loss_weight={prior_loss_weight}"

    random_crop = kwargs.get("random_crop")
    if random_crop:
        run_cmd += " --random_crop"

    reg_data_dir = kwargs.get("reg_data_dir")
    if reg_data_dir and len(reg_data_dir):
        run_cmd += f' --reg_data_dir="{reg_data_dir}"'

    resume = kwargs.get("resume")
    if resume:
        run_cmd += f' --resume="{resume}"'

    save_every_n_epochs = kwargs.get("save_every_n_epochs")
    if save_every_n_epochs:
        run_cmd += f' --save_every_n_epochs="{int(save_every_n_epochs)}"'

    save_every_n_steps = kwargs.get("save_every_n_steps")
    if save_every_n_steps and int(save_every_n_steps) > 0:
        run_cmd += f' --save_every_n_steps="{int(save_every_n_steps)}"'

    save_last_n_steps = kwargs.get("save_last_n_steps")
    if save_last_n_steps and int(save_last_n_steps) > 0:
        run_cmd += f' --save_last_n_steps="{int(save_last_n_steps)}"'

    save_last_n_steps_state = kwargs.get("save_last_n_steps_state")
    if save_last_n_steps_state and int(save_last_n_steps_state) > 0:
        run_cmd += f' --save_last_n_steps_state="{int(save_last_n_steps_state)}"'

    save_model_as = kwargs.get("save_model_as")
    if save_model_as and not save_model_as == "same as source model":
        run_cmd += f" --save_model_as={save_model_as}"

    save_precision = kwargs.get("save_precision")
    if save_precision:
        run_cmd += f' --save_precision="{save_precision}"'

    save_state = kwargs.get("save_state")
    if save_state:
        run_cmd += " --save_state"

    scale_v_pred_loss_like_noise_pred = kwargs.get("scale_v_pred_loss_like_noise_pred")
    if scale_v_pred_loss_like_noise_pred:
        run_cmd += " --scale_v_pred_loss_like_noise_pred"

    scale_weight_norms = kwargs.get("scale_weight_norms")
    if scale_weight_norms and scale_weight_norms > 0.0:
        run_cmd += f' --scale_weight_norms="{scale_weight_norms}"'

    seed = kwargs.get("seed")
    if seed and seed != "":
        run_cmd += f' --seed="{seed}"'

    shuffle_caption = kwargs.get("shuffle_caption")
    if shuffle_caption:
        run_cmd += " --shuffle_caption"

    stop_text_encoder_training = kwargs.get("stop_text_encoder_training")
    if stop_text_encoder_training and stop_text_encoder_training > 0:
        run_cmd += f' --stop_text_encoder_training="{stop_text_encoder_training}"'

    text_encoder_lr = kwargs.get("text_encoder_lr")
    if text_encoder_lr and (float(text_encoder_lr) > 0):
        run_cmd += f" --text_encoder_lr={text_encoder_lr}"

    train_batch_size = kwargs.get("train_batch_size")
    if train_batch_size:
        run_cmd += f' --train_batch_size="{train_batch_size}"'

    training_comment = kwargs.get("training_comment")
    if training_comment and len(training_comment):
        run_cmd += f' --training_comment="{training_comment}"'

    train_data_dir = kwargs.get("train_data_dir")
    if train_data_dir:
        run_cmd += f' --train_data_dir="{train_data_dir}"'

    train_text_encoder = kwargs.get("train_text_encoder")
    if train_text_encoder:
        run_cmd += " --train_text_encoder"

    unet_lr = kwargs.get("unet_lr")
    if unet_lr and (float(unet_lr) > 0):
        run_cmd += f" --unet_lr={unet_lr}"

    use_wandb = kwargs.get("use_wandb")
    if use_wandb:
        run_cmd += " --log_with wandb"

    v_parameterization = kwargs.get("v_parameterization")
    if v_parameterization:
        run_cmd += " --v_parameterization"

    v_pred_like_loss = kwargs.get("v_pred_like_loss")
    if v_pred_like_loss and float(v_pred_like_loss) > 0:
        run_cmd += f' --v_pred_like_loss="{float(v_pred_like_loss)}"'

    v2 = kwargs.get("v2")
    if v2:
        run_cmd += " --v2"

    vae = kwargs.get("vae")
    if vae and not vae == "":
        run_cmd += f' --vae="{vae}"'

    vae_batch_size = kwargs.get("vae_batch_size")
    if vae_batch_size and int(vae_batch_size) > 0:
        run_cmd += f' --vae_batch_size="{int(vae_batch_size)}"'

    wandb_api_key = kwargs.get("wandb_api_key")
    if wandb_api_key:
        run_cmd += f' --wandb_api_key="{wandb_api_key}"'

    weighted_captions = kwargs.get("weighted_captions")
    if weighted_captions:
        run_cmd += " --weighted_captions"

    xformers = kwargs.get("xformers")
    if xformers and xformers == "xformers":
        run_cmd += " --xformers"
    elif xformers and xformers == "sdpa":
        run_cmd += " --sdpa"

    return run_cmd


def verify_image_folder_pattern(folder_path):
    false_response = True  # temporarily set to true to prevent stopping training in case of false positive
    true_response = True

    # Check if the folder exists
    if not os.path.isdir(folder_path):
        log.error(
            f"The provided path '{folder_path}' is not a valid folder. Please follow the folder structure documentation found at docs\image_folder_structure.md ..."
        )
        return false_response

    # Create a regular expression pattern to match the required sub-folder names
    # The pattern should start with one or more digits (\d+) followed by an underscore (_)
    # After the underscore, it should match one or more word characters (\w+), which can be letters, numbers, or underscores
    # Example of a valid pattern matching name: 123_example_folder
    pattern = r"^\d+_\w+"

    # Get the list of sub-folders in the directory
    subfolders = [
        os.path.join(folder_path, subfolder)
        for subfolder in os.listdir(folder_path)
        if os.path.isdir(os.path.join(folder_path, subfolder))
    ]

    # Check the pattern of each sub-folder
    matching_subfolders = [
        subfolder
        for subfolder in subfolders
        if re.match(pattern, os.path.basename(subfolder))
    ]

    # Print non-matching sub-folders
    non_matching_subfolders = set(subfolders) - set(matching_subfolders)
    if non_matching_subfolders:
        log.error(
            f"The following folders do not match the required pattern <number>_<text>: {', '.join(non_matching_subfolders)}"
        )
        log.error(
            f"Please follow the folder structure documentation found at docs\image_folder_structure.md ..."
        )
        return false_response

    # Check if no sub-folders exist
    if not matching_subfolders:
        log.error(
            f"No image folders found in {folder_path}. Please follow the folder structure documentation found at docs\image_folder_structure.md ..."
        )
        return false_response

    log.info(f"Valid image folder names found in: {folder_path}")
    return true_response


def SaveConfigFile(
    parameters,
    file_path: str,
    exclusion=["file_path", "save_as", "headless", "print_only"],
):
    # Return the values of the variables as a dictionary
    variables = {
        name: value
        for name, value in sorted(parameters, key=lambda x: x[0])
        if name not in exclusion
    }

    # Save the data to the selected file
    with open(file_path, "w") as file:
        json.dump(variables, file, indent=2)


def save_to_file(content):
    logs_directory = "logs"
    file_path = os.path.join(logs_directory, "print_command.txt")

    try:
        # Create the 'logs' directory if it does not exist
        if not os.path.exists(logs_directory):
            os.makedirs(logs_directory)

        with open(file_path, "a") as file:
            file.write(content + "\n")
    except IOError as e:
        print(f"Error: Could not write to file - {e}")
    except OSError as e:
        print(f"Error: Could not create 'logs' directory - {e}")


def check_duplicate_filenames(
    folder_path, image_extension=[".gif", ".png", ".jpg", ".jpeg", ".webp"]
):
    log.info("Checking for duplicate image filenames in training data directory...")
    for root, dirs, files in os.walk(folder_path):
        filenames = {}
        for file in files:
            filename, extension = os.path.splitext(file)
            if extension.lower() in image_extension:
                full_path = os.path.join(root, file)
                if filename in filenames:
                    existing_path = filenames[filename]
                    if existing_path != full_path:
                        print(
                            f"Warning: Same filename '{filename}' with different image extension found. This will cause training issues. Rename one of the file."
                        )
                        print(f"Existing file: {existing_path}")
                        print(f"Current file: {full_path}")
                else:
                    filenames[filename] = full_path


def is_file_writable(file_path):
    if not os.path.exists(file_path):
        # print(f"File '{file_path}' does not exist.")
        return True

    try:
        log.warning(f"File '{file_path}' already exist... it will be overwritten...")
        # Check if the file can be opened in write mode (which implies it's not open by another process)
        with open(file_path, "a"):
            pass
        return True
    except IOError:
        log.warning(f"File '{file_path}' can't be written to...")
        return False