lora_train_popup.py

import gc
import json
import time
from functools import partial
from typing import Union
import os
import tkinter as tk
from tkinter import filedialog as fd, ttk
from tkinter import simpledialog as sd
from tkinter import messagebox as mb

import torch.cuda
import train_network
import library.train_util as util
import argparse


class ArgStore:
    # Represents the entirety of all possible inputs for sd-scripts. they are ordered from most important to least
    def __init__(self):
        # Important, these are the most likely things you will modify
        self.base_model: str = r""  # example path, r"E:\sd\stable-diffusion-webui\models\Stable-diffusion\nai.ckpt"
        self.img_folder: str = r""  # is the folder path to your img folder, make sure to follow the guide here for folder setup: https://rentry.org/2chAI_LoRA_Dreambooth_guide_english#for-kohyas-script
        self.output_folder: str = r""  # just the folder all epochs/safetensors are output
        self.change_output_name: Union[str, None] = None  # changes the output name of the epochs
        self.save_json_folder: Union[str, None] = None  # OPTIONAL, saves a json folder of your config to whatever location you set here.
        self.load_json_path: Union[str, None] = None  # OPTIONAL, loads a json file partially changes the config to match. things like folder paths do not get modified.
        self.json_load_skip_list: Union[list[str], None] = ["save_json_folder", "reg_img_folder",
                                                            "lora_model_for_resume", "change_output_name",
                                                            "training_comment",
                                                            "json_load_skip_list"]  # OPTIONAL, allows the user to define what they skip when loading a json, by default it loads everything, including all paths, set it up like this ["base_model", "img_folder", "output_folder"]
        self.caption_dropout_rate: Union[float, None] = None  # The rate at which captions for files get dropped.
        self.caption_dropout_every_n_epochs: Union[int, None] = None  # Defines how often an epoch will completely ignore
        # captions, EX. 3 means it will ignore captions at epochs 3, 6, and 9
        self.caption_tag_dropout_rate: Union[float, None] = None  # Defines the rate at which a tag would be dropped, rather than the entire caption file
        self.noise_offset: Union[float, None] = None  # OPTIONAL, seems to help allow SD to gen better blacks and whites
                                                      # Kohya recommends, if you have it set, to use 0.1, not sure how
                                                      # high the value can be, I'm going to assume maximum of 1

        self.net_dim: int = 128  # network dimension, 128 is the most common, however you might be able to get lesser to work
        self.alpha: float = 128  # represents the scalar for training. the lower the alpha, the less gets learned per step. if you want the older way of training, set this to dim
        # list of schedulers: linear, cosine, cosine_with_restarts, polynomial, constant, constant_with_warmup
        self.scheduler: str = "cosine_with_restarts"  # the scheduler for learning rate. Each does something specific
        self.cosine_restarts: Union[int, None] = 1  # OPTIONAL, represents the number of times it restarts. Only matters if you are using cosine_with_restarts
        self.scheduler_power: Union[float, None] = 1  # OPTIONAL, represents the power of the polynomial. Only matters if you are using polynomial
        self.warmup_lr_ratio: Union[float, None] = None  # OPTIONAL, Calculates the number of warmup steps based on the ratio given. Make sure to set this if you are using constant_with_warmup, None to ignore
        self.learning_rate: Union[float, None] = 1e-4  # OPTIONAL, when not set, lr gets set to 1e-3 as per adamW. Personally, I suggest actually setting this as lower lr seems to be a small bit better.
        self.text_encoder_lr: Union[float, None] = None  # OPTIONAL, Sets a specific lr for the text encoder, this overwrites the base lr I believe, None to ignore
        self.unet_lr: Union[float, None] = None  # OPTIONAL, Sets a specific lr for the unet, this overwrites the base lr I believe, None to ignore
        self.num_workers: int = 1  # The number of threads that are being used to load images, lower speeds up the start of epochs, but slows down the loading of data. The assumption here is that it increases the training time as you reduce this value
        self.persistent_workers: bool = True  # makes workers persistent, further reduces/eliminates the lag in between epochs. however it may increase memory usage

        self.batch_size: int = 1  # The number of images that get processed at one time, this is directly proportional to your vram and resolution. with 12gb of vram, at 512 reso, you can get a maximum of 6 batch size
        self.num_epochs: int = 1  # The number of epochs, if you set max steps this value is ignored as it doesn't calculate steps.
        self.save_every_n_epochs: Union[int, None] = None  # OPTIONAL, how often to save epochs, None to ignore
        self.shuffle_captions: bool = False  # OPTIONAL, False to ignore
        self.keep_tokens: Union[int, None] = None  # OPTIONAL, None to ignore
        self.max_steps: Union[int, None] = None  # OPTIONAL, if you have specific steps you want to hit, this allows you to set it directly. None to ignore
        self.tag_occurrence_txt_file: bool = False  # OPTIONAL, creates a txt file that has the entire occurrence of all tags in your dataset
                                                    # the metadata will also have this so long as you have metadata on, so no reason to have this on by default
                                                    # will automatically output to the same folder as your output checkpoints
        self.sort_tag_occurrence_alphabetically: bool = False  # OPTIONAL, only applies if tag_occurrence_txt_file is also true
                                                               # Will change the output to be alphabetically vs being occurrence based

        # These are the second most likely things you will modify
        self.train_resolution: int = 512
        self.min_bucket_resolution: int = 320
        self.max_bucket_resolution: int = 960
        self.lora_model_for_resume: Union[str, None] = None  # OPTIONAL, takes an input lora to continue training from, not exactly the way it *should* be, but it works, None to ignore
        self.save_state: bool = False  # OPTIONAL, is the intended way to save a training state to use for continuing training, False to ignore
        self.load_previous_save_state: Union[str, None] = None  # OPTIONAL, is the intended way to load a training state to use for continuing training, None to ignore
        self.training_comment: Union[str, None] = None  # OPTIONAL, great way to put in things like activation tokens right into the metadata. seems to not work at this point and time
        self.unet_only: bool = False  # OPTIONAL, set it to only train the unet
        self.text_only: bool = False  # OPTIONAL, set it to only train the text encoder

        # These are the least likely things you will modify
        self.reg_img_folder: Union[str, None] = None  # OPTIONAL, None to ignore
        self.clip_skip: int = 2  # If you are training on a model that is anime based, keep this at 2 as most models are designed for that
        self.test_seed: int = 23  # this is the "reproducable seed", basically if you set the seed to this, you should be able to input a prompt from one of your training images and get a close representation of it
        self.prior_loss_weight: float = 1  # is the loss weight much like Dreambooth, is required for LoRA training
        self.gradient_checkpointing: bool = False  # OPTIONAL, enables gradient checkpointing
        self.gradient_acc_steps: Union[int, None] = None  # OPTIONAL, not sure exactly what this means
        self.mixed_precision: str = "fp16"  # If you have the ability to use bf16, do it, it's better
        self.save_precision: str = "fp16"  # You can also save in bf16, but because it's not universally supported, I suggest you keep saving at fp16
        self.save_as: str = "safetensors"  # list is pt, ckpt, safetensors
        self.caption_extension: str = ".txt"  # the other option is .captions, but since wd1.4 tagger outputs as txt files, this is the default
        self.max_clip_token_length = 150  # can be 75, 150, or 225 I believe, there is no reason to go higher than 150 though
        self.buckets: bool = True
        self.xformers: bool = True
        self.use_8bit_adam: bool = True
        self.cache_latents: bool = True
        self.color_aug: bool = False  # IMPORTANT: Clashes with cache_latents, only have one of the two on!
        self.flip_aug: bool = False
        self.vae: Union[str, None] = None  # Seems to only make results worse when not using that specific vae, should probably not use
        self.no_meta: bool = False  # This removes the metadata that now gets saved into safetensors, (you should keep this on)
        self.log_dir: Union[str, None] = None  # output of logs, not useful to most people.
        self.v2: bool = False  # Sets up training for SD2.1
        self.v_parameterization: bool = False  # Only is used when v2 is also set and you are using the 768x version of v2

    # Creates the dict that is used for the rest of the code, to facilitate easier json saving and loading
    @staticmethod
    def convert_args_to_dict():
        return ArgStore().__dict__


