feat: finishing v15

allwellll · Apr 10, 2023 · 903d470 · 903d470
1 parent 2b01086
commit 903d470
Showing 1 changed file with 66 additions and 40 deletions.
diff --git a/kohya-trainer.ipynb b/kohya-trainer.ipynb
@@ -18,13 +18,12 @@
         "id": "RL2cTWJ66iz-"
       },
       "source": [
-        "| Notebook Name | Description | Link | Old Commit |\n",
+        "| Notebook Name | Description | Link | V14 |\n",
         "| --- | --- | --- | --- |\n",
-        "| [Kohya LoRA Dreambooth](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-dreambooth.ipynb) | LoRA Training (Dreambooth method) | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-dreambooth.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/bc0892647cb17492a106ad1d05716e091eda13f6/kohya-LoRA-dreambooth.ipynb) | \n",
-        "| [Kohya LoRA Fine-Tuning](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-finetuner.ipynb) | LoRA Training (Fine-tune method) | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-finetuner.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/fb96280233d3434819ba5850b2c968150c4720f7/kohya-LoRA-finetuner.ipynb) | \n",
-        "| [Kohya Trainer](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-trainer.ipynb) | Native Training | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-trainer.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/21ad4942a917d3fd1ad6c03d87d16677b427254b/kohya-trainer.ipynb) | \n",
-        "| [Kohya Dreambooth](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-dreambooth.ipynb) | Dreambooth Training | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-dreambooth.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/9c7f891981bee92cc7690f2094f892c46feb99e2/kohya-dreambooth.ipynb) | \n",
-        "| [Fast Kohya Trainer](https://github.com/Linaqruf/kohya-trainer/blob/main/fast-kohya-trainer.ipynb) `NEW`| Easy 1-click LoRA & Native Training| [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/fast-kohya-trainer.ipynb) |\n",
+        "| [Kohya LoRA Dreambooth](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-dreambooth.ipynb) | LoRA Training (Dreambooth method) | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-dreambooth.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/ff701379c65380c967cd956e4e9e8f6349563878/kohya-LoRA-dreambooth.ipynb) | \n",
+        "| [Kohya LoRA Fine-Tuning](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-finetuner.ipynb) | LoRA Training (Fine-tune method) | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-finetuner.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/ff701379c65380c967cd956e4e9e8f6349563878/kohya-LoRA-finetuner.ipynb) | \n",
+        "| [Kohya Trainer](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-trainer.ipynb) | Native Training | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-trainer.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/ff701379c65380c967cd956e4e9e8f6349563878/kohya-trainer.ipynb) | \n",
+        "| [Kohya Dreambooth](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-dreambooth.ipynb) | Dreambooth Training | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-dreambooth.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/ff701379c65380c967cd956e4e9e8f6349563878/kohya-dreambooth.ipynb) | \n",
         "| [Cagliostro Colab UI](https://github.com/Linaqruf/sd-notebook-collection/blob/main/cagliostro-colab-ui.ipynb) `NEW`| A Customizable Stable Diffusion Web UI| [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/sd-notebook-collection/blob/main/cagliostro-colab-ui.ipynb) | "
       ]
     },
