kohya_config_webui.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/WSH032/kohya-config-webui/blob/main/kohya_config_webui.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "| ![visitors](https://visitor-badge.glitch.me/badge?page_id=wsh.kohya_config_webui) | [![GitHub Repo stars](https://img.shields.io/github/stars/WSH032/kohya-config-webui?style=social)](https://github.com/WSH032/kohya-config-webui)</a> |\n",
        "\n",
        "#A WebUI for making config files used by kohya_sd_script\n",
        "\n",
        "Created by [WSH](https://space.bilibili.com/8417436)"
      ],
      "metadata": {
        "id": "7aje0w7w2qsc"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "cellView": "form",
        "id": "ec1xuZuQmvTg"
      },
      "outputs": [],
      "source": [
        "#@title 安装依赖\n",
        "!pip install gradio   > /dev/null 2>&1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {
        "id": "oa6oEre6KC_B",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title 函数部分\n",
        "\n",
        "#A WebUI for making config files used by kohya_sd_script\n",
        "\n",
        "#Created by [WSH](https://space.bilibili.com/8417436)\n",
        "#[![GitHub Repo stars](https://img.shields.io/github/stars/WSH032/kohya-config-webui?style=social)](https://github.com/WSH032/kohya-config-webui)\n",
        "\n",
        "import os\n",
        "import toml\n",
        "import warnings\n",
        "import gradio as gr\n",
        "\n",
        "common_parameter_dict_key_list=[]\n",
        "sample_parameter_dict_key_list=[]\n",
        "plus_parameter_dict_key_list=[]\n",
        "all_parameter_dict_key_list=[]  #后面会有一次all_parameter_dict_key_list = common_parameter_dict_key_list + sample_parameter_dict_key_list + plus_parameter_dict_key_list\n",
        "    \n",
        "\n",
        "common_parameter_dict=({})\n",
        "sample_parameter_dict=({})\n",
        "plus_parameter_dict=({})\n",
        "\n",
        "common_confirm_flag = False   #必须要确认常规参数一次才允许写入toml\n",
        "\n",
        "parameter_len_dict={\"common\":0, \"sample\":0, \"plus\":0}\n",
        "\n",
        "random_symbol = '\\U0001f3b2\\ufe0f'  # 🎲️\n",
        "reuse_symbol = '\\u267b\\ufe0f'  # ♻️\n",
        "paste_symbol = '\\u2199\\ufe0f'  # ↙\n",
        "refresh_symbol = '\\U0001f504'  # 🔄\n",
        "save_style_symbol = '\\U0001f4be'  # 💾\n",
        "apply_style_symbol = '\\U0001f4cb'  # 📋\n",
        "clear_prompt_symbol = '\\U0001f5d1\\ufe0f'  # 🗑️\n",
        "extra_networks_symbol = '\\U0001F3B4'  # 🎴\n",
        "switch_values_symbol = '\\U000021C5' # ⇅\n",
        "folder_symbol = '\\U0001f4c2'  # 📂\n",
        "\n",
        "\n",
        "def check_len_and_2dict(args, parameter_len_dict_value, parameter_dict_key_list, func_name=\"\"):\n",
        "    if len(args) != parameter_len_dict_value:\n",
        "        warnings.warn(f\"传入{func_name}的参数长度不匹配\", UserWarning)\n",
        "    if len(parameter_dict_key_list) != parameter_len_dict_value:\n",
        "        warnings.warn(f\" {func_name}内部字典赋值关键字列表的长度不匹配\", UserWarning)\n",
        "    parameter_dict = dict(zip(parameter_dict_key_list, args))\n",
        "    return parameter_dict\n",
        "\n",
        "def common_parameter_get(*args):\n",
        "    global common_parameter_dict, common_confirm_flag\n",
        "    common_confirm_flag = True    #必须要确认常规参数一次才允许写入toml\n",
        "    common_parameter_dict = check_len_and_2dict(args, parameter_len_dict[\"common\"], common_parameter_dict_key_list, func_name=\"common_parameter_get\")\n",
        "    common_parameter_toml = toml.dumps(common_parameter_dict)\n",
        "    common_parameter_title = \"基础参数配置确认\"\n",
        "    return common_parameter_toml,  common_parameter_title\n",
        "\n",
        "def sample_parameter_get(*args):\n",
        "    global sample_parameter_dict\n",
        "    sample_parameter_dict = check_len_and_2dict(args, parameter_len_dict[\"sample\"], sample_parameter_dict_key_list, func_name=\"sample_parameter_get\")\n",
        "    sample_parameter_toml = toml.dumps(sample_parameter_dict)\n",
        "    sample_parameter_title = \"采样配置确认\"\n",
        "    return sample_parameter_toml,  sample_parameter_title\n",
        "\n",
        "\n",
        "def plus_parameter_get(*args):\n",
        "    global plus_parameter_dict\n",
        "    plus_parameter_dict = check_len_and_2dict(args, parameter_len_dict[\"plus\"], plus_parameter_dict_key_list, func_name=\"plus_parameter_get\")\n",
        "    plus_parameter_toml = toml.dumps(plus_parameter_dict)\n",
        "    plus_parameter_title = \"进阶参数配置确认\"\n",
        "    return plus_parameter_toml,  plus_parameter_title\n",
        "\n",
        "\n",
        "def all_parameter_get(*args):\n",
        "    if len(args) != sum( parameter_len_dict.values() ):\n",
        "         warnings.warn(f\"传入all_parameter_get的参数长度不匹配\", UserWarning)\n",
        "    common_parameter_toml,  common_parameter_title = common_parameter_get( *args[ : parameter_len_dict[\"common\"] ] )\n",
        "    sample_parameter_toml,  sample_parameter_title = sample_parameter_get( *args[ parameter_len_dict[\"common\"] : parameter_len_dict[\"common\"] + parameter_len_dict[\"sample\"] ] )\n",
        "    plus_parameter_toml,  plus_parameter_title = plus_parameter_get( *args[ -parameter_len_dict[\"plus\"] : ] )\n",
        "    return common_parameter_toml, sample_parameter_toml, plus_parameter_toml,  \"全部参数确认\"\n",
        "\n",
        "                      \n",
        "def save_webui_config(save_webui_config_dir, save_webui_config_name, write_files_dir):\n",
        "    os.makedirs(save_webui_config_dir, exist_ok=True)\n",
        "    \n",
        "    other = {\"write_files_dir\":write_files_dir}\n",
        "    param = {**common_parameter_dict, **sample_parameter_dict, **plus_parameter_dict}\n",
        "    dict = { \"other\":other, \"param\":param }\n",
        "\n",
        "    save_webui_config_path = os.path.join(save_webui_config_dir, save_webui_config_name)\n",
        "    with open(save_webui_config_path, \"w\", encoding=\"utf-8\") as f:\n",
        "        webui_config_str = toml.dumps( dict )\n",
        "        f.write(webui_config_str)\n",
        "    return f\"保存webui配置成功，文件在{save_webui_config_path}\"\n",
        "\n",
        "def read_webui_config_get(read_webui_config_dir):\n",
        "    try:\n",
        "        files = [f for f in os.listdir(read_webui_config_dir) if f.endswith(\".toml\") ]\n",
        "        if files:\n",
        "            return gr.update( choices=files,value=files[0] )\n",
        "        else:\n",
        "            return gr.update( choices=[],value=\"没有找到webui配置文件\" )\n",
        "    except Exception as e:\n",
        "        return gr.update( choices=[], value=f\"错误的文件夹路径:{e}\" )\n",
        "\n",
        "def read_webui_config(read_webui_config_dir, read_webui_config_name, write_files_dir, *args):\n",
        "    dir_change_flag = False\n",
        "    param_len = sum( parameter_len_dict.values() )\n",
        "    if len(args) != param_len:\n",
        "        warnings.warn(f\"传入read_webui_config的*args长度不匹配\", UserWarning)\n",
        "    \n",
        "    read_webui_config_path = os.path.join(read_webui_config_dir, read_webui_config_name)\n",
        "    #能打开就正常操作\n",
        "    try:\n",
        "        with open(read_webui_config_path, \"r\", encoding=\"utf-8\") as f:\n",
        "            config_dict = toml.loads( f.read() )\n",
        "        \n",
        "        #能读到[\"other\"].[\"write_files_dir\"]就改，读不到就用原写入地址\n",
        "        try:\n",
        "            if config_dict[\"other\"][\"write_files_dir\"] != write_files_dir:\n",
        "                write_files_dir = config_dict[\"other\"][\"write_files_dir\"]\n",
        "                dir_change_flag = True\n",
        "        except KeyError:\n",
        "            pass\n",
        "        \n",
        "        param_dict_key_list = list( config_dict.get(\"param\",{}).keys() )\n",
        "        #找出共有的key进行赋值，非共有的报错\n",
        "        both_key = set(all_parameter_dict_key_list) & set(param_dict_key_list)\n",
        "        parameter_unique_key = set(all_parameter_dict_key_list) - set(both_key)\n",
        "        config_unique_key = set(param_dict_key_list) - set(both_key)\n",
        "        #赋值\n",
        "        count = 0\n",
        "        if both_key:\n",
        "            args = list(args)\n",
        "            for key in both_key:\n",
        "                index = all_parameter_dict_key_list.index(key)\n",
        "                args[ index ] = config_dict[\"param\"][key]\n",
        "                count += 1\n",
        "            args = tuple(args)\n",
        "        read_done = f\"\\n读取完成,WebUI中共有{param_len}项参数,更新了其中{count}项\\n\" + f\"写入文件夹发生改变:{write_files_dir}\" if dir_change_flag else \"\"\n",
        "        config_warning = f\"\\nwebui-config文件中以下参数可能已经失效或错误：\\n{config_unique_key}\\n\" if config_unique_key else \"\"\n",
        "        parameter_warning = f\"\\nWebUI中以下参数在webui-config文件中未找到，不发生修改：\\n{parameter_unique_key}\\n\" if parameter_unique_key else \"\"\n",
        "        str = read_done + config_warning + parameter_warning\n",
        "        return  str, write_files_dir, *args\n",
        "\n",
        "    #打不开就返回原值\n",
        "    except FileNotFoundError:\n",
        "        return \"文件或目录不存在\", write_files_dir, *args\n",
        "    except PermissionError:\n",
        "        return \"没有权限访问文件或目录\", write_files_dir, *args\n",
        "    except OSError as e:\n",
        "        return f\"something wrong：{e}\", write_files_dir, *args\n",
        "    \n",
        "    \n",
        "\n",
        "def model_get(model_dir):\n",
        "    try:\n",
        "        files = [f for f in os.listdir(model_dir) if os.path.isfile(os.path.join(model_dir, f))]\n",
        "        if files:\n",
        "            return gr.update( choices=files,value=files[0] )\n",
        "        else:\n",
        "            return gr.update( choices=[],value=\"没有找到模型\" )\n",
        "    except Exception as e:\n",
        "        return gr.update( choices=[], value=f\"错误的文件夹路径:{e}\" )\n",
        "\n",
        "\n",
        "def write_files(write_files_dir):\n",
        "\n",
        "    if not common_confirm_flag:\n",
        "        return \"必须要确认常规参数一次才允许写入toml\"\n",
        "\n",
        "    write_files_dir = write_files_dir if write_files_dir else os.getcwd()\n",
        "    os.makedirs(write_files_dir, exist_ok=True)\n",
        "    config_file_toml_path = os.path.join(write_files_dir, \"config_file.toml\")\n",
        "    sample_prompts_txt_path = os.path.join(write_files_dir, \"sample_prompts.txt\")\n",
        "\n",
        "    all = {**common_parameter_dict, **sample_parameter_dict, **plus_parameter_dict}\n",
        "\n",
        "    def parameter2toml():\n",
        "\n",
        "        #生成config_file.toml的字典\n",
        "\n",
        "        #model_arguments部分\n",
        "        model_arguments = { key: all.get(key) for key in [\"v2\", \"v_parameterization\"] }\n",
        "        \"\"\" 生成底模路径 \"\"\"\n",
        "        base_model_path = os.path.join( all.get(\"base_model_dir\"), all.get(\"base_model_name\") )\n",
        "        model_arguments.update( {\"pretrained_model_name_or_path\": base_model_path} )\n",
        "        \"\"\" 生成vae路径 \"\"\"\n",
        "        if all.get(\"use_vae\"):\n",
        "            vae_model_path = os.path.join( all.get(\"vae_model_dir\"), all.get(\"vae_model_name\") )\n",
        "            model_arguments.update( {\"vae\": vae_model_path} )\n",
        "\n",
        "        #additional_network_arguments部分\n",
        "        additional_network_arguments = { key: all.get(key) for key in [\"unet_lr\", \"text_encoder_lr\", \"network_dim\",\\\n",
        "                                            \"network_alpha\", \"network_train_unet_only\",\\\n",
        "                                            \"network_train_text_encoder_only\"] }\n",
        "        \"\"\" 生成如network_module = \"locon.locon_kohya\" \"\"\"\n",
        "        #[\"LoRA-LierLa\", \"LoRA-C3Lier\", \"LoCon_Lycoris\", \"LoHa_Lycoris\", \"DyLoRa-LierLa\", \"DyLoRa-C3Lier\"]\n",
        "        #主要负责network_module的参数生成\n",
        "        def network_module_param(train_method):\n",
        "            conv_dim = all.get(\"conv_dim\") if train_method != \"DyLoRa-C3Lier\" else all.get(\"network_dim\")\n",
        "            conv_alpha = all.get(\"conv_alpha\")\n",
        "            algo = \"lora\" if train_method == \"LoCon_Lycoris\" else \"loha\"\n",
        "            unit = all.get(\"unit\")\n",
        "            if train_method in [\"LoRA-LierLa\", \"LoRA-C3Lier\"]:\n",
        "                network_module = \"networks.lora\"\n",
        "                if train_method == \"LoRA-C3Lier\":\n",
        "                    network_module_args = [f\"conv_dim={conv_dim}\", f\"conv_alpha={conv_alpha}\"]\n",
        "                else:\n",
        "                    network_module_args = []\n",
        "            elif train_method in [\"LoCon_Lycoris\", \"LoHa_Lycoris\"]:\n",
        "                network_module = \"lycoris.kohya\"\n",
        "                network_module_args = [f\"conv_dim={conv_dim}\", f\"conv_alpha={conv_alpha}\", f\"algo={algo}\"]\n",
        "            elif train_method in [\"DyLoRa-LierLa\", \"DyLoRa-C3Lier\"]:\n",
        "                network_module = \"networks.dylora\"\n",
        "                if train_method == \"DyLoRa-C3Lier\":\n",
        "                    network_module_args = [f\"conv_dim={conv_dim}\", f\"conv_alpha={conv_alpha}\", f\"unit={unit}\"]\n",
        "                else:\n",
        "                    network_module_args = [f\"unit={unit}\"]\n",
        "            else: \n",
        "                warnings.warn(f\"训练方法参数生成出错\", UserWarning)\n",
        "            return network_module, network_module_args\n",
        "        network_module, network_module_args = network_module_param( all.get(\"train_method\") )\n",
        "        #更多network_args部分（主要为分层训练）\n",
        "        network_lr_weight_args = [ f\"{name}={all.get(name)}\" for name in [\"up_lr_weight\", \"mid_lr_weight\", \"down_lr_weight\"] if all.get(name) ]\n",
        "\n",
        "        def network_block_param(train_method):\n",
        "            lst = [\"block_dims\", \"block_alphas\", \"conv_block_dims\", \"conv_block_alphas\"]\n",
        "            if train_method == \"LoRA-LierLa\":\n",
        "                return [ f\"{name}={all.get(name)}\" for name in lst[0:1] if all.get(name) ]\n",
        "            if train_method in [\"LoRA-C3Lier\", \"LoCon_Lycoris\", \"LoHa_Lycoris\"]:\n",
        "                return [ f\"{name}={all.get(name)}\" for name in lst if all.get(name) ]\n",
        "            else:\n",
        "                return []\n",
        "        network_block_args = network_block_param( all.get(\"train_method\") )\n",
        "        \n",
        "\n",
        "        network_args = []\n",
        "        network_args.extend(network_module_args)\n",
        "        network_args.extend(network_lr_weight_args)\n",
        "        network_args.extend(network_block_args)\n",
        "\n",
        "        additional_network_arguments.update( { \"network_module\":network_module } )\n",
        "        additional_network_arguments.update( {\"network_args\":network_args} )          \n",
        "\n",
        "        #optimizer_arguments部分\n",
        "        optimizer_arguments = { key: all.get(key) for key in [\"optimizer_type\", \"lr_scheduler\", \"lr_warmup_steps\"] }\n",
        "        \"\"\"只有余弦重启调度器指定重启次数\"\"\"\n",
        "        if all.get(\"lr_scheduler\") == \"cosine_with_restarts\":\n",
        "            optimizer_arguments.update( {\"lr_restart_cycles\":all.get(\"lr_restart_cycles\")} )\n",
        "        \"\"\"学习率lr指定=unet_lr\"\"\"\n",
        "        optimizer_arguments.update( {\"learning_rate\":all.get(\"unet_lr\")} )\n",
        "            #optimizer_args（待添加）\n",
        "\n",
        "        #dataset_arguments部分\n",
        "        dataset_arguments = { key: all.get(key) for key in [\"cache_latents\", \"shuffle_caption\", \"enable_bucket\"] }\n",
        "        \n",
        "        #training_arguments部分\n",
        "        training_arguments = { key: all.get(key) for key in [\"batch_size\", \"noise_offset\", \"keep_tokens\",\\\n",
        "                                      \"min_bucket_reso\", \"max_bucket_reso\",\\\n",
        "                                      \"caption_extension\", \"max_token_length\", \"seed\",\\\n",
        "                                      \"xformers\", \"lowram\"]\n",
        "        }\n",
        "        \"\"\"min_snr_gamma大于零才生效\"\"\"\n",
        "        if all.get(\"min_snr_gamma\") > 0:\n",
        "            training_arguments.update( { \"min_snr_gamma\":all.get(\"min_snr_gamma\") } )\n",
        "        \"\"\" 最大训练时间 \"\"\"\n",
        "        training_arguments.update( { all.get(\"max_train_method\"):all.get(\"max_train_value\") } )\n",
        "        \"\"\" 训练分辨率 \"\"\"\n",
        "        training_arguments.update( { \"resolution\":f\"{all.get('width')},{all.get('height')}\" } )\n",
        "        \"\"\" 如果v2开启，则不指定clip_skip \"\"\"\n",
        "        if not all.get(\"v2\"):\n",
        "            training_arguments.update( { \"clip_skip\":all.get(\"clip_skip\") } )\n",
        "        \"\"\" 重训练模块 \"\"\"\n",
        "        if all.get(\"use_retrain\") == \"model\":\n",
        "            training_arguments.update( { \"network_weights\":all.get(\"retrain_dir\") } )\n",
        "        elif all.get(\"use_retrain\") == \"state\":\n",
        "            training_arguments.update( { \"resume\":all.get(\"retrain_dir\") } )\n",
        "        \"\"\"  训练精度、保存精度 \"\"\"\n",
        "        training_arguments.update( { \"mixed_precision\":\"fp16\" } )\n",
        "        training_arguments.update( { \"save_precision\":\"fp16\" } )\n",
        "        \n",
        "\n",
        "\n",
        "        #sample_prompt_arguments部分（采样间隔，采样文件地址待添加）\n",
        "        sample_prompt_arguments = { key: all.get(key) for key in [\"sample_sampler\"] }\n",
        "        if all.get(\"sample_every_n_type\"):    #如果采样部分没确认过一次，会出现all.get(\"sample_every_n_type\")=None:None的字典造成报错\n",
        "            sample_prompt_arguments.update( {all.get(\"sample_every_n_type\"):all.get(\"sample_every_n_type_value\")} )\n",
        "\n",
        "        #dreambooth_arguments部分\n",
        "        dreambooth_arguments = { key: all.get(key) for key in [\"train_data_dir\", \"reg_data_dir\", \"prior_loss_weight\"] }\n",
        "\n",
        "        #saving_arguments部分\n",
        "        saving_arguments = { key: all.get(key) for key in [\"output_dir\",\\\n",
        "                                      \"output_name\", \"save_every_n_epochs\", \"save_n_epoch_ratio\",\\\n",
        "                                      \"save_last_n_epochs\", \"save_state\", \"save_model_as\" ]\n",
        "        }\n",
        "        \"\"\" 指定log输出目录与output相同 \"\"\"\n",
        "        saving_arguments.update( { \"logging_dir\":os.path.join( all.get(\"output_dir\"), \"logs\" ) } )\n",
        "        \"\"\" 指定log前缀和输出名字相同 \"\"\"\n",
        "        saving_arguments.update( { \"log_prefix\":all.get(\"output_name\") } )\n",
        "        \n",
        "\n",
        "        toml_dict = {\"model_arguments\":model_arguments,\n",
        "               \"additional_network_arguments\":additional_network_arguments,\n",
        "               \"optimizer_arguments\":optimizer_arguments,\n",
        "               \"dataset_arguments\":dataset_arguments,\n",
        "               \"training_arguments\":training_arguments,\n",
        "               \"sample_prompt_arguments\":sample_prompt_arguments,\n",
        "               \"dreambooth_arguments\":dreambooth_arguments,\n",
        "               \"saving_arguments\":saving_arguments,\n",
        "        }\n",
        "        toml_str = toml.dumps(toml_dict)\n",
        "        return toml_str\n",
        "    def sample_parameter2txt():\n",
        "        #key_list = [\"prompt\", \"negative\", \"sample_width\", \"sample_height\", \"sample_scale\", \"sample_steps\", \"sample_seed\"]\n",
        "\n",
        "        if not all.get('sample_seed'):    #如果采样部分没确认过，会出现all.get('sample_seed')=None > 0造成报错\n",
        "            return \"\"\n",
        "        sample_str = f\"\"\"{all.get(\"prompt\")}  \\\n",
        "--n {all.get(\"negative\")}  \\\n",
        "--w {all.get(\"sample_width\")}  \\\n",
        "--h {all.get(\"sample_height\")}  \\\n",
        "--l {all.get(\"sample_scale\")}  \\\n",
        "--s {all.get(\"sample_steps\")}  \\\n",
        "{f\"--d {all.get('sample_seed')}\" if all.get('sample_seed') > 0 else \"\"}\"\"\"\n",
        "        return sample_str\n",
        "\n",
        "    def write(content, path):\n",
        "        with open(path, \"w\", encoding=\"utf-8\") as f:\n",
        "            f.write(content)\n",
        "\n",
        "    write(parameter2toml(), config_file_toml_path)\n",
        "    write(sample_parameter2txt(), sample_prompts_txt_path)\n",
        "    write_files_title = f\"写入成功, 训练配置文件在{config_file_toml_path}, 采样参数文件在{sample_prompts_txt_path}\"\n",
        "    return write_files_title\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LT7NaOl8Jfgf",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title WebUI部分\n",
        "\n",
        "\n",
        "with gr.Blocks() as demo:\n",
        "    with gr.Accordion(\"保存、读取\\nwebui配置\", open=False):\n",
        "        save_read_webui_config_title = gr.Markdown(\"保存或读取\")\n",
        "        with gr.Row():\n",
        "            save_webui_config_button = gr.Button(\"保存\")\n",
        "        with gr.Row():\n",
        "            save_webui_config_dir = gr.Textbox(lines=1, label=\"保存目录\", value=os.path.join(os.getcwd(),\"kohya_config_webui_save\") )\n",
        "            save_webui_config_name = gr.Textbox(lines=1, label=\"保存名字（以toml为扩展名，否则不会被读取）\", value=\"kohya_config_webui_save.toml\" )\n",
        "        with gr.Row():\n",
        "            read_webui_config_get_button = gr.Button(refresh_symbol)\n",
        "            read_webui_config_button = gr.Button(\"读取\")\n",
        "        with gr.Row():\n",
        "            read_webui_config_dir = gr.Textbox(lines=1, label=\"读取目录\", value=os.path.join(os.getcwd(),\"kohya_config_webui_save\") )  \n",
        "            read_webui_config_name = gr.Dropdown(choices=[], label=\"读取文件\", value=\"\" )          \n",
        "    with gr.Row():\n",
        "        write_files_button = gr.Button(\"生成toml参数与采样配置文件\")\n",
        "        all_parameter_get_button = gr.Button(\"全部参数确认\")\n",
        "        write_files_dir = gr.Textbox( lines=1, label=\"写入文件夹\", placeholder=\"一般填kohya_script目录，留空就默认根目录\", value=\"\" )\n",
        "    write_files_title = gr.Markdown(\"生成适用于kohya/train_network.py的配置文件\")\n",
        "    with gr.Tabs():\n",
        "        with gr.TabItem(\"基础参数\"):\n",
        "            common_parameter_get_button = gr.Button(\"确定\")\n",
        "            common_parameter_title = gr.Markdown(\"\")\n",
        "            with gr.Accordion(\"当前基础参数配置\", open=False):\n",
        "                common_parameter_toml = gr.Textbox(label=\"toml形式\", placeholder=\"基础参数\", value=\"\")\n",
        "            with gr.Row():\n",
        "                train_data_dir = gr.Textbox(lines=1, label=\"train_data_dir\", placeholder=\"训练集路径\", value=\"\")\n",
        "            with gr.Accordion(\"使用正则化(可选)\", open=False):\n",
        "                with gr.Row():\n",
        "                    reg_data_dir = gr.Textbox(lines=1, label=\"reg_data_dir\", placeholder=\"正则化集路径（填入意味着启用正则化）\", value=\"\")\n",
        "                    prior_loss_weight = gr.Slider(0, 1, step=0.01, value=0.3, label=\"正则化权重\")\n",
        "            with gr.Row():\n",
        "                base_model_dir = gr.Textbox(label=\"底模文件夹地址\", placeholder=\"文件夹路径\", value=\"\")\n",
        "                base_model_name = gr.Dropdown(choices=[],label=\"底模\",value=\"\")\n",
        "                base_model_get_button = gr.Button(refresh_symbol)\n",
        "            with gr.Accordion(\"使用vae(可选)\", open=False):\n",
        "                with gr.Row():\n",
        "                    use_vae = gr.Checkbox(label=\"是否使用vae\",value=False)\n",
        "                with gr.Row():\n",
        "                    vae_model_dir = gr.Textbox(label=\"vae文件夹地址\", placeholder=\"文件夹路径\", value=\"\")\n",
        "                    vae_model_name = gr.Dropdown(choices=[],label=\"vae\", value=\"\")\n",
        "                    vae_model_get_button = gr.Button(refresh_symbol)\n",
        "            with gr.Row():\n",
        "                width = gr.Slider(64, 1920, step=64, value=512, label=\"训练分辨率（宽）width\")\n",
        "                height = gr.Slider(64, 1920, step=64, value=512, label=\"训练分辨率（高）height\")\n",
        "                batch_size = gr.Slider(1, 24, step=1, value=1, label=\"batch大小\")\n",
        "            with gr.Row():\n",
        "                noise_offset = gr.Slider(0, 1, step=0.01, value=0.05, label=\"noise_offset\")\n",
        "                keep_tokens = gr.Slider(0, 225, step=1, value=0, label=\"keep_tokens\")\n",
        "                min_snr_gamma = gr.Slider(0, 100, step=0.1, value=5, label=\"min_snr_gamma(设置为0则不生效)\")\n",
        "            \"\"\"\n",
        "            with gr.Row():\n",
        "                gr.Markdown(\"repeat * 图片数 = 每个epoch的steps数\")\n",
        "            \"\"\"\n",
        "            with gr.Row():\n",
        "                max_train_method = gr.Dropdown([\"max_train_epochs\",\"max_train_steps\"], label=\"以epochs或steps来指定最大训练时间\", value=\"max_train_epochs\")\n",
        "                max_train_value = gr.Number(label=\"最大训练epochs\\steps数\", value=10, precision=0)\n",
        "            with gr.Accordion(\"输出设置\", open=True):\n",
        "                with gr.Row():\n",
        "                    output_dir = gr.Textbox( label=\"模型、log日志输出地址（自行修改）\", placeholder=\"文件夹路径\",value=os.path.join(os.getcwd(),\"output\") )\n",
        "                    output_name = gr.Textbox(label=\"输出模型名称（自行修改）\", placeholder=\"名称\",value=\"output_name\")\n",
        "                    save_model_as = gr.Dropdown([\"safetensors\",\"ckpt\",\"pt\"], label=\"保存模型格式\", value=\"safetensors\")\n",
        "                with gr.Row():\n",
        "                    save_every_n_epochs = gr.Slider(1, 499, step=1, value=1, label=\"每n个epoch保存一次\")\n",
        "                    save_n_epoch_ratio = gr.Slider(1, 499, step=1, value=0, label=\"等间隔保存n个(如不为0，会覆盖每n个epoch保存一次)\")\n",
        "                    save_last_n_epochs = gr.Slider(1, 499, step=1, value=499, label=\"最多保存n个（后面的出来就会把前面删了,优先级最高）\")\n",
        "                with gr.Row():   \n",
        "                    save_state = gr.Checkbox(label=\"保存学习状态\",value=False)\n",
        "            with gr.Row():\n",
        "                optimizer_type = gr.Dropdown([\"AdamW8bit\", \"Lion\", \"DAdaptation\", \"AdamW\", \"SGDNesterov\", \"SGDNesterov8bit\", \"AdaFactor\"],\\\n",
        "                                label=\"optimizer_type优化器类型\", value=\"AdamW8bit\")\n",
        "                unet_lr = gr.Number(label=\"unet学习率\", value=1e-4)\n",
        "                text_encoder_lr = gr.Number(label=\"text_encoder学习率\", value=1e-5)\n",
        "            with gr.Row():\n",
        "                lr_scheduler = gr.Dropdown([\"cosine_with_restarts\",\"cosine\",\"polynomial\",\"linear\",\"constant_with_warmup\",\"constant\"],\\\n",
        "                               label=\"lr_scheduler学习率调度器\", value=\"cosine_with_restarts\")\n",
        "                lr_warmup_steps = gr.Number(label=\"升温步数\", value=0, precision=0)\n",
        "                lr_restart_cycles = gr.Number(label=\"退火重启次数\", value=1, precision=0)\n",
        "            with gr.Row():\n",
        "                train_method = gr.Dropdown([\"LoRA-LierLa\", \"LoRA-C3Lier\",\\\n",
        "                                \"LoCon_Lycoris\",\"LoHa_Lycoris\",\\\n",
        "                                \"DyLoRa-LierLa\", \"DyLoRa-C3Lier\"],\\\n",
        "                                label=\"train_method训练方法\", value=\"LoRA-LierLa\")\n",
        "                network_dim = gr.Number(label=\"线性dim\", value=32, precision=0)\n",
        "                network_alpha = gr.Number(label=\"线性alpha（可以为小数）\", value=16)\n",
        "            with gr.Accordion(\"额外网络参数(LoRA-C3Lier、LoCon、LoHa、DyLoRa-C3Lier都属于卷积,unit为DyLoRa专用)\", open=True):\n",
        "                with gr.Row():\n",
        "                    with gr.Column():\n",
        "                        conv_dim = gr.Number(label=\"卷积dim\", info=\"使用DyLoRa-C3Lier时会被设置为等于基础dim\", value=8, precision=0)\n",
        "                    with gr.Column():\n",
        "                        conv_alpha = gr.Number(label=\"卷积alpha\", info=\"可以为小数\", value=1)\n",
        "                    with gr.Column():\n",
        "                        unit = gr.Number(label=\"分割单位unit(整数)\", info=\"使用DyLoRa时，请让dim为unit的倍数\", value=1, precision=0)\n",
        "            with gr.Row():          \n",
        "                v2 = gr.Checkbox(label=\"v2\")\n",
        "                v_parameterization = gr.Checkbox(label=\"v_parameterization\")\n",
        "                lowram = gr.Checkbox(label=\"lowram\")\n",
        "                xformers = gr.Checkbox(label=\"xformers\",value=True)\n",
        "                cache_latents = gr.Checkbox(label=\"cache_latents\",value=True)\n",
        "                shuffle_caption = gr.Checkbox(label=\"shuffle_caption\",value=True)\n",
        "                enable_bucket = gr.Checkbox(label=\"enable_bucket\",value=True)\n",
        "        with gr.TabItem(\"采样参数\"):\n",
        "            sample_parameter_get_button = gr.Button(\"确定\")\n",
        "            sample_parameter_title = gr.Markdown(\"\")\n",
        "            with gr.Accordion(\"当前采样配置\", open=False):\n",
        "                sample_parameter_toml = gr.Textbox(label=\"toml形式\", placeholder=\"采样配置\", value=\"\")\n",
        "            with gr.Row():\n",
        "                #enable_sample = gr.Checkbox(label=\"是否启用采样功能\")\n",
        "                sample_every_n_type = gr.Dropdown([\"sample_every_n_epochs\", \"sample_every_n_steps\"], label=\"sample_every_n_type\", value=\"sample_every_n_epochs\")\n",
        "                sample_every_n_type_value = gr.Number(label=\"sample_every_n_type_value\", value=1, precision=0)\n",
        "            with gr.Row():\n",
        "                sample_sampler = gr.Dropdown([\"ddim\", \"pndm\", \"lms\", \"euler\", \"euler_a\", \"heun\",\\\n",
        "                            \"dpm_2\", \"dpm_2_a\", \"dpmsolver\",\"dpmsolver++\", \"dpmsingle\",\\\n",
        "                            \"k_lms\", \"k_euler\", \"k_euler_a\", \"k_dpm_2\", \"k_dpm_2_a\"],\\\n",
        "                            label=\"采样器\", value=\"euler_a\")\n",
        "                sample_seed = gr.Number(label=\"采样种子(-1不是随机，大于0才生效)\", value=-1, precision=0)\n",
        "            with gr.Row():\n",
        "                sample_width = gr.Slider(64, 1920, step=64, value=512, label=\"采样图片宽\")\n",
        "                sample_height = gr.Slider(64, 1920, step=64, value=768, label=\"采样图片高\")\n",
        "                sample_scale = gr.Slider(1, 30, step=0.5, value=7, label=\"提示词相关性\")\n",
        "                sample_steps = gr.Slider(1, 150, step=1, value=24, label=\"采样迭代步数\")\n",
        "            with gr.Row():\n",
        "                prompt = gr.Textbox(lines=10, label=\"prompt\", placeholder=\"正面提示词\", value=\"(masterpiece, best quality, hires:1.2), 1girl, solo,\")\n",
        "                default_negative = (\"(worst quality, bad quality:1.4), \"\n",
        "                          \"lowres, bad anatomy, bad hands, text, error, \"\n",
        "                          \"missing fingers, extra digit, fewer digits, \"\n",
        "                          \"cropped, worst quality, low quality, normal quality, \"\n",
        "                          \"jpeg artifacts,signature, watermark, username, blurry,\")\n",
        "                negative = gr.Textbox(lines=10, label=\"negative\", placeholder=\"负面提示词\", value=default_negative)\n",
        "        with gr.TabItem(\"进阶参数\"):\n",
        "            plus_parameter_get_button = gr.Button(\"确定\")\n",
        "            plus_parameter_title = gr.Markdown(\"\")\n",
        "            with gr.Accordion(\"当前进阶参数配置\", open=False):\n",
        "                plus_parameter_toml = gr.Textbox(label=\"toml形式\", placeholder=\"进阶参数\", value=\"\")\n",
        "            with gr.Row():\n",
        "                use_retrain = gr.Dropdown([\"no\",\"model\",\"state\"], label=\"是否使用重训练\", value=\"no\")\n",
        "                retrain_dir = gr.Textbox(lines=1, label=\"重训练路径\", placeholder=\"模型或者状态路径\", value=\"\")\n",
        "            with gr.Row():\n",
        "                min_bucket_reso = gr.Slider(64, 1920, step=64, value=256, label=\"最低桶分辨率\")\n",
        "                max_bucket_reso = gr.Slider(64, 1920, step=64, value=1024, label=\"最高桶分辨率\")\n",
        "                clip_skip = gr.Slider(0, 25, step=1, value=2, label=\"跳过层数\")\n",
        "                max_token_length = gr.Slider(75, 225, step=75, value=225, label=\"训练最大token数\")\n",
        "                caption_extension = gr.Textbox(lines=1, label=\"标签文件扩展名\", placeholder=\"一般填.txt或.cap\", value=\".txt\")\n",
        "                seed = gr.Number(label=\"种子\", value=1337, precision=0)\n",
        "            with gr.Row():\n",
        "                network_train_unet_only= gr.Checkbox(label=\"仅训练unet网络\",value=False)\n",
        "                network_train_text_encoder_only = gr.Checkbox(label=\"仅训练text_encoder网络\",value=False)\n",
        "            with gr.Accordion(\"分层学习模块\", open=True):\n",
        "                gr.Markdown(\"学习率分层，为不同层的结构指定不同学习率倍数； 如果某一层权重为0，那该层不会被创建\")\n",
        "                with gr.Row():\n",
        "                    with gr.Column(scale=15):\n",
        "                        up_lr_weight = gr.Textbox(lines=1, label=\"上层学习率权重\", placeholder=\"留空则不启用\",\\\n",
        "                                      info=\"15层，例如1.5,1.5,1.5,1.5,1.0,1.0,1.0,1.0,0.5,0.5,0.5,0.5\", value=\"\")\n",
        "                    with gr.Column(scale=1):\n",
        "                        mid_lr_weight = gr.Textbox(lines=1, label=\"中层学习率权重\", placeholder=\"留空则不启用\",\\\n",
        "                                      info=\"1层，例如2.0\", value=\"\")\n",
        "                    with gr.Column(scale=15):\n",
        "                        down_lr_weight = gr.Textbox(lines=1, label=\"下层学习率权重\", placeholder=\"留空则不启用\",\\\n",
        "                                      info=\"15层，例如0.5,0.5,0.5,0.5,1.0,1.0,1.0,1.0,1.5,1.5,1.5,1.5\", value=\"\")\n",
        "                gr.Markdown(\"dim和alpha分层，为不同层的结构指定不同的dim和alpha（`DyLoRa`无法使用，卷积分层只有`LoRa-C3Lier、LoCon、LoHa`可以使用）\")\n",
        "                with gr.Row():\n",
        "                        block_dims = gr.Textbox(lines=1, label=\"线性dim分层\", placeholder=\"留空则不启用\",\\\n",
        "                                      info=\"25层（上中下），例如2,4,4,4,8,8,8,8,12,12,12,12,16,12,12,12,12,8,8,8,8,4,4,4,2\", value=\"\")\n",
        "                        block_alphas = gr.Textbox(lines=1, label=\"线性alpha分层\", placeholder=\"留空则不启用\",\\\n",
        "                                      info=\"25层（上中下），例如2,2,2,2,4,4,4,4,6,6,6,6,8,6,6,6,6,4,4,4,4,2,2,2,2\", value=\"\")\n",
        "                with gr.Row():\n",
        "                        conv_block_dims = gr.Textbox(lines=1, label=\"卷积dim分层\", placeholder=\"留空则不启用\",\\\n",
        "                                        info=\"25层（上中下），例如2,2,2,2,4,4,4,4,6,6,6,6,8,6,6,6,6,4,4,4,4,2,2,2,2\", value=\"\")\n",
        "                        conv_block_alphas = gr.Textbox(lines=1, label=\"卷积alpha分层\", placeholder=\"留空则不启用\",\\\n",
        "                                        info=\"25层（上中下），例如2,2,2,2,4,4,4,4,6,6,6,6,8,6,6,6,6,4,4,4,4,2,2,2,2\", value=\"\")\n",
        "\n",
        "\n",
        "    def dict_key_list_2_list(dict_key_list):\n",
        "        list = []\n",
        "        for key in dict_key_list:\n",
        "            try:\n",
        "                list.append(globals()[key])\n",
        "            except KeyError:\n",
        "                print(f\"Error: parameter_dict_key_list中{key}不存在\")\n",
        "        list_len = len(list)\n",
        "        return list, list_len\n",
        "\n",
        "    common_parameter_dict_key_list = [\"train_data_dir\",\n",
        "                      \"reg_data_dir\",\n",
        "                      \"prior_loss_weight\",\n",
        "                      \"base_model_dir\",\n",
        "                      \"base_model_name\",\n",
        "                      \"use_vae\",\n",
        "                      \"vae_model_dir\",\n",
        "                      \"vae_model_name\",\n",
        "                      \"width\",\n",
        "                      \"height\",\n",
        "                      \"batch_size\",\n",
        "                      \"noise_offset\",\n",
        "                      \"keep_tokens\",\n",
        "                      \"min_snr_gamma\",\n",
        "                      \"max_train_method\",\n",
        "                      \"max_train_value\",\n",
        "                      \"output_dir\",\n",
        "                      \"output_name\",\n",
        "                      \"save_model_as\",\n",
        "                      \"save_every_n_epochs\",\n",
        "                      \"save_n_epoch_ratio\",\n",
        "                      \"save_last_n_epochs\",\n",
        "                      \"save_state\",\n",
        "                      \"optimizer_type\",\n",
        "                      \"unet_lr\",\n",
        "                      \"text_encoder_lr\",\n",
        "                      \"lr_scheduler\",\n",
        "                      \"lr_warmup_steps\",\n",
        "                      \"lr_restart_cycles\",\n",
        "                      \"train_method\",\n",
        "                      \"network_dim\",\n",
        "                      \"network_alpha\",\n",
        "                      \"conv_dim\",\n",
        "                      \"conv_alpha\",\n",
        "                      \"unit\",\n",
        "                      \"v2\",\n",
        "                      \"v_parameterization\",\n",
        "                      \"lowram\",\n",
        "                      \"xformers\",\n",
        "                      \"cache_latents\",\n",
        "                      \"shuffle_caption\",\n",
        "                      \"enable_bucket\"]\n",
        "    common_parameter_list, parameter_len_dict[\"common\"] = dict_key_list_2_list(common_parameter_dict_key_list)\n",
        "    sample_parameter_dict_key_list = [\"sample_every_n_type\",\n",
        "                      \"sample_every_n_type_value\",\n",
        "                      \"sample_sampler\",\n",
        "                      \"sample_seed\",\n",
        "                      \"sample_width\",\n",
        "                      \"sample_height\",\n",
        "                      \"sample_scale\",\n",
        "                      \"sample_steps\",\n",
        "                      \"prompt\",\n",
        "                      \"negative\"]\n",
        "    sample_parameter_list, parameter_len_dict[\"sample\"] = dict_key_list_2_list(sample_parameter_dict_key_list)\n",
        "    plus_parameter_dict_key_list = [\"use_retrain\",\n",
        "                    \"retrain_dir\",\n",
        "                    \"min_bucket_reso\",\n",
        "                    \"max_bucket_reso\",\n",
        "                    \"clip_skip\",\n",
        "                    \"max_token_length\",\n",
        "                    \"caption_extension\",\n",
        "                    \"seed\",\n",
        "                    \"network_train_unet_only\",\n",
        "                    \"network_train_text_encoder_only\",\n",
        "                    \"up_lr_weight\",\n",
        "                    \"mid_lr_weight\",\n",
        "                    \"down_lr_weight\",\n",
        "                    \"block_dims\",\n",
        "                    \"block_alphas\",\n",
        "                    \"conv_block_dims\",\n",
        "                    \"conv_block_alphas\"]\n",
        "    plus_parameter_list, parameter_len_dict[\"plus\"] = dict_key_list_2_list(plus_parameter_dict_key_list)\n",
        "\n",
        "    #注意，这几个list相加的顺序不能错\n",
        "    all_parameter_list = common_parameter_list + sample_parameter_list + plus_parameter_list\n",
        "    all_parameter_dict_key_list = common_parameter_dict_key_list + sample_parameter_dict_key_list + plus_parameter_dict_key_list\n",
        "\n",
        "    save_webui_config_button.click(fn=save_webui_config,\n",
        "                    inputs=[save_webui_config_dir, save_webui_config_name, write_files_dir],\n",
        "                    outputs=save_read_webui_config_title      \n",
        "    )\n",
        "    read_webui_config_get_button.click(fn=read_webui_config_get,\n",
        "                      inputs=[read_webui_config_dir],\n",
        "                      outputs=[read_webui_config_name]       \n",
        "    )\n",
        "    read_webui_config_button.click(fn=read_webui_config,\n",
        "                    inputs=[read_webui_config_dir, read_webui_config_name, write_files_dir] + all_parameter_list,\n",
        "                    outputs=[save_read_webui_config_title, write_files_dir] + all_parameter_list\n",
        "    )\n",
        "    common_parameter_get_button.click(fn=common_parameter_get,\n",
        "                    inputs=common_parameter_list,\n",
        "                    outputs=[common_parameter_toml,  common_parameter_title]\n",
        "    )\n",
        "    sample_parameter_get_button.click(fn=sample_parameter_get,\n",
        "                    inputs=sample_parameter_list,\n",
        "                    outputs=[sample_parameter_toml,  sample_parameter_title]\n",
        "    )\n",
        "    plus_parameter_get_button.click(fn=plus_parameter_get,\n",
        "                    inputs=plus_parameter_list,\n",
        "                    outputs=[plus_parameter_toml,  plus_parameter_title]\n",
        "    )\n",
        "    all_parameter_get_button.click(fn=all_parameter_get,\n",
        "                    inputs=all_parameter_list,\n",
        "                    outputs=[common_parameter_toml, sample_parameter_toml, plus_parameter_toml,  write_files_title]\n",
        "    )\n",
        "    base_model_get_button.click(fn=model_get,\n",
        "                  inputs=[base_model_dir],\n",
        "                  outputs=[base_model_name]\n",
        "    )\n",
        "    vae_model_get_button.click(fn=model_get,\n",
        "                  inputs=[vae_model_dir],\n",
        "                  outputs=[vae_model_name]\n",
        "    )\n",
        "    write_files_button.click(fn=write_files,\n",
        "                inputs=[write_files_dir],\n",
        "                outputs=[write_files_title]\n",
        "    )\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    demo.launch(share=False,inbrowser=False,inline=True,debug=True)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNv6whrNjHPpxmgtHuC8ypu",
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}