def main():
    parser = argparse.ArgumentParser()
    setup_args(parser)
    pre_args = parser.parse_args()
    queues = 0
    args_queue = []
    cont = True
    while cont:
        arg_dict = ArgStore.convert_args_to_dict()
        ret = mb.askyesno(message="Do you want to load a json config file?")
        if ret:
            load_json(ask_file("select json to load from", {"json"}), arg_dict)
            arg_dict = ask_elements_trunc(arg_dict)
        else:
            arg_dict = ask_elements(arg_dict)
        if pre_args.save_json_path or arg_dict["save_json_folder"]:
            save_json(pre_args.save_json_path if pre_args.save_json_path else arg_dict['save_json_folder'], arg_dict)
        args = create_arg_space(arg_dict)
        args = parser.parse_args(args)
        queues += 1
        args_queue.append(args)
        if arg_dict['tag_occurrence_txt_file']:
            get_occurrence_of_tags(arg_dict)
        ret = mb.askyesno(message="Do you want to queue another training?")
        if not ret:
            cont = False
    for args in args_queue:
        try:
            train_network.train(args)
        except Exception as e:
            print(f"Failed to train this set of args.\nSkipping this training session.\nError is: {e}")
        gc.collect()
        torch.cuda.empty_cache()


def create_arg_space(args: dict) -> [str]:
    # This is the list of args that are to be used regardless of setup
    output = ["--network_module=networks.lora", f"--pretrained_model_name_or_path={args['base_model']}",
              f"--train_data_dir={args['img_folder']}", f"--output_dir={args['output_folder']}",
              f"--prior_loss_weight={args['prior_loss_weight']}", f"--caption_extension=" + args['caption_extension'],
              f"--resolution={args['train_resolution']}", f"--train_batch_size={args['batch_size']}",
              f"--mixed_precision={args['mixed_precision']}", f"--save_precision={args['save_precision']}",
              f"--network_dim={args['net_dim']}", f"--save_model_as={args['save_as']}",
              f"--clip_skip={args['clip_skip']}", f"--seed={args['test_seed']}",
              f"--max_token_length={args['max_clip_token_length']}", f"--lr_scheduler={args['scheduler']}",
              f"--network_alpha={args['alpha']}", f"--max_data_loader_n_workers={args['num_workers']}"]
    if not args['max_steps']:
        output.append(f"--max_train_epochs={args['num_epochs']}")
        output += create_optional_args(args, find_max_steps(args))
    else:
        output.append(f"--max_train_steps={args['max_steps']}")
        output += create_optional_args(args, args['max_steps'])
    return output