@@ -454,9 +453,17 @@
       "outputs": [],
       "source": [
         "# @title ## 3.1. Locating Train Data Directory\n",
-        "# @markdown Define location of your training data. This cell will also create a fo/content/fine_tune your input.\n",
-        "# @markdown This folder will serve as the /content/fine_tune for scraping, tagging,/content/fine_tunend training in the next cell.\n",
-        "import/content/fine_tuner/content/fine_tunedir = \"/content/fine_tune/train_data\"  # @param {'type' : 'stri/content/fine_tunerain_data_/content/fine_tuneirs(trai/content/fine_tunexist_ok=T/content/fine_tuneYour train data directo/content/fine_tuneata_dir}\")"
+        "# @markdown Define location of your training data. This cell will also create a folder based on your input.\n",
+        "# @markdown This folder will serve as the target folder for scraping, tagging, bucketing, and training in the next cell.\n",
+        "import os\n",
+        "\n",
+        "%store -r\n",
+        "\n",
+        "train_data_dir = \"/content/LoRA/train_data\"  # @param {'type' : 'string'}\n",
+        "%store train_data_dir\n",
+        "\n",
+        "os.makedirs(train_data_dir, exist_ok=True)\n",
+        "print(f\"Your train data directory : {train_data_dir}\")"
       ]
     },
     {
@@ -566,7 +573,7 @@
         "write_tags = True #@param {type: \"boolean\"}\n",
         "additional_arguments = \"--filename /O --no-part\" #@param {type: \"string\"}\n",
         "#@markdown Set `with_aria_2c` to `True` to scrape images using aria2c.\n",
-        "with_aria_2c = True #@param {type: \"boolean\"}\n",
+        "with_aria_2c = False #@param {type: \"boolean\"}\n",
         "\n",
         "tags = prompt.split(',')\n",
         "tags = '+'.join(tags)\n",
@@ -1046,15 +1053,17 @@
         "meta_cap_dict = {\n",
         "    \"_train_data_dir\": train_data_dir,\n",
         "    \"_out_json\": meta_cap,\n",
+        "    \"caption_extension\": \".caption\",\n",
         "}\n",
         "meta_cap_dd_dict = {\n",
         "    \"_train_data_dir\": train_data_dir,\n",
         "    \"_out_json\": meta_cap_dd,\n",
+        "    \"caption_extension\": \".txt\",\n",
         "}\n",
         "merge_metadata_dict = {\n",
         "    \"_train_data_dir\": train_data_dir,\n",
-        "    \"in_json\": meta_cap,\n",
         "    \"_out_json\": meta_cap_dd,\n",
+        "    \"in_json\": meta_cap,\n",
         "}\n",
         "clean_metadata_dict = {\n",
         "    \"_in_json\": meta_cap_dd if os.path.exists(meta_cap_dd) else meta_cap,\n",
@@ -1270,7 +1279,7 @@
         "# @markdown Set `learning_rate` to `1.0` if you use `DAdaptation` optimizer, as it's a [free learning rate](https://github.com/facebookresearch/dadaptation) algorithm.\n",
         "# @markdown You probably need to specify `optimizer_args` for custom optimizer, like using `[\"decouple=true\",\"weight_decay=0.6\"]` for `DAdaptation`.\n",
         "learning_rate = 2e-6  # @param {'type':'number'}\n",
-        "stop_train_text_encoder = -1 #@param {'type':'number'}\n",
+        "train_text_encoder = False  # @param {'type':'boolean'}\n",
         "lr_scheduler = \"constant\"  # @param [\"linear\", \"cosine\", \"cosine_with_restarts\", \"polynomial\", \"constant\", \"constant_with_warmup\", \"adafactor\"] {allow-input: false}\n",
         "lr_warmup_steps = 0  # @param {'type':'number'}\n",
         "# @markdown You can define `num_cycles` value for `cosine_with_restarts` or `power` value for `polynomial` in the field below.\n",
@@ -1282,8 +1291,8 @@
         "if optimizer_args:\n",
         "  print(f\"Optimizer Args :\", optimizer_args)\n",
         "print(\"Learning rate: \", learning_rate)\n",
-        "if stop_train_text_encoder > 0:\n",
-        "  print(f\"Text Encoder training stopped at {stop_train_text_encoder} steps\")\n",
+        "if train_text_encoder:\n",
+        "    print(f\"Train Text Encoder\")\n",
         "print(\"Learning rate warmup steps: \", lr_warmup_steps)\n",
         "print(\"Learning rate Scheduler:\", lr_scheduler)\n",
         "if lr_scheduler == \"cosine_with_restarts\":\n",
@@ -1356,6 +1365,19 @@
         "    },\n",
         "    \"dataset_arguments\": {\n",
         "        \"debug_dataset\": False,\n",
+        "        \"in_json\": in_json,\n",
+        "        \"train_data_dir\": train_data_dir,\n",
+        "        \"dataset_repeats\": dataset_repeats,\n",
+        "        \"shuffle_caption\": True,\n",
+        "        \"keep_tokens\": keep_tokens,\n",
+        "        \"resolution\": resolution,\n",
+        "        \"caption_dropout_rate\": 0,\n",
+        "        \"caption_tag_dropout_rate\": 0,\n",
+        "        \"caption_dropout_every_n_epochs\": 0,\n",
+        "        \"color_aug\": False,\n",
+        "        \"face_crop_aug_range\": None,\n",
+        "        \"token_warmup_min\": 1,\n",
+        "        \"token_warmup_step\": 0,\n",
         "    },\n",
         "    \"training_arguments\": {\n",
         "        \"output_dir\": output_dir,\n",
@@ -1465,7 +1487,7 @@
         "\n",
         "accelerate_args = train(accelerate_conf)\n",
         "train_args = train(train_conf)\n",
-        "final_args = f\"accelerate launch {accelerate_args} train_network.py {train_args}\"\n",
+        "final_args = f\"accelerate launch {accelerate_args} fine_tune.py {train_args}\"\n",
         "\n",
         "os.chdir(repo_dir)\n",
         "!{final_args}"
@@ -1513,7 +1535,6 @@
         "\n",
         "v2 = False  # @param {type:\"boolean\"}\n",
         "v_parameterization = False  # @param {type:\"boolean\"}\n",
-        "instance_prompt = \"\"  # @param {type: \"string\"}\n",
         "prompt = \"masterpiece, best quality, 1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, hoop earrings, jewelry, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt\"  # @param {type: \"string\"}\n",
         "negative = \"lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry\"  # @param {type: \"string\"}\n",
         "model = \"\"  # @param {type: \"string\"}\n",
@@ -1530,7 +1551,7 @@
         "clip_skip = 2  # @param {type: \"slider\", min: 1, max: 40}\n",
         "seed = -1  # @param {type: \"integer\"}\n",
         "\n",
-        "final_prompt = f\"{instance_prompt}, {prompt} --n {negative}\" if instance_prompt else f\"{prompt} --n {negative}\"\n",
+        "final_prompt = f\"{prompt} --n {negative}\"\n",
         "\n",
         "config = {\n",
         "    \"v2\": v2,\n",
@@ -1601,7 +1622,7 @@
         "control_dir = os.path.join(webui_dir, \"models/ControlNet\")\n",
         "\n",
         "webui_models_dir = os.path.join(webui_dir, \"models/Stable-diffusion\")\n",
-        "# webui_lora_dir = os.path.join(webui_dir, \"models/Lora\")\n",
+        "webui_lora_dir = os.path.join(webui_dir, \"models/Lora\")\n",
         "webui_vaes_dir = os.path.join(webui_dir, \"models/VAE\")\n",
         "\n",
         "control_net_max_models_num = 2\n",
@@ -1929,7 +1950,7 @@
         "\n",
         "sd_to_diffusers_dict = {\n",
         "    \"_model_to_load\": model_to_load,\n",
-        "    \"_model_to_save\": model_to_load,\n",
+        "    \"_model_to_save\": model_to_save,\n",
         "    \"v2\": True if v2 else False,\n",
         "    \"v1\": True if not v2 else False,\n",
         "    \"global_step\": global_step,\n",
@@ -1939,7 +1960,7 @@
         "    \"reference_model\": reference_model\n",
         "}\n",
         "\n",
-        "def convert(config):\n",
+        "def convert_dict(config):\n",
         "    args = \"\"\n",
         "    for k, v in config.items():\n",
         "        if k.startswith(\"_\"):\n",
@@ -1956,31 +1977,36 @@
         "    return args\n",
         "\n",
         "def run_script(script_name, script_args):\n",
-        "    !python {script_name} {script_args}}\n",
+        "    !python {script_name} {script_args}\n",
         "\n",
         "def download(output, url, save_dir):\n",
         "    !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d '{save_dir}' -o '{output}' {url}\n",
         "\n",
-        "diffusers_to_sd_args = convert(diffusers_to_sd)\n",
-        "sd_to_diffusers_args = convert(sd_to_diffusers)\n",
+        "diffusers_to_sd_args = convert_dict(diffusers_to_sd_dict)\n",
+        "sd_to_diffusers_args = convert_dict(sd_to_diffusers_dict)\n",
         "\n",
         "if convert == \"diffusers_to_checkpoint\":\n",
-        "    if \".ckpt\" or \".safetensors\" not in model_to_load:\n",
+        "    if model_to_load.endswith((\"ckpt\",\"safetensors\")):\n",
+        "        print(f\"{os.path.basename(model_to_load)} is not in diffusers format\")\n",
+        "    else:\n",
         "        run_script(\"convert_diffusers20_original_sd.py\", diffusers_to_sd_args)\n",
-        "else:    \n",
-        "    run_script(\"convert_diffusers20_original_sd.py\", sd_to_diffusers_args)\n",
-        "\n",
-        "    if feature_extractor:\n",
-        "        save_dir = os.path.join(model_to_save, \"feature_extractor\")\n",
-        "        os.makedirs(save_dir, exist_ok=True)\n",
-        "        output, url = urls[0]\n",
-        "        download(output, url, save_dir)\n",
-        "        \n",
-        "    if safety_checker:\n",
-        "        save_dir = os.path.join(model_to_save, \"safety_checker\")\n",
-        "        os.makedirs(save_dir, exist_ok=True)\n",
-        "        for output, url in urls[1:]:\n",
-        "            download(output, url, save_dir)"
+        "else:\n",
+        "    if not model_to_load.endswith((\"ckpt\",\"safetensors\")):\n",
+        "        print(f\"{os.path.basename(model_to_load)} is not in ckpt/safetensors format\")\n",
+        "    else:     \n",
+        "        run_script(\"convert_diffusers20_original_sd.py\", sd_to_diffusers_args)\n",
+        "\n",
+        "        if feature_extractor:\n",
+        "            save_dir = os.path.join(model_to_save, \"feature_extractor\")\n",
+        "            os.makedirs(save_dir, exist_ok=True)\n",
+        "            output, url = urls[0]\n",
+        "            download(output, url, save_dir)\n",
+        "            \n",
+        "        if safety_checker:\n",
+        "            save_dir = os.path.join(model_to_save, \"safety_checker\")\n",
+        "            os.makedirs(save_dir, exist_ok=True)\n",
+        "            for output, url in urls[1:]:\n",
+        "                download(output, url, save_dir)"
       ]
     },
     {
@@ -2036,7 +2062,7 @@
         "\n",
         "print(f\"Loading model from {model_path}\")\n",
         "\n",
-        "model_path_path = os.path.dirname(model_path)\n",
+        "dir_name = os.path.dirname(model_path)\n",
         "base_name = os.path.basename(model_path)\n",
         "output_name = base_name.split('.')[0]\n",
         "\n",
@@ -2046,7 +2072,7 @@
         "        output_name += suffix\n",
         "        \n",
         "output_name += '-pruned'\n",
-        "output_path = os.path.join(model_path, output_name + ('.ckpt' if config['model_path'].endswith(\".ckpt\") else \".safetensors\"))\n",
+        "output_path = os.path.join(dir_name, output_name + ('.ckpt' if model_path.endswith(\".ckpt\") else \".safetensors\"))\n",
         "\n",
         "args = \"\"\n",
         "for k, v in config.items():\n",