def create_optional_args(args: dict, steps):
    output = []
    if args["reg_img_folder"]:
        output.append(f"--reg_data_dir={args['reg_img_folder']}")

    if args['lora_model_for_resume']:
        output.append(f"--network_weights={args['lora_model_for_resume']}")

    if args['save_every_n_epochs']:
        output.append(f"--save_every_n_epochs={args['save_every_n_epochs']}")
    else:
        output.append("--save_every_n_epochs=999999")

    if args['shuffle_captions']:
        output.append("--shuffle_caption")

    if args['keep_tokens'] and args['keep_tokens'] > 0:
        output.append(f"--keep_tokens={args['keep_tokens']}")

    if args['buckets']:
        output.append("--enable_bucket")
        output.append(f"--min_bucket_reso={args['min_bucket_resolution']}")
        output.append(f"--max_bucket_reso={args['max_bucket_resolution']}")

    if args['use_8bit_adam']:
        output.append("--use_8bit_adam")

    if args['xformers']:
        output.append("--xformers")

    if args['color_aug']:
        if args['cache_latents']:
            print("color_aug and cache_latents conflict with one another. Please select only one")
            quit(1)
        output.append("--color_aug")

    if args['flip_aug']:
        output.append("--flip_aug")

    if args['cache_latents']:
        output.append("--cache_latents")

    if args['warmup_lr_ratio'] and args['warmup_lr_ratio'] > 0:
        warmup_steps = int(steps * args['warmup_lr_ratio'])
        output.append(f"--lr_warmup_steps={warmup_steps}")

    if args['gradient_checkpointing']:
        output.append("--gradient_checkpointing")

    if args['gradient_acc_steps'] and args['gradient_acc_steps'] > 0 and args['gradient_checkpointing']:
        output.append(f"--gradient_accumulation_steps={args['gradient_acc_steps']}")

    if args['learning_rate'] and args['learning_rate'] > 0:
        output.append(f"--learning_rate={args['learning_rate']}")

    if args['text_encoder_lr'] and args['text_encoder_lr'] > 0:
        output.append(f"--text_encoder_lr={args['text_encoder_lr']}")

    if args['unet_lr'] and args['unet_lr'] > 0:
        output.append(f"--unet_lr={args['unet_lr']}")

    if args['vae']:
        output.append(f"--vae={args['vae']}")

    if args['no_meta']:
        output.append("--no_metadata")

    if args['save_state']:
        output.append("--save_state")

    if args['load_previous_save_state']:
        output.append(f"--resume={args['load_previous_save_state']}")

    if args['change_output_name']:
        output.append(f"--output_name={args['change_output_name']}")

    if args['training_comment']:
        output.append(f"--training_comment={args['training_comment']}")

    if args['cosine_restarts'] and args['scheduler'] == "cosine_with_restarts":
        output.append(f"--lr_scheduler_num_cycles={args['cosine_restarts']}")

    if args['scheduler_power'] and args['scheduler'] == "polynomial":
        output.append(f"--lr_scheduler_power={args['scheduler_power']}")

    if args['persistent_workers']:
        output.append(f"--persistent_data_loader_workers")

    if args['unet_only']:
        output.append("--network_train_unet_only")

    if args['text_only'] and not args['unet_only']:
        output.append("--network_train_text_encoder_only")

    if args["log_dir"]:
        output.append(f"--logging_dir={args['log_dir']}")

    if args['caption_dropout_rate']:
        output.append(f"--caption_dropout_rate={args['caption_dropout_rate']}")

    if args['caption_dropout_every_n_epochs']:
        output.append(f"--caption_dropout_every_n_epochs={args['caption_dropout_every_n_epochs']}")

    if args['caption_tag_dropout_rate']:
        output.append(f"--caption_tag_dropout_rate={args['caption_tag_dropout_rate']}")

    if args['v2']:
        output.append("--v2")

    if args['v2'] and args['v_parameterization']:
        output.append("--v_parameterization")

    if args['noise_offset']:
        output.append(f"--noise_offset={args['noise_offset']}")
    return output


def find_max_steps(args: dict) -> int:
    total_steps = 0
    folders = os.listdir(args["img_folder"])
    for folder in folders:
        if not os.path.isdir(os.path.join(args["img_folder"], folder)):
            continue
        num_repeats = folder.split("_")
        if len(num_repeats) < 2:
            print(f"folder {folder} is not in the correct format. Format is x_name. skipping")
            continue
        try:
            num_repeats = int(num_repeats[0])
        except ValueError:
            print(f"folder {folder} is not in the correct format. Format is x_name. skipping")
            continue
        imgs = 0
        for file in os.listdir(os.path.join(args["img_folder"], folder)):
            if os.path.isdir(file):
                continue
            ext = file.split(".")
            if ext[-1].lower() in {"png", "bmp", "gif", "jpeg", "jpg", "webp"}:
                imgs += 1
        total_steps += (num_repeats * imgs)
    total_steps = int((total_steps / args["batch_size"]) * args["num_epochs"])
    return total_steps


def add_misc_args(parser):
    parser.add_argument("--save_json_path", type=str, default=None,
                        help="Path to save a configuration json file to")
    parser.add_argument("--load_json_path", type=str, default=None,
                        help="Path to a json file to configure things from")
    parser.add_argument("--no_metadata", action='store_true',
                        help="do not save metadata in output model / メタデータを出力先モデルに保存しない")
    parser.add_argument("--save_model_as", type=str, default="safetensors", choices=[None, "ckpt", "pt", "safetensors"],
                        help="format to save the model (default is .safetensors) / モデル保存時の形式（デフォルトはsafetensors）")

    parser.add_argument("--unet_lr", type=float, default=None, help="learning rate for U-Net / U-Netの学習率")
    parser.add_argument("--text_encoder_lr", type=float, default=None,
                        help="learning rate for Text Encoder / Text Encoderの学習率")
    parser.add_argument("--lr_scheduler_num_cycles", type=int, default=1,
                        help="Number of restarts for cosine scheduler with restarts / cosine with restartsスケジューラでのリスタート回数")
    parser.add_argument("--lr_scheduler_power", type=float, default=1,
                        help="Polynomial power for polynomial scheduler / polynomialスケジューラでのpolynomial power")

    parser.add_argument("--network_weights", type=str, default=None,
                        help="pretrained weights for network / 学習するネットワークの初期重み")
    parser.add_argument("--network_module", type=str, default=None,
                        help='network module to train / 学習対象のネットワークのモジュール')
    parser.add_argument("--network_dim", type=int, default=None,
                        help='network dimensions (depends on each network) / モジュールの次元数（ネットワークにより定義は異なります）')
    parser.add_argument("--network_alpha", type=float, default=1,
                        help='alpha for LoRA weight scaling, default 1 (same as network_dim for same behavior as old version) / LoRaの重み調整のalpha値、デフォルト1（旧バージョンと同じ動作をするにはnetwork_dimと同じ値を指定）')
    parser.add_argument("--network_args", type=str, default=None, nargs='*',
                        help='additional argmuments for network (key=value) / ネットワークへの追加の引数')
    parser.add_argument("--network_train_unet_only", action="store_true",
                        help="only training U-Net part / U-Net関連部分のみ学習する")
    parser.add_argument("--network_train_text_encoder_only", action="store_true",
                        help="only training Text Encoder part / Text Encoder関連部分のみ学習する")
    parser.add_argument("--training_comment", type=str, default=None,
                        help="arbitrary comment string stored in metadata / メタデータに記録する任意のコメント文字列")


def setup_args(parser):
    util.add_sd_models_arguments(parser)
    util.add_dataset_arguments(parser, True, True, True)
    util.add_training_arguments(parser, True)
    add_misc_args(parser)


def get_occurrence_of_tags(args):
    extension = args['caption_extension']
    img_folder = args['img_folder']
    output_folder = args['output_folder']
    occurrence_dict = {}
    print(img_folder)
    for folder in os.listdir(img_folder):
        print(folder)
        if not os.path.isdir(os.path.join(img_folder, folder)):
            continue
        for file in os.listdir(os.path.join(img_folder, folder)):
            if not os.path.isfile(os.path.join(img_folder, folder, file)):
                continue
            ext = os.path.splitext(file)[1]
            if ext != extension:
                continue
            get_tags_from_file(os.path.join(img_folder, folder, file), occurrence_dict)
    if not args['sort_tag_occurrence_alphabetically']:
        output_list = {k: v for k, v in sorted(occurrence_dict.items(), key=lambda item: item[1], reverse=True)}
    else:
        output_list = {k: v for k, v in sorted(occurrence_dict.items(), key=lambda item: item[0])}
    name = args['change_output_name'] if args['change_output_name'] else "last"
    with open(os.path.join(output_folder, f"{name}.txt"), "w") as f:
        f.write(f"Below is a list of keywords used during the training of {args['change_output_name']}:\n")
        for k, v in output_list.items():
            f.write(f"[{v}] {k}\n")
    print(f"Created a txt file named {name}.txt in the output folder")


def get_tags_from_file(file, occurrence_dict):
    f = open(file)
    temp = f.read().replace(", ", ",").split(",")
    f.close()
    for tag in temp:
        if tag in occurrence_dict:
            occurrence_dict[tag] += 1
        else:
            occurrence_dict[tag] = 1


def ask_file(message, accepted_ext_list, file_path=None):
    mb.showinfo(message=message)
    res = ""
    _initialdir = ""
    _initialfile = ""
    if file_path != None:
        _initialdir = os.path.dirname(file_path) if os.path.exists(file_path) else ""
        _initialfile = os.path.basename(file_path) if os.path.exists(file_path) else ""

    while res == "":
        res = fd.askopenfilename(title=message, initialdir=_initialdir, initialfile=_initialfile)
        if res == "" or type(res) == tuple:
            ret = mb.askretrycancel(message="Do you want to to cancel training?")
            if not ret:
                exit()
            continue
        elif not os.path.exists(res):
            res = ""
            continue
        _, name = os.path.split(res)
        split_name = name.split(".")
        if split_name[-1] not in accepted_ext_list:
            res = ""
    return res


def ask_dir(message, dir_path=None):
    mb.showinfo(message=message)
    res = ""
    _initialdir = ""
    if dir_path != None:
        _initialdir = dir_path if os.path.exists(dir_path) else ""
    while res == "":
        res = fd.askdirectory(title=message, initialdir=_initialdir)
        if res == "" or type(res) == tuple:
            ret = mb.askretrycancel(message="Do you want to to cancel training?")
            if not ret:
                exit()
            continue
        if not os.path.exists(res):
            res = ""
    return res


def ask_elements_trunc(args: dict):
    args['base_model'] = ask_file("Select your base model", {"ckpt", "safetensors"}, args['base_model'])
    args['img_folder'] = ask_dir("Select your image folder", args['img_folder'])
    args['output_folder'] = ask_dir("Select your output folder", args['output_folder'])

    ret = mb.askyesno(message="Do you want to save a json of your configuration?")
    if ret:
        args['save_json_folder'] = ask_dir("Select the folder to save json files to", args['save_json_folder'])
    else:
        args['save_json_folder'] = None

    ret = mb.askyesno(message="Are you training on a SD2 based model?")
    if ret:
        args['v2'] = True

    ret = mb.askyesno(message="Are you training on an realistic model?")
    if ret:
        args['clip_skip'] = 1

    if args['v2']:
        ret = mb.askyesno(message="Are you training on a model based on the 768x version of SD2?")
        if ret:
            args['v_parameterization'] = True

    ret = mb.askyesno(message="Do you want to use regularization images?")
    if ret:
        args['reg_img_folder'] = ask_dir("Select your regularization folder", args['reg_img_folder'])
    else:
        args['reg_img_folder'] = None

    ret = mb.askyesno(message="Do you want to continue from an earlier version?")
    if ret:
        args['lora_model_for_resume'] = ask_file("Select your lora model", {"ckpt", "pt", "safetensors"},
                                                 args['lora_model_for_resume'])
    else:
        args['lora_model_for_resume'] = None

    ret = mb.askyesno(message="Do you want to flip all of your images? It is supposed to reduce biases\n"
                              "within your dataset but it can also ruin learning an asymmetrical element\n")
    if ret:
        args['flip_aug'] = True

    ret = mb.askyesno(message="Do you want to change the name of output checkpoints?")
    if ret:
        ret = sd.askstring(title="output_name", prompt="What do you want your output name to be?\n"
                                                       "Cancel keeps outputs the original")
        if ret:
            args['change_output_name'] = ret
        else:
            args['change_output_name'] = None

    ret = sd.askstring(title="comment",
                       prompt="Do you want to set a comment that gets put into the metadata?\nA good use of this would "
                              "be to include how to use, such as activation keywords.\nCancel will leave empty")
    if ret is None:
        args['training_comment'] = ret
    else:
        args['training_comment'] = None

    ret = mb.askyesno(message="Do you want to train only one of unet and text encoder?")
    if ret:
        button = ButtonBox("Which do you want to train with?", ["unet_only", "text_only"])
        button.window.mainloop()
        if button.current_value != "":
            args[button.current_value] = True

    ret = mb.askyesno(message="Do you want to save a txt file that contains a list\n"
                              "of all tags that you have used in your training data?\n")
    if ret:
        args['tag_occurrence_txt_file'] = True
        button = ButtonBox("How do you want tags to be ordered?", ["alphabetically", "occurrence-ly"])
        button.window.mainloop()
        if button.current_value == "alphabetically":
            args['sort_tag_occurrence_alphabetically'] = True

    ret = mb.askyesno(message="Do you want to use caption dropout?")
    if ret:
        ret = mb.askyesno(message="Do you want full caption files to dropout randomly?")
        if ret:
            ret = sd.askinteger(title="Caption_File_Dropout",
                                prompt="How often do you want caption files to drop out?\n"
                                       "enter a number from 0 to 100 that is the percentage chance of dropout\n"
                                       "Cancel sets to 0")
            if ret and 0 <= ret <= 100:
                args['caption_dropout_rate'] = ret / 100.0

        ret = mb.askyesno(message="Do you want to have full epochs have no captions?")
        if ret:
            ret = sd.askinteger(title="Caption_epoch_dropout", prompt="The number set here is how often you will have an"
                                                                      "epoch with no captions\nSo if you set 3, then every"
                                                                      "three epochs will not have captions (3, 6, 9)\n"
                                                                      "Cancel will set to None")
            if ret:
                args['caption_dropout_every_n_epochs'] = ret

        ret = mb.askyesno(message="Do you want to have tags to randomly drop?")
        if ret:
            ret = sd.askinteger(title="Caption_tag_dropout", prompt="How often do you want tags to randomly drop out?\n"
                                                                    "Enter a number between 0 and 100, that is the percentage"
                                                                    "chance of dropout.\nCancel sets to 0")
            if ret and 0 <= ret <= 100:
                args['caption_tag_dropout_rate'] = ret / 100.0

    ret = mb.askyesno(message="Do you want to use noise offset? Noise offset seems to allow for SD to better generate\n"
                              "darker or lighter images using this than normal.")
    if ret:
        ret = sd.askfloat(title="noise_offset", prompt="What value do you want to set? recommended value is 0.1,\n"
                                                       "but it can go higher. Cancel defaults to 0.1")
        if ret:
            args['noise_offset'] = ret
        else:
            args['noise_offset'] = 0.1
    return args


def ask_elements(args: dict):
    # start with file dialog
    args['base_model'] = ask_file("Select your base model", {"ckpt", "safetensors"}, args['base_model'])
    args['img_folder'] = ask_dir("Select your image folder", args['img_folder'])
    args['output_folder'] = ask_dir("Select your output folder", args['output_folder'])

    # optional file dialog
    ret = mb.askyesno(message="Do you want to save a json of your configuration?")
    if ret:
        args['save_json_folder'] = ask_dir("Select the folder to save json files to", args['save_json_folder'])
    else:
        args['save_json_folder'] = None

    ret = mb.askyesno(message="Are you training on a SD2 based model?")
    if ret:
        args['v2'] = True

    ret = mb.askyesno(message="Are you training on an realistic model?")
    if ret:
        args['clip_skip'] = 1

    if args['v2']:
        ret = mb.askyesno(message="Are you training on a model based on the 768x version of SD2?")
        if ret:
            args['v_parameterization'] = True

    ret = mb.askyesno(message="Do you want to use regularization images?")
    if ret:
        args['reg_img_folder'] = ask_dir("Select your regularization folder", args['reg_img_folder'])
    else:
        args['reg_img_folder'] = None

    ret = mb.askyesno(message="Do you want to continue from an earlier version?")
    if ret:
        args['lora_model_for_resume'] = ask_file("Select your lora model", {"ckpt", "pt", "safetensors"},
                                                 args['lora_model_for_resume'])
    else:
        args['lora_model_for_resume'] = None

    ret = mb.askyesno(message="Do you want to flip all of your images? It is supposed to reduce biases\n"
                              "within your dataset but it can also ruin learning an asymmetrical element\n")
    if ret:
        args['flip_aug'] = True

    # text based required elements
    ret = sd.askinteger(title="batch_size",
                        prompt="The number of images that get processed at one time, this is directly proportional to "
                               "your vram and resolution. with 12gb of vram, at 512 reso, you can get a maximum of 6 "
                               "batch size\nHow large is your batch size going to be?\nCancel will default to 1")
    if ret is None:
        args['batch_size'] = 1
    else:
        args['batch_size'] = ret

    ret = sd.askinteger(title="num_epochs", prompt="How many epochs do you want?\nCancel will default to 1")
    if ret is None:
        args['num_epochs'] = 1
    else:
        args['num_epochs'] = ret

    ret = sd.askinteger(title="network_dim", prompt="What is the dim size you want to use?\nCancel will default to 128")
    if ret is None:
        args['net_dim'] = 128
    else:
        args['net_dim'] = ret

    ret = sd.askfloat(title="alpha", prompt="Alpha is the scalar of the training, generally a good starting point is "
                                            "0.5x dim size\nWhat Alpha do you want?\nCancel will default to equal to "
                                            "0.5 x network_dim")
    if ret is None:
        args['alpha'] = args['net_dim'] / 2
    else:
        args['alpha'] = ret

    ret = sd.askinteger(title="resolution", prompt="How large of a resolution do you want to train at?\n"
                                                   "Cancel will default to 512")
    if ret is None:
        args['train_resolution'] = 512
    else:
        args['train_resolution'] = ret

    ret = sd.askfloat(title="learning_rate", prompt="What learning rate do you want to use?\n"
                                                    "Cancel will default to 1e-4")
    if ret is None:
        args['learning_rate'] = 1e-4
    else:
        args['learning_rate'] = ret

    ret = sd.askfloat(title="text_encoder_lr", prompt="Do you want to set the text_encoder_lr?\n"
                                                      "Cancel will default to None")
    if ret is None:
        args['text_encoder_lr'] = None
    else:
        args['text_encoder_lr'] = ret

    ret = sd.askfloat(title="unet_lr", prompt="Do you want to set the unet_lr?\nCancel will default to None")
    if ret is None:
        args['unet_lr'] = None
    else:
        args['unet_lr'] = ret

    button = ButtonBox("Which scheduler do you want?", ["cosine_with_restarts", "cosine", "polynomial",
                                                        "constant", "constant_with_warmup", "linear"])
    button.window.mainloop()
    args['scheduler'] = button.current_value if button.current_value != "" else "cosine_with_restarts"

    if args['scheduler'] == "cosine_with_restarts":
        ret = sd.askinteger(title="Cycle Count",
                            prompt="How many times do you want cosine to restart?\nThis is the entire amount of times "
                                   "it will restart for the entire training\nCancel will default to 1")
        if ret is None:
            args['cosine_restarts'] = 1
        else:
            args['cosine_restarts'] = ret

    if args['scheduler'] == "polynomial":
        ret = sd.askfloat(title="Poly Strength",
                          prompt="What power do you want to set your polynomial to?\nhigher power means that the "
                                 "model reduces the learning more more aggressively from initial training.\n1 = "
                                 "linear\nCancel sets to 1")
        if ret is None:
            args['scheduler_power'] = 1
        else:
            args['scheduler_power'] = ret

    ret = mb.askyesno(message="Do you want to save epochs as it trains?")
    if ret:
        ret = sd.askinteger(title="save_epoch",
                            prompt="How often do you want to save epochs?\nCancel will default to 1")
        if ret is None:
            args['save_every_n_epochs'] = 1
        else:
            args['save_every_n_epochs'] = ret

    ret = mb.askyesno(message="Do you want to shuffle captions?")
    if ret:
        args['shuffle_captions'] = True
    else:
        args['shuffle_captions'] = False

    ret = mb.askyesno(message="Do you want to keep some tokens at the front of your captions?")
    if ret:
        ret = sd.askinteger(title="keep_tokens", prompt="How many do you want to keep at the front?"
                                                        "\nCancel will default to 1")
        if ret is None:
            args['keep_tokens'] = 1
        else:
            args['keep_tokens'] = ret

    ret = mb.askyesno(message="Do you want to have a warmup ratio?")
    if ret:
        ret = sd.askfloat(title="warmup_ratio", prompt="What is the ratio of steps to use as warmup "
                                                       "steps?\nCancel will default to None")
        if ret is None:
            args['warmup_lr_ratio'] = None
        else:
            args['warmup_lr_ratio'] = ret

    ret = mb.askyesno(message="Do you want to change the name of output checkpoints?")
    if ret:
        ret = sd.askstring(title="output_name", prompt="What do you want your output name to be?\n"
                                                       "Cancel keeps outputs the original")
        if ret:
            args['change_output_name'] = ret
        else:
            args['change_output_name'] = None

    ret = sd.askstring(title="comment",
                       prompt="Do you want to set a comment that gets put into the metadata?\nA good use of this would "
                              "be to include how to use, such as activation keywords.\nCancel will leave empty")
    if ret is None:
        args['training_comment'] = ret
    else:
        args['training_comment'] = None

    ret = mb.askyesno(message="Do you want to train only one of unet and text encoder?")
    if ret:
        if ret:
            button = ButtonBox("Which do you want to train with?", ["unet_only", "text_only"])
            button.window.mainloop()
            if button.current_value != "":
                args[button.current_value] = True

    ret = mb.askyesno(message="Do you want to save a txt file that contains a list\n"
                              "of all tags that you have used in your training data?\n")
    if ret:
        args['tag_occurrence_txt_file'] = True
        button = ButtonBox("How do you want tags to be ordered?", ["alphabetically", "occurrence-ly"])
        button.window.mainloop()
        if button.current_value == "alphabetically":
            args['sort_tag_occurrence_alphabetically'] = True

    ret = mb.askyesno(message="Do you want to use caption dropout?")
    if ret:
        ret = mb.askyesno(message="Do you want full caption files to dropout randomly?")
        if ret:
            ret = sd.askinteger(title="Caption_File_Dropout",
                                prompt="How often do you want caption files to drop out?\n"
                                       "enter a number from 0 to 100 that is the percentage chance of dropout\n"
                                       "Cancel sets to 0")
            if ret and 0 <= ret <= 100:
                args['caption_dropout_rate'] = ret / 100.0

        ret = mb.askyesno(message="Do you want to have full epochs have no captions?")
        if ret:
            ret = sd.askinteger(title="Caption_epoch_dropout", prompt="The number set here is how often you will have an"
                                                                      "epoch with no captions\nSo if you set 3, then every"
                                                                      "three epochs will not have captions (3, 6, 9)\n"
                                                                      "Cancel will set to None")
            if ret:
                args['caption_dropout_every_n_epochs'] = ret

        ret = mb.askyesno(message="Do you want to have tags to randomly drop?")
        if ret:
            ret = sd.askinteger(title="Caption_tag_dropout", prompt="How often do you want tags to randomly drop out?\n"
                                                                    "Enter a number between 0 and 100, that is the percentage"
                                                                    "chance of dropout.\nCancel sets to 0")
            if ret and 0 <= ret <= 100:
                args['caption_tag_dropout_rate'] = ret / 100.0

    ret = mb.askyesno(message="Do you want to use noise offset? Noise offset seems to allow for SD to better generate\n"
                              "darker or lighter images using this than normal.")
    if ret:
        ret = sd.askfloat(title="noise_offset", prompt="What value do you want to set? recommended value is 0.1,\n"
                                                       "but it can go higher. Cancel defaults to 0.1")
        if ret:
            args['noise_offset'] = ret
        else:
            args['noise_offset'] = 0.1
    return args


def save_json(path, obj: dict) -> None:
    fp = open(os.path.join(path, f"config-{time.time()}.json"), "w")
    json.dump(obj, fp=fp, indent=4)
    fp.close()


def load_json(path, obj: dict) -> dict:
    with open(path) as f:
        json_obj = json.loads(f.read())
    print("loaded json, setting variables...")
    ui_name_scheme = {"pretrained_model_name_or_path": "base_model", "logging_dir": "log_dir",
                      "train_data_dir": "img_folder", "reg_data_dir": "reg_img_folder",
                      "output_dir": "output_folder", "max_resolution": "train_resolution",
                      "lr_scheduler": "scheduler", "lr_warmup": "warmup_lr_ratio",
                      "train_batch_size": "batch_size", "epoch": "num_epochs",
                      "save_at_n_epochs": "save_every_n_epochs", "num_cpu_threads_per_process": "num_workers",
                      "enable_bucket": "buckets", "save_model_as": "save_as", "shuffle_caption": "shuffle_captions",
                      "resume": "load_previous_save_state", "network_dim": "net_dim",
                      "gradient_accumulation_steps": "gradient_acc_steps", "output_name": "change_output_name",
                      "network_alpha": "alpha", "lr_scheduler_num_cycles": "cosine_restarts",
                      "lr_scheduler_power": "scheduler_power"}

    for key in list(json_obj):
        if key in ui_name_scheme:
            json_obj[ui_name_scheme[key]] = json_obj[key]
            if ui_name_scheme[key] in {"batch_size", "num_epochs"}:
                try:
                    json_obj[ui_name_scheme[key]] = int(json_obj[ui_name_scheme[key]])
                except ValueError:
                    print(f"attempting to load {key} from json failed as input isn't an integer")
                    quit(1)

    for key in list(json_obj):
        if obj["json_load_skip_list"] and key in obj["json_load_skip_list"]:
            continue
        if key in obj:
            if key in {"keep_tokens", "warmup_lr_ratio"}:
                json_obj[key] = int(json_obj[key]) if json_obj[key] is not None else None
            if key in {"learning_rate", "unet_lr", "text_encoder_lr"}:
                json_obj[key] = float(json_obj[key]) if json_obj[key] is not None else None
            if obj[key] != json_obj[key]:
                print_change(key, obj[key], json_obj[key])
                obj[key] = json_obj[key]
    print("completed changing variables.")
    return obj


def print_change(value, old, new):
    print(f"{value} changed from {old} to {new}")


class ButtonBox:
    def __init__(self, label: str, button_name_list: list[str]) -> None:
        self.window = tk.Tk()
        self.button_list = []
        self.current_value = ""

        self.window.attributes("-topmost", True)
        self.window.resizable(False, False)
        self.window.eval('tk::PlaceWindow . center')

        def del_window():
            self.window.quit()
            self.window.destroy()

        self.window.protocol("WM_DELETE_WINDOW", del_window)
        tk.Label(text=label, master=self.window).pack()
        for button in button_name_list:
            self.button_list.append(ttk.Button(text=button, master=self.window,
                                               command=partial(self.set_current_value, button)))
            self.button_list[-1].pack()

    def set_current_value(self, value):
        self.current_value = value
        self.window.quit()
        self.window.destroy()


root = tk.Tk()
root.attributes('-topmost', True)
root.withdraw()

if __name__ == "__main__":
    main()