kohya-LoRA-dreambooth.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "slgjeYgd6pWp"
      },
      "source": [
        "![visitors](https://visitor-badge.glitch.me/badge?page_id=linaqruf.lora-dreambooth) [![](https://dcbadge.vercel.app/api/shield/850007095775723532?style=flat)](https://lookup.guru/850007095775723532) [![ko-fi](https://img.shields.io/badge/Support%20me%20on%20Ko--fi-F16061?logo=ko-fi&logoColor=white&style=flat)](https://ko-fi.com/linaqruf) <a href=\"https://saweria.co/linaqruf\"><img alt=\"Saweria\" src=\"https://img.shields.io/badge/Saweria-7B3F00?style=flat&logo=ko-fi&logoColor=white\"/></a>\n",
        "\n",
        "# **Kohya LoRA Dreambooth**\n",
        "A Colab Notebook For LoRA Training (Dreambooth Method)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "D_6E01tRBfQU"
      },
      "source": [
        "| 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/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) | "
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "tTVqCAgSmie4"
      },
      "source": [
        "# I. Install Kohya Trainer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_u3q60di584x"
      },
      "outputs": [],
      "source": [
        "# @title ## 1.1. Install Dependencies\n",
        "# @markdown Clone Kohya Trainer from GitHub and check for updates. Use textbox below if you want to checkout other branch or old commit. Leave it empty to stay the HEAD on main.  This will also install the required libraries.\n",
        "import os\n",
        "import zipfile\n",
        "import shutil\n",
        "import time\n",
        "from subprocess import getoutput\n",
        "from IPython.utils import capture\n",
        "from google.colab import drive\n",
        "\n",
        "%store -r\n",
        "\n",
        "# root_dir\n",
        "root_dir = \"/content\"\n",
        "deps_dir = os.path.join(root_dir, \"deps\")\n",
        "repo_dir = os.path.join(root_dir, \"kohya-trainer\")\n",
        "training_dir = os.path.join(root_dir, \"LoRA\")\n",
        "pretrained_model = os.path.join(root_dir, \"pretrained_model\")\n",
        "vae_dir = os.path.join(root_dir, \"vae\")\n",
        "config_dir = os.path.join(training_dir, \"config\")\n",
        "\n",
        "# repo_dir\n",
        "accelerate_config = os.path.join(repo_dir, \"accelerate_config/config.yaml\")\n",
        "tools_dir = os.path.join(repo_dir, \"tools\")\n",
        "finetune_dir = os.path.join(repo_dir, \"finetune\")\n",
        "\n",
        "for store in [\n",
        "    \"root_dir\",\n",
        "    \"deps_dir\",\n",
        "    \"repo_dir\",\n",
        "    \"training_dir\",\n",
        "    \"pretrained_model\",\n",
        "    \"vae_dir\",\n",
        "    \"accelerate_config\",\n",
        "    \"tools_dir\",\n",
        "    \"finetune_dir\",\n",
        "    \"config_dir\",\n",
        "]:\n",
        "    with capture.capture_output() as cap:\n",
        "        %store {store}\n",
        "        del cap\n",
        "\n",
        "repo_url = \"https://github.com/Linaqruf/kohya-trainer\"\n",
        "bitsandytes_main_py = \"/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py\"\n",
        "branch = \"\"  # @param {type: \"string\"}\n",
        "mount_drive = False  # @param {type: \"boolean\"}\n",
        "verbose = False # @param {type: \"boolean\"}\n",
        "\n",
        "def read_file(filename):\n",
        "    with open(filename, \"r\") as f:\n",
        "        contents = f.read()\n",
        "    return contents\n",
        "\n",
        "\n",
        "def write_file(filename, contents):\n",
        "    with open(filename, \"w\") as f:\n",
        "        f.write(contents)\n",
        "\n",
        "\n",
        "def clone_repo(url):\n",
        "    if not os.path.exists(repo_dir):\n",
        "        os.chdir(root_dir)\n",
        "        !git clone {url} {repo_dir}\n",
        "    else:\n",
        "        os.chdir(repo_dir)\n",
        "        !git pull origin {branch} if branch else !git pull\n",
        "\n",
        "\n",
        "def install_dependencies():\n",
        "    s = getoutput('nvidia-smi')\n",
        "\n",
        "    if 'T4' in s:\n",
        "        !sed -i \"s@cpu@cuda@\" library/model_util.py\n",
        "\n",
        "    !pip install {'-q' if not verbose else ''} --upgrade -r requirements.txt\n",
        "        \n",
        "    from accelerate.utils import write_basic_config\n",
        "\n",
        "    if not os.path.exists(accelerate_config):\n",
        "        write_basic_config(save_location=accelerate_config)\n",
        "\n",
        "\n",
        "def remove_bitsandbytes_message(filename):\n",
        "    welcome_message = \"\"\"\n",
        "def evaluate_cuda_setup():\n",
        "    print('')\n",
        "    print('='*35 + 'BUG REPORT' + '='*35)\n",
        "    print('Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues')\n",
        "    print('For effortless bug reporting copy-paste your error into this form: https://docs.google.com/forms/d/e/1FAIpQLScPB8emS3Thkp66nvqwmjTEgxp8Y9ufuWTzFyr9kJ5AoI47dQ/viewform?usp=sf_link')\n",
        "    print('='*80)\"\"\"\n",
        "\n",
        "    new_welcome_message = \"\"\"\n",
        "def evaluate_cuda_setup():\n",
        "    import os\n",
        "    if 'BITSANDBYTES_NOWELCOME' not in os.environ or str(os.environ['BITSANDBYTES_NOWELCOME']) == '0':\n",
        "        print('')\n",
        "        print('=' * 35 + 'BUG REPORT' + '=' * 35)\n",
        "        print('Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues')\n",
        "        print('For effortless bug reporting copy-paste your error into this form: https://docs.google.com/forms/d/e/1FAIpQLScPB8emS3Thkp66nvqwmjTEgxp8Y9ufuWTzFyr9kJ5AoI47dQ/viewform?usp=sf_link')\n",
        "        print('To hide this message, set the BITSANDBYTES_NOWELCOME variable like so: export BITSANDBYTES_NOWELCOME=1')\n",
        "        print('=' * 80)\"\"\"\n",
        "\n",
        "    contents = read_file(filename)\n",
        "    new_contents = contents.replace(welcome_message, new_welcome_message)\n",
        "    write_file(filename, new_contents)\n",
        "\n",
        "\n",
        "def main():\n",
        "    os.chdir(root_dir)\n",
        "\n",
        "    if mount_drive:\n",
        "        if not os.path.exists(\"/content/drive\"):\n",
        "            drive.mount(\"/content/drive\")\n",
        "\n",
        "    for dir in [\n",
        "        deps_dir, \n",
        "        training_dir, \n",
        "        config_dir, \n",
        "        pretrained_model, \n",
        "        vae_dir\n",
        "    ]:\n",
        "        os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "    clone_repo(repo_url)\n",
        "\n",
        "    if branch:\n",
        "        os.chdir(repo_dir)\n",
        "        status = os.system(f\"git checkout {branch}\")\n",
        "        if status != 0:\n",
        "            raise Exception(\"Failed to checkout branch or commit\")\n",
        "\n",
        "    os.chdir(repo_dir)\n",
        "    \n",
        "    !apt install aria2 {'-qq' if not verbose else ''}\n",
        "\n",
        "    install_dependencies()\n",
        "    time.sleep(3)\n",
        "    \n",
        "    remove_bitsandbytes_message(bitsandytes_main_py)\n",
        "\n",
        "    os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
        "    os.environ[\"BITSANDBYTES_NOWELCOME\"] = \"1\"  \n",
        "    os.environ[\"SAFETENSORS_FAST_GPU\"] = \"1\"\n",
        "\n",
        "    cuda_path = \"/usr/local/cuda-11.8/targets/x86_64-linux/lib/\"\n",
        "    ld_library_path = os.environ.get(\"LD_LIBRARY_PATH\", \"\")\n",
        "    os.environ[\"LD_LIBRARY_PATH\"] = f\"{ld_library_path}:{cuda_path}\"\n",
        "\n",
        "main()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ZmIRAxgEQESm"
      },
      "outputs": [],
      "source": [
        "# @title ## 1.2. Start `File Explorer`\n",
        "# @markdown This will work in real-time even when you run other cells\n",
        "import threading\n",
        "from google.colab import output\n",
        "from imjoy_elfinder.app import main\n",
        "\n",
        "open_in_new_tab = False  # @param {type:\"boolean\"}\n",
        "\n",
        "def start_file_explorer(root_dir=root_dir, port=8765):\n",
        "    try:\n",
        "        main([\"--root-dir=\" + root_dir, \"--port=\" + str(port)])\n",
        "    except Exception as e:\n",
        "        print(\"Error starting file explorer:\", str(e))\n",
        "\n",
        "\n",
        "def open_file_explorer(open_in_new_tab=False, root_dir=root_dir, port=8765):\n",
        "    thread = threading.Thread(target=start_file_explorer, args=[root_dir, port])\n",
        "    thread.start()\n",
        "\n",
        "    if open_in_new_tab:\n",
        "        output.serve_kernel_port_as_window(port)\n",
        "    else:\n",
        "        output.serve_kernel_port_as_iframe(port, height=\"500\")\n",
        "\n",
        "open_file_explorer(open_in_new_tab=open_in_new_tab, root_dir=root_dir, port=8765)\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "3gob9_OwTlwh"
      },
      "source": [
        "# II. Pretrained Model Selection"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8VT6NLv-2u6q"
      },
      "outputs": [],
      "source": [
        "# @title ## 2.1. Download Available Model\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "models = {\n",
        "    \"Animefull-final-pruned\": \"https://huggingface.co/Linaqruf/personal-backup/resolve/main/models/animefull-final-pruned.ckpt\",\n",
        "    \"Anything-v3-1\": \"https://huggingface.co/cag/anything-v3-1/resolve/main/anything-v3-1.safetensors\",\n",
        "    \"AnyLoRA\": \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/AnyLoRA_noVae_fp16-pruned.safetensors\",\n",
        "    \"AnyLoRA-anime-mix\": \"https://huggingface.co/Lykon/AnyLoRA/resolve/main/AAM_Anylora_AnimeMix.safetensors\",\n",
        "    \"AnimePastelDream\": \"https://huggingface.co/Lykon/AnimePastelDream/resolve/main/AnimePastelDream_Soft_noVae_fp16.safetensors\",\n",
        "    \"Chillout-mix\": \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/chillout_mix-pruned.safetensors\",\n",
        "    \"OpenJourney-v4\": \"https://huggingface.co/prompthero/openjourney-v4/resolve/main/openjourney-v4.ckpt\",\n",
        "    \"Stable-Diffusion-v1-5\": \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/stable_diffusion_1_5-pruned.safetensors\",\n",
        "}\n",
        "\n",
        "v2_models = {\n",
        "    \"stable-diffusion-2-1-base\": \"https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-ema-pruned.safetensors\",\n",
        "    \"stable-diffusion-2-1-768v\": \"https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors\",\n",
        "    \"plat-diffusion-v1-3-1\": \"https://huggingface.co/p1atdev/pd-archive/resolve/main/plat-v1-3-1.safetensors\",\n",
        "    \"replicant-v1\": \"https://huggingface.co/gsdf/Replicant-V1.0/resolve/main/Replicant-V1.0.safetensors\",\n",
        "    \"illuminati-diffusion-v1-0\": \"https://huggingface.co/IlluminatiAI/Illuminati_Diffusion_v1.0/resolve/main/illuminati_diffusion_v1.0.safetensors\",\n",
        "    \"illuminati-diffusion-v1-1\": \"https://huggingface.co/4eJIoBek/Illuminati-Diffusion-v1-1/resolve/main/illuminatiDiffusionV1_v11.safetensors\",\n",
        "    \"waifu-diffusion-1-4-anime-e2\": \"https://huggingface.co/hakurei/waifu-diffusion-v1-4/resolve/main/wd-1-4-anime_e2.ckpt\",\n",
        "    \"waifu-diffusion-1-5-e2\": \"https://huggingface.co/waifu-diffusion/wd-1-5-beta2/resolve/main/checkpoints/wd-1-5-beta2-fp32.safetensors\",\n",
        "    \"waifu-diffusion-1-5-e2-aesthetic\": \"https://huggingface.co/waifu-diffusion/wd-1-5-beta2/resolve/main/checkpoints/wd-1-5-beta2-aesthetic-fp32.safetensors\",\n",
        "}\n",
        "\n",
        "installModels = []\n",
        "installv2Models = []\n",
        "\n",
        "# @markdown ### SD1.x model\n",
        "model_name = \"AnyLoRA\"  # @param [\"\", \"Animefull-final-pruned\", \"Anything-v3-1\", \"AnyLoRA\", \"AnyLoRA-anime-mix\", \"AnimePastelDream\", \"Chillout-mix\", \"OpenJourney-v4\", \"Stable-Diffusion-v1-5\"]\n",
        "# @markdown ### SD2.x model\n",
        "v2_model_name = \"\"  # @param [\"\", \"stable-diffusion-2-1-base\", \"stable-diffusion-2-1-768v\", \"plat-diffusion-v1-3-1\", \"replicant-v1\", \"illuminati-diffusion-v1-0\", \"illuminati-diffusion-v1-1\", \"waifu-diffusion-1-4-anime-e2\", \"waifu-diffusion-1-5-e2\", \"waifu-diffusion-1-5-e2-aesthetic\"]\n",
        "\n",
        "if model_name:\n",
        "    model_url = models.get(model_name)\n",
        "    if model_url:\n",
        "        installModels.append((model_name, model_url))\n",
        "\n",
        "if v2_model_name:\n",
        "    v2_model_url = v2_models.get(v2_model_name)\n",
        "    if v2_model_url:\n",
        "        installv2Models.append((v2_model_name, v2_model_url))\n",
        "\n",
        "\n",
        "def install(checkpoint_name, url):\n",
        "    ext = \"ckpt\" if url.endswith(\".ckpt\") else \"safetensors\"\n",
        "\n",
        "    hf_token = \"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\"\n",
        "    user_header = f'\"Authorization: Bearer {hf_token}\"'\n",
        "    !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {pretrained_model} -o {checkpoint_name}.{ext} \"{url}\"\n",
        "\n",
        "\n",
        "def install_checkpoint():\n",
        "    for model in installModels:\n",
        "        install(model[0], model[1])\n",
        "    for v2model in installv2Models:\n",
        "        install(v2model[0], v2model[1])\n",
        "\n",
        "\n",
        "install_checkpoint()\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3LWn6GzNQ4j5"
      },
      "outputs": [],
      "source": [
        "# @title ## 2.2. Download Custom Model\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "# @markdown ### Custom model\n",
        "modelUrls = \"\"  # @param {'type': 'string'}\n",
        "\n",
        "def install(url):\n",
        "    base_name = os.path.basename(url)\n",
        "\n",
        "    if \"drive.google.com\" in url:\n",
        "        os.chdir(pretrained_model)\n",
        "        !gdown --fuzzy {url}\n",
        "    elif \"huggingface.co\" in url:\n",
        "        if \"/blob/\" in url:\n",
        "            url = url.replace(\"/blob/\", \"/resolve/\")\n",
        "        # @markdown Change this part with your own huggingface token if you need to download your private model\n",
        "        hf_token = \"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\"  # @param {type:\"string\"}\n",
        "        user_header = f'\"Authorization: Bearer {hf_token}\"'\n",
        "        !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {pretrained_model} -o {base_name} {url}\n",
        "    else:\n",
        "        !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {pretrained_model} {url}\n",
        "\n",
        "if modelUrls:\n",
        "    urls = modelUrls.split(\",\")\n",
        "    for url in urls:\n",
        "        install(url.strip())\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "qrY4KtfL6Dqp"
      },
      "outputs": [],
      "source": [
        "# @title ## 2.3. Download Available VAE (Optional)\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "vaes = {\n",
        "    \"none\": \"\",\n",
        "    \"anime.vae.pt\": \"https://huggingface.co/Linaqruf/personal-backup/resolve/main/vae/animevae.pt\",\n",
        "    \"waifudiffusion.vae.pt\": \"https://huggingface.co/hakurei/waifu-diffusion-v1-4/resolve/main/vae/kl-f8-anime.ckpt\",\n",
        "    \"stablediffusion.vae.pt\": \"https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt\",\n",
        "}\n",
        "install_vaes = []\n",
        "\n",
        "# @markdown Select one of the VAEs to download, select `none` for not download VAE:\n",
        "vae_name = \"anime.vae.pt\"  # @param [\"none\", \"anime.vae.pt\", \"waifudiffusion.vae.pt\", \"stablediffusion.vae.pt\"]\n",
        "\n",
        "if vae_name in vaes:\n",
        "    vae_url = vaes[vae_name]\n",
        "    if vae_url:\n",
        "        install_vaes.append((vae_name, vae_url))\n",
        "\n",
        "\n",
        "def install(vae_name, url):\n",
        "    hf_token = \"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\"\n",
        "    user_header = f'\"Authorization: Bearer {hf_token}\"'\n",
        "    !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {vae_dir} -o {vae_name} \"{url}\"\n",
        "\n",
        "\n",
        "def install_vae():\n",
        "    for vae in install_vaes:\n",
        "        install(vae[0], vae[1])\n",
        "\n",
        "\n",
        "install_vae()\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "En9UUwGNMRMM"
      },
      "source": [
        "# III. Data Acquisition\n",
        "\n",
        "You have three options for acquiring your dataset:\n",
        "\n",
        "1. Uploading it to Colab's local files.\n",
        "2. Bulk downloading images from Danbooru using the `Simple Booru Scraper`.\n",
        "3. Locating your dataset from Google Drive.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "kh7CeDqK4l3Y"
      },
      "outputs": [],
      "source": [
        "# @title ## 3.1. Locating Train Data Directory\n",
        "# @markdown Define the location of your training data. This cell will also create a folder based on your input. Regularization Images is optional and can be skipped.\n",
        "import os\n",
        "from IPython.utils import capture\n",
        "\n",
        "%store -r\n",
        "\n",
        "train_data_dir = \"/content/LoRA/train_data\"  # @param {type:'string'}\n",
        "reg_data_dir = \"/content/LoRA/reg_data\"  # @param {type:'string'}\n",
        "\n",
        "for dir in [train_data_dir, reg_data_dir]:\n",
        "    if dir:\n",
        "        with capture.capture_output() as cap:\n",
        "            os.makedirs(dir, exist_ok=True)\n",
        "            %store dir\n",
        "            del cap\n",
        "\n",
        "print(f\"Your train data directory : {train_data_dir}\")\n",
        "if reg_data_dir:\n",
        "    print(f\"Your reg data directory : {reg_data_dir}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "t17ZfiMB8GWZ"
      },
      "outputs": [],
      "source": [
        "# @title ## 3.2. Unzip Dataset\n",
        "\n",
        "import os\n",
        "import shutil\n",
        "from pathlib import Path\n",
        "\n",
        "#@title ## Unzip Dataset\n",
        "# @markdown Use this section if your dataset is in a `zip` file and has been uploaded somewhere. This code cell will download your dataset and automatically extract it to the `train_data_dir` if the `unzip_to` variable is empty.\n",
        "zipfile_url = \"\" #@param {type:\"string\"}\n",
        "zipfile_name = \"zipfile.zip\"\n",
        "unzip_to = \"\" #@param {type:\"string\"}\n",
        "\n",
        "hf_token = \"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\"\n",
        "user_header = f'\"Authorization: Bearer {hf_token}\"'\n",
        "\n",
        "if unzip_to:\n",
        "    os.makedirs(unzip_to, exist_ok=True)\n",
        "else:\n",
        "    unzip_to = train_data_dir\n",
        "\n",
        "\n",
        "def download_dataset(url):\n",
        "    if url.startswith(\"/content\"):\n",
        "        return url\n",
        "    elif \"drive.google.com\" in url:\n",
        "        os.chdir(root_dir)\n",
        "        !gdown --fuzzy {url}\n",
        "        return f\"{root_dir}/{zipfile_name}\"\n",
        "    elif \"huggingface.co\" in url:\n",
        "        if \"/blob/\" in url:\n",
        "            url = url.replace(\"/blob/\", \"/resolve/\")\n",
        "        !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {root_dir} -o {zipfile_name} {url}\n",
        "        return f\"{root_dir}/{zipfile_name}\"\n",
        "    else:\n",
        "        !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {root_dir} -o {zipfile_name} {url}\n",
        "        return f\"{root_dir}/{zipfile_name}\"\n",
        "\n",
        "\n",
        "def extract_dataset(zip_file, output_path):\n",
        "    if zip_file.startswith(\"/content\"):\n",
        "        !unzip -j -o {zip_file} -d \"{output_path}\"\n",
        "    else:\n",
        "        !unzip -j -o \"{zip_file}\" -d \"{output_path}\"\n",
        "\n",
        "\n",
        "def remove_files(train_dir, files_to_move):\n",
        "    for filename in os.listdir(train_dir):\n",
        "        file_path = os.path.join(train_dir, filename)\n",
        "        if filename in files_to_move:\n",
        "            if not os.path.exists(file_path):\n",
        "                shutil.move(file_path, training_dir)\n",
        "            else:\n",
        "                os.remove(file_path)\n",
        "\n",
        "\n",
        "zip_file = download_dataset(zipfile_url)\n",
        "extract_dataset(zip_file, unzip_to)\n",
        "os.remove(zip_file)\n",
        "\n",
        "files_to_move = (\n",
        "    \"meta_cap.json\",\n",
        "    \"meta_cap_dd.json\",\n",
        "    \"meta_lat.json\",\n",
        "    \"meta_clean.json\",\n",
        ")\n",
        "\n",
        "remove_files(train_data_dir, files_to_move)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "A0t1dfnU5Xkq"
      },
      "outputs": [],
      "source": [
        "#@title ## 3.3. Image Scraper (Optional)\n",
        "import os\n",
        "import html\n",
        "from IPython.utils import capture\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "#@markdown Use `gallery-dl` to scrape images from an imageboard site. Specify the `prompt(s)` by separating them with commas, e.g., `hito_komoru, touhou`.\n",
        "booru = \"Danbooru\" #@param [\"Danbooru\", \"Gelbooru\", \"Safebooru\"]\n",
        "prompt = \"\" #@param {type: \"string\"}\n",
        "#@markdown  You can also specify a `custom_url` instead of using a predefined site.\n",
        "custom_url = \"\" #@param {type: \"string\"}\n",
        "#@markdown `sub_folder` option can be used to organize the downloaded images into separate folders based on their concept or category.\n",
        "sub_folder = \"\" #@param {type: \"string\"}\n",
        "user_agent = \"gdl/1.24.5\"\n",
        "#@markdown You can limit the number of images to download by using the `--range` option followed by the desired range. For example `1-200`.\n",
        "range = \"1-200\" #@param {type: \"string\"}\n",
        "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",
        "\n",
        "tags = prompt.split(',')\n",
        "tags = '+'.join(tags)\n",
        "\n",
        "replacement_dict = {\" \": \"\", \"(\": \"%28\", \")\": \"%29\", \":\": \"%3a\"}\n",
        "tags = ''.join(replacement_dict.get(c, c) for c in tags)\n",
        "\n",
        "if sub_folder == \"\":\n",
        "    image_dir = train_data_dir\n",
        "elif sub_folder.startswith(\"/content\"):\n",
        "    image_dir = sub_folder\n",
        "else:\n",
        "    image_dir = os.path.join(train_data_dir, sub_folder)\n",
        "    os.makedirs(image_dir, exist_ok=True)\n",
        "\n",
        "if booru == \"Danbooru\":\n",
        "    url = \"https://danbooru.donmai.us/posts?tags={}\".format(tags)\n",
        "elif booru == \"Gelbooru\":\n",
        "    url = \"https://gelbooru.com/index.php?page=post&s=list&tags={}\".format(tags)\n",
        "else:\n",
        "    url = \"https://safebooru.org/index.php?page=post&s=list&tags={}\".format(tags)\n",
        "\n",
        "valid_url = custom_url if custom_url else url\n",
        "\n",
        "def scrape(config):\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "\n",
        "    return args\n",
        "\n",
        "def pre_process_tags(directory):\n",
        "    for item in os.listdir(directory):\n",
        "        item_path = os.path.join(directory, item)\n",
        "        if os.path.isfile(item_path) and item.endswith(\".txt\"):\n",
        "            old_path = item_path\n",
        "            new_file_name = os.path.splitext(os.path.splitext(item)[0])[0] + \".txt\"\n",
        "            new_path = os.path.join(directory, new_file_name)\n",
        "\n",
        "            os.rename(old_path, new_path)\n",
        "\n",
        "            with open(new_path, \"r\") as f:\n",
        "                contents = f.read()\n",
        "\n",
        "            contents = html.unescape(contents)\n",
        "            contents = contents.replace(\"_\", \" \")\n",
        "            contents = \", \".join(contents.split(\"\\n\"))\n",
        "\n",
        "            with open(new_path, \"w\") as f:\n",
        "                f.write(contents)\n",
        "\n",
        "        elif os.path.isdir(item_path):\n",
        "            pre_process_tags(item_path)\n",
        "\n",
        "get_url_config = {\n",
        "    \"get-urls\" : True,\n",
        "    \"range\" : range if range else None,\n",
        "    \"user-agent\" : user_agent\n",
        "}\n",
        "\n",
        "scrape_config = {\n",
        "    \"directory\" : image_dir,\n",
        "    \"write-tags\" : write_tags,\n",
        "    \"range\" : range if range else None,\n",
        "    \"user-agent\" : user_agent\n",
        "}\n",
        "\n",
        "if with_aria_2c:\n",
        "    scraper_text = os.path.join(root_dir, \"scrape_this.txt\")\n",
        "    with capture.capture_output() as cap:\n",
        "        args = scrape(get_url_config)\n",
        "        !gallery-dl \"{valid_url}\" {args} {additional_arguments}\n",
        "    with open(scraper_text, \"w\") as f:\n",
        "        f.write(cap.stdout)\n",
        "\n",
        "    os.chdir(image_dir)\n",
        "    !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -i {scraper_text}\n",
        "\n",
        "else:\n",
        "    args = scrape(scrape_config)\n",
        "    !gallery-dl \"{valid_url}\" {args} {additional_arguments}\n",
        "\n",
        "if write_tags:\n",
        "    pre_process_tags(train_data_dir)\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "T-0qKyEgTchp"
      },
      "source": [
        "# IV. Data Preprocessing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Jz2emq6vWnPu"
      },
      "outputs": [],
      "source": [
        "# @title ## 4.1. Data Cleaning\n",
        "import os\n",
        "import random\n",
        "import concurrent.futures\n",
        "from tqdm import tqdm\n",
        "from PIL import Image\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "test = os.listdir(train_data_dir)\n",
        "# @markdown This section will delete unnecessary files and unsupported media such as `.mp4`, `.webm`, and `.gif`. \n",
        "# @markdown Set the `convert` parameter to convert your transparent dataset with an alpha channel (RGBA) to RGB and give it a white background. \n",
        "convert = False  # @param {type:\"boolean\"}\n",
        "# @markdown You can choose to give it a `random_color` background instead of white by checking the corresponding option.\n",
        "random_color = False  # @param {type:\"boolean\"}\n",
        "# @markdown Use the `recursive` option to preprocess subfolders as well.\n",
        "recursive = False #  @param {type:\"boolean\"}\n",
        " \n",
        "\n",
        "batch_size = 32\n",
        "supported_types = [\n",
        "    \".png\",\n",
        "    \".jpg\",\n",
        "    \".jpeg\",\n",
        "    \".webp\",\n",
        "    \".bmp\",\n",
        "    \".caption\",\n",
        "    \".npz\",\n",
        "    \".txt\",\n",
        "    \".json\",\n",
        "]\n",
        "\n",
        "background_colors = [\n",
        "    (255, 255, 255),\n",
        "    (0, 0, 0),\n",
        "    (255, 0, 0),\n",
        "    (0, 255, 0),\n",
        "    (0, 0, 255),\n",
        "    (255, 255, 0),\n",
        "    (255, 0, 255),\n",
        "    (0, 255, 255),\n",
        "]\n",
        "\n",
        "def clean_directory(directory):\n",
        "    for item in os.listdir(directory):\n",
        "        file_path = os.path.join(directory, item)\n",
        "        if os.path.isfile(file_path):\n",
        "            file_ext = os.path.splitext(item)[1]\n",
        "            if file_ext not in supported_types:\n",
        "                print(f\"Deleting file {item} from {directory}\")\n",
        "                os.remove(file_path)\n",
        "        elif os.path.isdir(file_path) and recursive:\n",
        "            clean_directory(file_path)\n",
        "\n",
        "def process_image(image_path):\n",
        "    img = Image.open(image_path)\n",
        "    img_dir, image_name = os.path.split(image_path)\n",
        "\n",
        "    if img.mode in (\"RGBA\", \"LA\"):\n",
        "        if random_color:\n",
        "            background_color = random.choice(background_colors)\n",
        "        else:\n",
        "            background_color = (255, 255, 255)\n",
        "        bg = Image.new(\"RGB\", img.size, background_color)\n",
        "        bg.paste(img, mask=img.split()[-1])\n",
        "\n",
        "        if image_name.endswith(\".webp\"):\n",
        "            bg = bg.convert(\"RGB\")\n",
        "            new_image_path = os.path.join(img_dir, image_name.replace(\".webp\", \".jpg\"))\n",
        "            bg.save(new_image_path, \"JPEG\")\n",
        "            os.remove(image_path)\n",
        "            print(f\" Converted image: {image_name} to {os.path.basename(new_image_path)}\")\n",
        "        else:\n",
        "            bg.save(image_path, \"PNG\")\n",
        "            print(f\" Converted image: {image_name}\")\n",
        "    else:\n",
        "        if image_name.endswith(\".webp\"):\n",
        "            new_image_path = os.path.join(img_dir, image_name.replace(\".webp\", \".jpg\"))\n",
        "            img.save(new_image_path, \"JPEG\")\n",
        "            os.remove(image_path)\n",
        "            print(f\" Converted image: {image_name} to {os.path.basename(new_image_path)}\")\n",
        "        else:\n",
        "            img.save(image_path, \"PNG\")\n",
        "\n",
        "def find_images(directory):\n",
        "    images = []\n",
        "    for root, _, files in os.walk(directory):\n",
        "        for file in files:\n",
        "            if file.endswith(\".png\") or file.endswith(\".webp\"):\n",
        "                images.append(os.path.join(root, file))\n",
        "    return images\n",
        "\n",
        "clean_directory(train_data_dir)\n",
        "images = find_images(train_data_dir)\n",
        "num_batches = len(images) // batch_size + 1\n",
        "\n",
        "if convert:\n",
        "    with concurrent.futures.ThreadPoolExecutor() as executor:\n",
        "        for i in tqdm(range(num_batches)):\n",
        "            start = i * batch_size\n",
        "            end = start + batch_size\n",
        "            batch = images[start:end]\n",
        "            executor.map(process_image, batch)\n",
        "\n",
        "    print(\"All images have been converted\")\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "qdISafLeyklg"
      },
      "source": [
        "## 4.2. Data Annotation\n",
        "You can choose to train a model using captions. We're using [BLIP](https://huggingface.co/spaces/Salesforce/BLIP) for image captioning and [Waifu Diffusion 1.4 Tagger](https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags) for image tagging similar to Danbooru.\n",
        "- Use BLIP Captioning for: `General Images`\n",
        "- Use Waifu Diffusion 1.4 Tagger V2 for: `Anime and Manga-style Images`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "xvGx2Ikhc8iy"
      },
      "outputs": [],
      "source": [
        "#@title ### 4.2.1. BLIP Captioning\n",
        "#@markdown BLIP is a pre-training framework for unified vision-language understanding and generation, which achieves state-of-the-art results on a wide range of vision-language tasks. It can be used as a tool for image captioning, for example, `astronaut riding a horse in space`. \n",
        "import os\n",
        "\n",
        "os.chdir(finetune_dir)\n",
        "\n",
        "batch_size = 8 #@param {type:'number'}\n",
        "max_data_loader_n_workers = 2 #@param {type:'number'}\n",
        "beam_search = True #@param {type:'boolean'}\n",
        "min_length = 5 #@param {type:\"slider\", min:0, max:100, step:5.0}\n",
        "max_length = 75 #@param {type:\"slider\", min:0, max:100, step:5.0}\n",
        "#@markdown Use the `recursive` option to process subfolders as well, useful for multi-concept training.\n",
        "recursive = False #@param {type:\"boolean\"} \n",
        "#@markdown Debug while captioning, it will print your image file with generated captions.\n",
        "verbose_logging = True #@param {type:\"boolean\"}\n",
        "\n",
        "config = {\n",
        "    \"_train_data_dir\" : train_data_dir,\n",
        "    \"batch_size\" : batch_size,\n",
        "    \"beam_search\" : beam_search,\n",
        "    \"min_length\" : min_length,\n",
        "    \"max_length\" : max_length,\n",
        "    \"debug\" : verbose_logging,\n",
        "    \"caption_extension\" : \".caption\",\n",
        "    \"max_data_loader_n_workers\" : max_data_loader_n_workers,\n",
        "    \"recursive\" : recursive\n",
        "}\n",
        "\n",
        "args = \"\"\n",
        "for k, v in config.items():\n",
        "    if k.startswith(\"_\"):\n",
        "        args += f'\"{v}\" '\n",
        "    elif isinstance(v, str):\n",
        "        args += f'--{k}=\"{v}\" '\n",
        "    elif isinstance(v, bool) and v:\n",
        "        args += f\"--{k} \"\n",
        "    elif isinstance(v, float) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "    elif isinstance(v, int) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "\n",
        "final_args = f\"python make_captions.py {args}\"\n",
        "\n",
        "os.chdir(finetune_dir)\n",
        "!{final_args}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-BdXV7rAy2ag"
      },
      "outputs": [],
      "source": [
        "#@title ### 4.2.2. Waifu Diffusion 1.4 Tagger V2\n",
        "import os\n",
        "%store -r\n",
        "\n",
        "os.chdir(finetune_dir)\n",
        "\n",
        "#@markdown [Waifu Diffusion 1.4 Tagger V2](https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags) is a Danbooru-styled image classification model developed by SmilingWolf. It can also be useful for general image tagging, for example, `1girl, solo, looking_at_viewer, short_hair, bangs, simple_background`.\n",
        "batch_size = 8 #@param {type:'number'}\n",
        "max_data_loader_n_workers = 2 #@param {type:'number'}\n",
        "model = \"SmilingWolf/wd-v1-4-convnextv2-tagger-v2\" #@param [\"SmilingWolf/wd-v1-4-convnextv2-tagger-v2\", \"SmilingWolf/wd-v1-4-swinv2-tagger-v2\", \"SmilingWolf/wd-v1-4-convnext-tagger-v2\", \"SmilingWolf/wd-v1-4-vit-tagger-v2\"]\n",
        "#@markdown Use the `recursive` option to process subfolders as well, useful for multi-concept training.\n",
        "recursive = False #@param {type:\"boolean\"} \n",
        "#@markdown Debug while tagging, it will print your image file with general tags and character tags.\n",
        "verbose_logging = True #@param {type:\"boolean\"}\n",
        "#@markdown Separate `undesired_tags` with comma `(,)` if you want to remove multiple tags, e.g. `1girl,solo,smile`.\n",
        "undesired_tags = \"\" #@param {type:'string'}\n",
        "#@markdown  Adjust `general_threshold` for pruning tags (less tags, less flexible). `character_threshold` is useful if you want to train with character tags, e.g. `hakurei reimu`.\n",
        "general_threshold = 0.35 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "character_threshold = 0.35 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "\n",
        "config = {\n",
        "    \"_train_data_dir\": train_data_dir,\n",
        "    \"batch_size\": batch_size,\n",
        "    \"repo_id\": model,\n",
        "    \"recursive\": recursive,\n",
        "    \"remove_underscore\": True,\n",
        "    \"general_threshold\": general_threshold,\n",
        "    \"character_threshold\": character_threshold,\n",
        "    \"caption_extension\": \".txt\",\n",
        "    \"max_data_loader_n_workers\": max_data_loader_n_workers,\n",
        "    \"debug\": verbose_logging,\n",
        "    \"undesired_tags\": undesired_tags\n",
        "}\n",
        "\n",
        "args = \"\"\n",
        "for k, v in config.items():\n",
        "    if k.startswith(\"_\"):\n",
        "        args += f'\"{v}\" '\n",
        "    elif isinstance(v, str):\n",
        "        args += f'--{k}=\"{v}\" '\n",
        "    elif isinstance(v, bool) and v:\n",
        "        args += f\"--{k} \"\n",
        "    elif isinstance(v, float) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "    elif isinstance(v, int) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "\n",
        "final_args = f\"python tag_images_by_wd14_tagger.py {args}\"\n",
        "\n",
        "os.chdir(finetune_dir)\n",
        "!{final_args}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_mLVURhM9PFE"
      },
      "outputs": [],
      "source": [
        "# @title ### 4.2.3. Custom Caption/Tag\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "# @markdown Add or remove custom tags here. You can refer to this [cheatsheet](https://rentry.org/kohyaminiguide#c-custom-tagscaption) for more information.\n",
        "extension = \".txt\"  # @param [\".txt\", \".caption\"]\n",
        "custom_tag = \"\"  # @param {type:\"string\"}\n",
        "# @markdown Use `sub_folder` option to specify a subfolder for multi-concept training. \n",
        "# @markdown > Specify `--all` to process all subfolders/`recursive`\n",
        "sub_folder = \"\" #@param {type: \"string\"}\n",
        "# @markdown Enable this to append custom tags at the end of lines.\n",
        "append = False  # @param {type:\"boolean\"}\n",
        "# @markdown Enable this if you want to remove captions/tags instead.\n",
        "remove_tag = False  # @param {type:\"boolean\"}\n",
        "recursive = False\n",
        "\n",
        "if sub_folder == \"\":\n",
        "    image_dir = train_data_dir\n",
        "elif sub_folder == \"--all\":\n",
        "    image_dir = train_data_dir\n",
        "    recursive = True\n",
        "elif sub_folder.startswith(\"/content\"):\n",
        "    image_dir = sub_folder\n",
        "else:\n",
        "    image_dir = os.path.join(train_data_dir, sub_folder)\n",
        "    os.makedirs(image_dir, exist_ok=True)\n",
        "\n",
        "def read_file(filename):\n",
        "    with open(filename, \"r\") as f:\n",
        "        contents = f.read()\n",
        "    return contents\n",
        "\n",
        "def write_file(filename, contents):\n",
        "    with open(filename, \"w\") as f:\n",
        "        f.write(contents)\n",
        "\n",
        "def process_tags(filename, custom_tag, append, remove_tag):\n",
        "    contents = read_file(filename)\n",
        "    tags = [tag.strip() for tag in contents.split(',')]\n",
        "    custom_tags = [tag.strip() for tag in custom_tag.split(',')]\n",
        "\n",
        "    for custom_tag in custom_tags:\n",
        "        custom_tag = custom_tag.replace(\"_\", \" \")\n",
        "        if remove_tag:\n",
        "            while custom_tag in tags:\n",
        "                tags.remove(custom_tag)\n",
        "        else:\n",
        "            if custom_tag not in tags:\n",
        "                if append:\n",
        "                    tags.append(custom_tag)\n",
        "                else:\n",
        "                    tags.insert(0, custom_tag)\n",
        "\n",
        "    contents = ', '.join(tags)\n",
        "    write_file(filename, contents)\n",
        "\n",
        "def process_directory(image_dir, tag, append, remove_tag, recursive):\n",
        "    for filename in os.listdir(image_dir):\n",
        "        file_path = os.path.join(image_dir, filename)\n",
        "        \n",
        "        if os.path.isdir(file_path) and recursive:\n",
        "            process_directory(file_path, tag, append, remove_tag, recursive)\n",
        "        elif filename.endswith(extension):\n",
        "            process_tags(file_path, tag, append, remove_tag)\n",
        "\n",
        "tag = custom_tag\n",
        "\n",
        "if not any(\n",
        "    [filename.endswith(extension) for filename in os.listdir(image_dir)]\n",
        "):\n",
        "    for filename in os.listdir(image_dir):\n",
        "        if filename.endswith((\".png\", \".jpg\", \".jpeg\", \".webp\", \".bmp\")):\n",
        "            open(\n",
        "                os.path.join(image_dir, filename.split(\".\")[0] + extension),\n",
        "                \"w\",\n",
        "            ).close()\n",
        "\n",
        "if custom_tag:\n",
        "    process_directory(image_dir, tag, append, remove_tag, recursive)\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "yHNbl3O_NSS0"
      },
      "source": [
        "# V. Training Model\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "H_Q23fUEJhnC"
      },
      "outputs": [],
      "source": [
        "# @title ## 5.1. Model Config\n",
        "from google.colab import drive\n",
        "\n",
        "v2 = False  # @param {type:\"boolean\"}\n",
        "v_parameterization = False  # @param {type:\"boolean\"}\n",
        "project_name = \"\"  # @param {type:\"string\"}\n",
        "if not project_name:\n",
        "    project_name = \"last\"\n",
        "%store project_name\n",
        "pretrained_model_name_or_path = \"/content/pretrained_model/AnyLoRA.safetensors\"  # @param {type:\"string\"}\n",
        "vae = \"\"  # @param {type:\"string\"}\n",
        "output_dir = \"/content/LoRA/output\"  # @param {'type':'string'}\n",
        "\n",
        "# @markdown `output_to_drive` sets default `output_dir` to `/content/drive/MyDrive/LoRA/output`. This will override the `output_dir` variable defined above.\n",
        "output_to_drive = False  # @param {'type':'boolean'}\n",
        "\n",
        "if output_to_drive:\n",
        "    output_dir = \"/content/drive/MyDrive/LoRA/output\"\n",
        "\n",
        "    if not os.path.exists(\"/content/drive\"):\n",
        "        drive.mount(\"/content/drive\")\n",
        "\n",
        "sample_dir = os.path.join(output_dir, \"sample\")\n",
        "for dir in [output_dir, sample_dir]:\n",
        "    os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "print(\"Project Name: \", project_name)\n",
        "print(\"Model Version: Stable Diffusion V1.x\") if not v2 else \"\"\n",
        "print(\"Model Version: Stable Diffusion V2.x\") if v2 and not v_parameterization else \"\"\n",
        "print(\"Model Version: Stable Diffusion V2.x 768v\") if v2 and v_parameterization else \"\"\n",
        "print(\n",
        "    \"Pretrained Model Path: \", pretrained_model_name_or_path\n",
        ") if pretrained_model_name_or_path else print(\"No Pretrained Model path specified.\")\n",
        "print(\"VAE Path: \", vae) if vae else print(\"No VAE path specified.\")\n",
        "print(\"Output Path: \", output_dir)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "G5u_DhFeyJ6R"
      },
      "outputs": [],
      "source": [
        "# @title ## 5.2. Dataset Config\n",
        "import os\n",
        "import toml\n",
        "import glob\n",
        "\n",
        "dataset_repeats = 10  # @param {type:\"number\"}\n",
        "# @markdown `activation_word` is not used in training if you train with captions/tags, but it is still printed to metadata.\n",
        "activation_word = \"mksks\"  # @param {type:\"string\"}\n",
        "caption_extension = \".txt\"  # @param [\"none\", \".txt\", \".caption\"]\n",
        "# @markdown Please refer to `4.2.3. Custom Caption/Tag (Optional)` if you want to append `activation_word` to captions/tags\n",
        "resolution = 512  # @param {type:\"slider\", min:512, max:1024, step:128}\n",
        "flip_aug = False  # @param {type:\"boolean\"}\n",
        "keep_tokens = 0  # @param {type:\"number\"}\n",
        "\n",
        "def parse_folder_name(folder_name, default_num_repeats, default_class_token):\n",
        "    folder_name_parts = folder_name.split(\"_\")\n",
        "\n",
        "    if len(folder_name_parts) == 2:\n",
        "        if folder_name_parts[0].isdigit():\n",
        "            num_repeats = int(folder_name_parts[0])\n",
        "            class_token = folder_name_parts[1].replace(\"_\", \" \")\n",
        "        else:\n",
        "            num_repeats = default_num_repeats\n",
        "            class_token = default_class_token\n",
        "    else:\n",
        "        num_repeats = default_num_repeats\n",
        "        class_token = default_class_token\n",
        "\n",
        "    return num_repeats, class_token\n",
        "\n",
        "def find_image_files(path):\n",
        "    supported_extensions = (\".png\", \".jpg\", \".jpeg\", \".webp\", \".bmp\")\n",
        "    return [file for file in glob.glob(path + '/**/*', recursive=True) if file.lower().endswith(supported_extensions)]\n",
        "\n",
        "def process_data_dir(data_dir, default_num_repeats, default_class_token, is_reg=False):\n",
        "    subsets = []\n",
        "\n",
        "    images = find_image_files(data_dir)\n",
        "    if images:\n",
        "        subsets.append({\n",
        "            \"image_dir\": data_dir,\n",
        "            \"class_tokens\": default_class_token,\n",
        "            \"num_repeats\": default_num_repeats,\n",
        "            **({\"is_reg\": is_reg} if is_reg else {}),\n",
        "        })\n",
        "\n",
        "    for root, dirs, files in os.walk(data_dir):\n",
        "        for folder in dirs:\n",
        "            folder_path = os.path.join(root, folder)\n",
        "            images = find_image_files(folder_path)\n",
        "\n",
        "            if images:\n",
        "                num_repeats, class_token = parse_folder_name(folder, default_num_repeats, default_class_token)\n",
        "\n",
        "                subset = {\n",
        "                    \"image_dir\": folder_path,\n",
        "                    \"class_tokens\": class_token,\n",
        "                    \"num_repeats\": num_repeats,\n",
        "                }\n",
        "\n",
        "                if is_reg:\n",
        "                    subset[\"is_reg\"] = True\n",
        "\n",
        "                subsets.append(subset)\n",
        "\n",
        "    return subsets\n",
        "\n",
        "\n",
        "train_subsets = process_data_dir(train_data_dir, dataset_repeats, activation_word)\n",
        "reg_subsets = process_data_dir(reg_data_dir, dataset_repeats, activation_word, is_reg=True)\n",
        "\n",
        "subsets = train_subsets + reg_subsets\n",
        "\n",
        "config = {\n",
        "    \"general\": {\n",
        "        \"enable_bucket\": True,\n",
        "        \"caption_extension\": caption_extension,\n",
        "        \"shuffle_caption\": True,\n",
        "        \"keep_tokens\": keep_tokens,\n",
        "        \"bucket_reso_steps\": 64,\n",
        "        \"bucket_no_upscale\": False,\n",
        "    },\n",
        "    \"datasets\": [\n",
        "        {\n",
        "            \"resolution\": resolution,\n",
        "            \"min_bucket_reso\": 320 if resolution > 640 else 256,\n",
        "            \"max_bucket_reso\": 1280 if resolution > 640 else 1024,\n",
        "            \"caption_dropout_rate\": 0,\n",
        "            \"caption_tag_dropout_rate\": 0,\n",
        "            \"caption_dropout_every_n_epochs\": 0,\n",
        "            \"flip_aug\": flip_aug,\n",
        "            \"color_aug\": False,\n",
        "            \"face_crop_aug_range\": None,\n",
        "            \"subsets\": subsets,\n",
        "        }\n",
        "    ],\n",
        "}\n",
        "\n",
        "dataset_config = os.path.join(config_dir, \"dataset_config.toml\")\n",
        "\n",
        "for key in config:\n",
        "    if isinstance(config[key], dict):\n",
        "        for sub_key in config[key]:\n",
        "            if config[key][sub_key] == \"\":\n",
        "                config[key][sub_key] = None\n",
        "    elif config[key] == \"\":\n",
        "        config[key] = None\n",
        "\n",
        "config_str = toml.dumps(config)\n",
        "\n",
        "with open(dataset_config, \"w\") as f:\n",
        "    f.write(config_str)\n",
        "\n",
        "print(config_str)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iD5Ecamp4rVW"
      },
      "outputs": [],
      "source": [
        "# @title ## 5.3. LoRA and Optimizer Config\n",
        "\n",
        "# @markdown ### LoRA Config:\n",
        "network_category = \"LoRA\"  # @param [\"LoRA\", \"LoCon\", \"LoCon_Lycoris\", \"LoHa\"]\n",
        "\n",
        "# @markdown Recommended values:\n",
        "\n",
        "# @markdown | network_category | network_dim | network_alpha | conv_dim | conv_alpha |\n",
        "# @markdown | :---: | :---: | :---: | :---: | :---: |\n",
        "# @markdown | LoRA | 32 | 1 | - | - |\n",
        "# @markdown | LoCon | 16 | 8 | 8 | 1 |\n",
        "# @markdown | LoHa | 8 | 4 | 4 | 1 |\n",
        "\n",
        "# @markdown - Note that `dropout` and `cp_decomposition` are not available in this notebook.\n",
        "\n",
        "# @markdown `conv_dim` and `conv_alpha` are needed to train `LoCon` and `LoHa`; skip them if you are training normal `LoRA`. However, when in doubt, set `dim = alpha`.\n",
        "conv_dim = 32  # @param {'type':'number'}\n",
        "conv_alpha = 16  # @param {'type':'number'}\n",
        "# @markdown It's recommended not to set `network_dim` and `network_alpha` higher than 64, especially for `LoHa`.\n",
        "# @markdown If you want to use a higher value for `dim` or `alpha`, consider using a higher learning rate, as models with higher dimensions tend to learn faster.\n",
        "network_dim = 32  # @param {'type':'number'}\n",
        "network_alpha = 16  # @param {'type':'number'}\n",
        "# @markdown You can specify this field for resume training.\n",
        "network_weight = \"\"  # @param {'type':'string'}\n",
        "network_module = \"lycoris.kohya\" if network_category in [\"LoHa\", \"LoCon_Lycoris\"] else \"networks.lora\"\n",
        "network_args = \"\" if network_category == \"LoRA\" else [\n",
        "    f\"conv_dim={conv_dim}\", f\"conv_alpha={conv_alpha}\",\n",
        "    ]\n",
        "# @markdown ### <br>Optimizer Config:\n",
        "# @markdown `NEW` Gamma for reducing the weight of high-loss timesteps. Lower numbers have a stronger effect. The paper recommends 5. Read the paper [here](https://arxiv.org/abs/2303.09556).\n",
        "min_snr_gamma = -1 #@param {type:\"number\"}\n",
        "# @markdown `AdamW8bit` was the old `--use_8bit_adam`.\n",
        "optimizer_type = \"AdamW8bit\"  # @param [\"AdamW\", \"AdamW8bit\", \"Lion\", \"SGDNesterov\", \"SGDNesterov8bit\", \"DAdaptation\", \"AdaFactor\"]\n",
        "# @markdown Additional arguments for optimizer, e.g: `[\"decouple=True\",\"weight_decay=0.6\"]`\n",
        "optimizer_args = \"\"  # @param {'type':'string'}\n",
        "# @markdown Set `unet_lr` to `1.0` if you use `DAdaptation` optimizer, because it's a [free learning rate](https://github.com/facebookresearch/dadaptation) algorithm.\n",
        "# @markdown However, it is recommended to set `text_encoder_lr = 0.5 * unet_lr`.\n",
        "# @markdown Also, you don't need to specify `learning_rate` value if both `unet_lr` and `text_encoder_lr` are defined.\n",
        "train_unet = True  # @param {'type':'boolean'}\n",
        "unet_lr = 1e-4  # @param {'type':'number'}\n",
        "train_text_encoder = True  # @param {'type':'boolean'}\n",
        "text_encoder_lr = 5e-5  # @param {'type':'number'}\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",
        "lr_scheduler_num_cycles = 0  # @param {'type':'number'}\n",
        "lr_scheduler_power = 0  # @param {'type':'number'}\n",
        "\n",
        "if network_category == \"LoHa\":\n",
        "  network_args.append(\"algo=loha\")\n",
        "elif network_category == \"LoCon_Lycoris\":\n",
        "  network_args.append(\"algo=lora\")\n",
        "\n",
        "print(\"- LoRA Config:\")\n",
        "print(f\"  - Min-SNR Weighting: {min_snr_gamma}\") if not min_snr_gamma == -1 else \"\"\n",
        "print(f\"  - Loading network module: {network_module}\")\n",
        "if not network_category == \"LoRA\":\n",
        "  print(f\"  - network args: {network_args}\")\n",
        "print(f\"  - {network_module} linear_dim set to: {network_dim}\")\n",
        "print(f\"  - {network_module} linear_alpha set to: {network_alpha}\")\n",
        "if not network_category == \"LoRA\":\n",
        "  print(f\"  - {network_module} conv_dim set to: {conv_dim}\")\n",
        "  print(f\"  - {network_module} conv_alpha set to: {conv_alpha}\")\n",
        "\n",
        "if not network_weight:\n",
        "    print(\"  - No LoRA weight loaded.\")\n",
        "else:\n",
        "    if os.path.exists(network_weight):\n",
        "        print(f\"  - Loading LoRA weight: {network_weight}\")\n",
        "    else:\n",
        "        print(f\"  - {network_weight} does not exist.\")\n",
        "        network_weight = \"\"\n",
        "\n",
        "print(\"- Optimizer Config:\")\n",
        "print(f\"  - Additional network category: {network_category}\")\n",
        "print(f\"  - Using {optimizer_type} as Optimizer\")\n",
        "if optimizer_args:\n",
        "    print(f\"  - Optimizer Args: {optimizer_args}\")\n",
        "if train_unet and train_text_encoder:\n",
        "    print(\"  - Train UNet and Text Encoder\")\n",
        "    print(f\"    - UNet learning rate: {unet_lr}\")\n",
        "    print(f\"    - Text encoder learning rate: {text_encoder_lr}\")\n",
        "if train_unet and not train_text_encoder:\n",
        "    print(\"  - Train UNet only\")\n",
        "    print(f\"    - UNet learning rate: {unet_lr}\")\n",
        "if train_text_encoder and not train_unet:\n",
        "    print(\"  - Train Text Encoder only\")\n",
        "    print(f\"    - Text encoder learning rate: {text_encoder_lr}\")\n",
        "print(f\"  - Learning rate warmup steps: {lr_warmup_steps}\")\n",
        "print(f\"  - Learning rate Scheduler: {lr_scheduler}\")\n",
        "if lr_scheduler == \"cosine_with_restarts\":\n",
        "    print(f\"  - lr_scheduler_num_cycles: {lr_scheduler_num_cycles}\")\n",
        "elif lr_scheduler == \"polynomial\":\n",
        "    print(f\"  - lr_scheduler_power: {lr_scheduler_power}\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-Z4w3lfFKLjr"
      },
      "outputs": [],
      "source": [
        "# @title ## 5.4. Training Config\n",
        "\n",
        "import toml\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "lowram = True  # @param {type:\"boolean\"}\n",
        "enable_sample_prompt = True  # @param {type:\"boolean\"}\n",
        "sampler = \"ddim\"  # @param [\"ddim\", \"pndm\", \"lms\", \"euler\", \"euler_a\", \"heun\", \"dpm_2\", \"dpm_2_a\", \"dpmsolver\",\"dpmsolver++\", \"dpmsingle\", \"k_lms\", \"k_euler\", \"k_euler_a\", \"k_dpm_2\", \"k_dpm_2_a\"]\n",
        "noise_offset = 0.0  # @param {type:\"number\"}\n",
        "num_epochs = 10  # @param {type:\"number\"}\n",
        "vae_batch_size = 4  # @param {type:\"number\"}\n",
        "train_batch_size = 6  # @param {type:\"number\"}\n",
        "mixed_precision = \"fp16\"  # @param [\"no\",\"fp16\",\"bf16\"] {allow-input: false}\n",
        "save_precision = \"fp16\"  # @param [\"float\", \"fp16\", \"bf16\"] {allow-input: false}\n",
        "save_n_epochs_type = \"save_every_n_epochs\"  # @param [\"save_every_n_epochs\", \"save_n_epoch_ratio\"] {allow-input: false}\n",
        "save_n_epochs_type_value = 1  # @param {type:\"number\"}\n",
        "save_model_as = \"safetensors\"  # @param [\"ckpt\", \"pt\", \"safetensors\"] {allow-input: false}\n",
        "max_token_length = 225  # @param {type:\"number\"}\n",
        "clip_skip = 2  # @param {type:\"number\"}\n",
        "gradient_checkpointing = False  # @param {type:\"boolean\"}\n",
        "gradient_accumulation_steps = 1  # @param {type:\"number\"}\n",
        "seed = -1  # @param {type:\"number\"}\n",
        "logging_dir = \"/content/LoRA/logs\"\n",
        "prior_loss_weight = 1.0\n",
        "\n",
        "os.chdir(repo_dir)\n",
        "\n",
        "sample_str = f\"\"\"\n",
        "  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 \\\n",
        "  --n 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 \\\n",
        "  --w 512 \\\n",
        "  --h 768 \\\n",
        "  --l 7 \\\n",
        "  --s 28    \n",
        "\"\"\"\n",
        "\n",
        "config = {\n",
        "    \"model_arguments\": {\n",
        "        \"v2\": v2,\n",
        "        \"v_parameterization\": v_parameterization if v2 and v_parameterization else False,\n",
        "        \"pretrained_model_name_or_path\": pretrained_model_name_or_path,\n",
        "        \"vae\": vae,\n",
        "    },\n",
        "    \"additional_network_arguments\": {\n",
        "        \"no_metadata\": False,\n",
        "        \"unet_lr\": float(unet_lr) if train_unet else None,\n",
        "        \"text_encoder_lr\": float(text_encoder_lr) if train_text_encoder else None,\n",
        "        \"network_weights\": network_weight,\n",
        "        \"network_module\": network_module,\n",
        "        \"network_dim\": network_dim,\n",
        "        \"network_alpha\": network_alpha,\n",
        "        \"network_args\": network_args,\n",
        "        \"network_train_unet_only\": True if train_unet and not train_text_encoder else False,\n",
        "        \"network_train_text_encoder_only\": True if train_text_encoder and not train_unet else False,\n",
        "        \"training_comment\": None,\n",
        "    },\n",
        "    \"optimizer_arguments\": {\n",
        "        \"min_snr_gamma\": min_snr_gamma if not min_snr_gamma == -1 else None,\n",
        "        \"optimizer_type\": optimizer_type,\n",
        "        \"learning_rate\": unet_lr,\n",
        "        \"max_grad_norm\": 1.0,\n",
        "        \"optimizer_args\": eval(optimizer_args) if optimizer_args else None,\n",
        "        \"lr_scheduler\": lr_scheduler,\n",
        "        \"lr_warmup_steps\": lr_warmup_steps,\n",
        "        \"lr_scheduler_num_cycles\": lr_scheduler_num_cycles if lr_scheduler == \"cosine_with_restarts\" else None,\n",
        "        \"lr_scheduler_power\": lr_scheduler_power if lr_scheduler == \"polynomial\" else None,\n",
        "    },\n",
        "    \"dataset_arguments\": {\n",
        "        \"cache_latents\": True,\n",
        "        \"debug_dataset\": False,\n",
        "        \"vae_batch_size\": vae_batch_size,\n",
        "    },\n",
        "    \"training_arguments\": {\n",
        "        \"output_dir\": output_dir,\n",
        "        \"output_name\": project_name,\n",
        "        \"save_precision\": save_precision,\n",
        "        \"save_every_n_epochs\": save_n_epochs_type_value if save_n_epochs_type == \"save_every_n_epochs\" else None,\n",
        "        \"save_n_epoch_ratio\": save_n_epochs_type_value if save_n_epochs_type == \"save_n_epoch_ratio\" else None,\n",
        "        \"save_last_n_epochs\": None,\n",
        "        \"save_state\": None,\n",
        "        \"save_last_n_epochs_state\": None,\n",
        "        \"resume\": None,\n",
        "        \"train_batch_size\": train_batch_size,\n",
        "        \"max_token_length\": 225,\n",
        "        \"mem_eff_attn\": False,\n",
        "        \"xformers\": True,\n",
        "        \"max_train_epochs\": num_epochs,\n",
        "        \"max_data_loader_n_workers\": 8,\n",
        "        \"persistent_data_loader_workers\": True,\n",
        "        \"seed\": seed if seed > 0 else None,\n",
        "        \"gradient_checkpointing\": gradient_checkpointing,\n",
        "        \"gradient_accumulation_steps\": gradient_accumulation_steps,\n",
        "        \"mixed_precision\": mixed_precision,\n",
        "        \"clip_skip\": clip_skip if not v2 else None,\n",
        "        \"logging_dir\": logging_dir,\n",
        "        \"log_prefix\": project_name,\n",
        "        \"noise_offset\": noise_offset if noise_offset > 0 else None,\n",
        "        \"lowram\": lowram,\n",
        "    },\n",
        "    \"sample_prompt_arguments\": {\n",
        "        \"sample_every_n_steps\": None,\n",
        "        \"sample_every_n_epochs\": 1 if enable_sample_prompt else 999999,\n",
        "        \"sample_sampler\": sampler,\n",
        "    },\n",
        "    \"dreambooth_arguments\": {\n",
        "        \"prior_loss_weight\": 1.0,\n",
        "    },\n",
        "    \"saving_arguments\": {\n",
        "        \"save_model_as\": save_model_as\n",
        "    },\n",
        "}\n",
        "\n",
        "config_path = os.path.join(config_dir, \"config_file.toml\")\n",
        "prompt_path = os.path.join(config_dir, \"sample_prompt.txt\")\n",
        "\n",
        "for key in config:\n",
        "    if isinstance(config[key], dict):\n",
        "        for sub_key in config[key]:\n",
        "            if config[key][sub_key] == \"\":\n",
        "                config[key][sub_key] = None\n",
        "    elif config[key] == \"\":\n",
        "        config[key] = None\n",
        "\n",
        "config_str = toml.dumps(config)\n",
        "\n",
        "def write_file(filename, contents):\n",
        "    with open(filename, \"w\") as f:\n",
        "        f.write(contents)\n",
        "\n",
        "write_file(config_path, config_str)\n",
        "write_file(prompt_path, sample_str)\n",
        "    \n",
        "print(config_str)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "p_SHtbFwHVl1"
      },
      "outputs": [],
      "source": [
        "#@title ## 5.5. Start Training\n",
        "\n",
        "#@markdown Check your config here if you want to edit something: \n",
        "#@markdown - `sample_prompt` : /content/LoRA/config/sample_prompt.txt\n",
        "#@markdown - `config_file` : /content/LoRA/config/config_file.toml\n",
        "#@markdown - `dataset_config` : /content/LoRA/config/dataset_config.toml\n",
        "\n",
        "#@markdown Generated sample can be seen here: /content/LoRA/output/sample\n",
        "\n",
        "#@markdown You can import config from another session if you want.\n",
        "sample_prompt = \"/content/LoRA/config/sample_prompt.txt\" #@param {type:'string'}\n",
        "config_file = \"/content/LoRA/config/config_file.toml\" #@param {type:'string'}\n",
        "dataset_config = \"/content/LoRA/config/dataset_config.toml\" #@param {type:'string'}\n",
        "\n",
        "accelerate_conf = {\n",
        "    \"config_file\" : accelerate_config,\n",
        "    \"num_cpu_threads_per_process\" : 1,\n",
        "}\n",
        "\n",
        "train_conf = {\n",
        "    \"sample_prompts\" : sample_prompt,\n",
        "    \"dataset_config\" : dataset_config,\n",
        "    \"config_file\" : config_file\n",
        "}\n",
        "\n",
        "def train(config):\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "\n",
        "    return args\n",
        "\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",
        "\n",
        "os.chdir(repo_dir)\n",
        "!{final_args}"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "reMcN0bM_o53"
      },
      "source": [
        "# VI. Testing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "TLSQslfFcQde"
      },
      "outputs": [],
      "source": [
        "# @title ## 6.1. Visualize loss graph (Optional)\n",
        "training_logs_path = \"/content/LoRA/logs\"  # @param {type : \"string\"}\n",
        "\n",
        "os.chdir(repo_dir)\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir {training_logs_path}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Rt7CKCog_4tm"
      },
      "outputs": [],
      "source": [
        "# @title ## 6.2. Interrogating LoRA Weights\n",
        "# @markdown Now you can check if your LoRA trained properly.\n",
        "import os\n",
        "import torch\n",
        "import json\n",
        "from safetensors.torch import load_file\n",
        "from safetensors.torch import safe_open\n",
        "\n",
        "# @markdown If you used `clip_skip = 2` during training, the values of `lora_te_text_model_encoder_layers_11_*` will be `0.0`, this is normal. These layers are not trained at this value of `Clip Skip`.\n",
        "network_weight = \"\"  # @param {'type':'string'}\n",
        "verbose = False  # @param {type:\"boolean\"}\n",
        "\n",
        "def is_safetensors(path):\n",
        "    return os.path.splitext(path)[1].lower() == \".safetensors\"\n",
        "\n",
        "def load_weight_data(file_path):\n",
        "    if is_safetensors(file_path):\n",
        "        return load_file(file_path)\n",
        "    else:\n",
        "        return torch.load(file_path, map_location=\"cuda\")\n",
        "\n",
        "def extract_lora_weights(weight_data):\n",
        "    lora_weights = [\n",
        "        (key, weight_data[key])\n",
        "        for key in weight_data.keys()\n",
        "        if \"lora_up\" in key or \"lora_down\" in key\n",
        "    ]\n",
        "    return lora_weights\n",
        "\n",
        "def print_lora_weight_stats(lora_weights):\n",
        "    print(f\"Number of LoRA modules: {len(lora_weights)}\")\n",
        "\n",
        "    for key, value in lora_weights:\n",
        "        value = value.to(torch.float32)\n",
        "        print(f\"{key}, {torch.mean(torch.abs(value))}, {torch.min(torch.abs(value))}\")\n",
        "\n",
        "def print_metadata(file_path):\n",
        "    if is_safetensors(file_path):\n",
        "        with safe_open(file_path, framework=\"pt\") as f:\n",
        "            metadata = f.metadata()\n",
        "        if metadata is not None:\n",
        "            print(f\"\\nLoad metadata for: {file_path}\")\n",
        "            print(json.dumps(metadata, indent=4))\n",
        "    else:\n",
        "        print(\"No metadata saved, your model is not in safetensors format\")\n",
        "\n",
        "def main(file_path, verbose: bool):\n",
        "    weight_data = load_weight_data(file_path)\n",
        "\n",
        "    if verbose:\n",
        "        lora_weights = extract_lora_weights(weight_data)\n",
        "        print_lora_weight_stats(lora_weights)\n",
        "\n",
        "    print_metadata(file_path)\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main(network_weight, verbose)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "FKBrTDPrcNjP"
      },
      "outputs": [],
      "source": [
        "# @title ## 6.3. Inference\n",
        "%store -r\n",
        "\n",
        "# @markdown ### LoRA Config\n",
        "# @markdown Currently, `LoHa` and `LoCon_Lycoris` are not supported. Please run `Portable Web UI` instead\n",
        "network_weight = \"\"  # @param {'type':'string'}\n",
        "network_mul = 0.7  # @param {type:\"slider\", min:-1, max:2, step:0.05}\n",
        "network_module = \"networks.lora\"\n",
        "network_args = \"\"\n",
        "\n",
        "# @markdown ### <br> General Config\n",
        "v2 = False  # @param {type:\"boolean\"}\n",
        "v_parameterization = False  # @param {type:\"boolean\"}\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 = \"/content/pretrained_model/AnyLoRA.safetensors\"  # @param {type: \"string\"}\n",
        "vae = \"\"  # @param {type: \"string\"}\n",
        "outdir = \"/content/tmp\"  # @param {type: \"string\"}\n",
        "scale = 7  # @param {type: \"slider\", min: 1, max: 40}\n",
        "sampler = \"ddim\"  # @param [\"ddim\", \"pndm\", \"lms\", \"euler\", \"euler_a\", \"heun\", \"dpm_2\", \"dpm_2_a\", \"dpmsolver\",\"dpmsolver++\", \"dpmsingle\", \"k_lms\", \"k_euler\", \"k_euler_a\", \"k_dpm_2\", \"k_dpm_2_a\"]\n",
        "steps = 28  # @param {type: \"slider\", min: 1, max: 100}\n",
        "precision = \"fp16\"  # @param [\"fp16\", \"bf16\"] {allow-input: false}\n",
        "width = 512  # @param {type: \"integer\"}\n",
        "height = 768  # @param {type: \"integer\"}\n",
        "images_per_prompt = 4  # @param {type: \"integer\"}\n",
        "batch_size = 4  # @param {type: \"integer\"}\n",
        "clip_skip = 2  # @param {type: \"slider\", min: 1, max: 40}\n",
        "seed = -1  # @param {type: \"integer\"}\n",
        "\n",
        "final_prompt = f\"{prompt} --n {negative}\"\n",
        "\n",
        "config = {\n",
        "    \"v2\": v2,\n",
        "    \"v_parameterization\": v_parameterization,\n",
        "    \"network_module\": network_module,\n",
        "    \"network_weight\": network_weight,\n",
        "    \"network_mul\": float(network_mul),\n",
        "    \"network_args\": eval(network_args) if network_args else None,\n",
        "    \"ckpt\": model,\n",
        "    \"outdir\": outdir,\n",
        "    \"xformers\": True,\n",
        "    \"vae\": vae if vae else None,\n",
        "    \"fp16\": True,\n",
        "    \"W\": width,\n",
        "    \"H\": height,\n",
        "    \"seed\": seed if seed > 0 else None,\n",
        "    \"scale\": scale,\n",
        "    \"sampler\": sampler,\n",
        "    \"steps\": steps,\n",
        "    \"max_embeddings_multiples\": 3,\n",
        "    \"batch_size\": batch_size,\n",
        "    \"images_per_prompt\": images_per_prompt,\n",
        "    \"clip_skip\": clip_skip if not v2 else None,\n",
        "    \"prompt\": final_prompt,\n",
        "}\n",
        "\n",
        "args = \"\"\n",
        "for k, v in config.items():\n",
        "    if k.startswith(\"_\"):\n",
        "        args += f'\"{v}\" '\n",
        "    elif isinstance(v, str):\n",
        "        args += f'--{k}=\"{v}\" '\n",
        "    elif isinstance(v, bool) and v:\n",
        "        args += f\"--{k} \"\n",
        "    elif isinstance(v, float) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "    elif isinstance(v, int) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "\n",
        "final_args = f\"python gen_img_diffusers.py {args}\"\n",
        "\n",
        "os.chdir(repo_dir)\n",
        "!{final_args}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "dglQfu3A0oaj"
      },
      "outputs": [],
      "source": [
        "#@title ## 6.4. Launch Portable Web UI\n",
        "import os\n",
        "import random\n",
        "import shutil\n",
        "import zipfile\n",
        "import time\n",
        "import json\n",
        "from google.colab import drive\n",
        "from datetime import timedelta\n",
        "from subprocess import getoutput\n",
        "from IPython.display import clear_output, display, HTML\n",
        "from IPython.utils import capture\n",
        "from tqdm import tqdm\n",
        "\n",
        "webui_dir = os.path.join(root_dir, \"stable-diffusion-webui\")\n",
        "tmp_dir = os.path.join(root_dir, \"tmp\")\n",
        "patches_dir = os.path.join(root_dir, \"patches\")\n",
        "deps_dir = os.path.join(root_dir, \"deps\")\n",
        "extensions_dir = os.path.join(webui_dir, \"extensions\")\n",
        "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_vaes_dir = os.path.join(webui_dir, \"models/VAE\")\n",
        "\n",
        "control_net_max_models_num = 2\n",
        "theme = \"ogxBGreen\"\n",
        "\n",
        "default_prompt = \"masterpiece, best quality,\"\n",
        "default_neg_prompt = \"(worst quality, low quality:1.4)\"\n",
        "default_sampler = \"DPM++ 2M Karras\"\n",
        "default_steps = 20\n",
        "default_width = 512\n",
        "default_height = 768\n",
        "default_denoising_strength = 0.55\n",
        "default_cfg_scale = 7\n",
        "\n",
        "config_file = os.path.join(webui_dir, \"config.json\")\n",
        "ui_config_file = os.path.join(webui_dir, \"ui-config.json\")\n",
        "webui_style_path = os.path.join(webui_dir, \"style.css\")\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "for dir in [patches_dir, deps_dir]:\n",
        "    os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "package_url = [\n",
        "    f\"https://huggingface.co/Linaqruf/fast-repo/resolve/main/anapnoe-webui.tar.lz4\",\n",
        "    f\"https://huggingface.co/Linaqruf/fast-repo/resolve/main/anapnoe-webui-deps.tar.lz4\",\n",
        "    f\"https://huggingface.co/Linaqruf/fast-repo/resolve/main/anapnoe-webui-cache.tar.lz4\",\n",
        "]\n",
        "\n",
        "def pre_download(desc):\n",
        "    for package in tqdm(package_url, desc=desc):\n",
        "        with capture.capture_output() as cap:\n",
        "            package_name = os.path.basename(package)\n",
        "            !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {root_dir} -o {package_name} {package}\n",
        "            if package_name == f\"anapnoe-webui-deps.tar.lz4\":\n",
        "                !tar -xI lz4 -f {package_name} --overwrite-dir --directory=/usr/local/lib/python3.10/dist-packages/\n",
        "            else:\n",
        "                !tar -xI lz4 -f {package_name} --directory=/\n",
        "            os.remove(package_name)\n",
        "            del cap\n",
        "\n",
        "    if os.path.exists(\"/usr/local/lib/python3.10/dist-packages/ffmpy-0.3.0.dist-info\"):\n",
        "        shutil.rmtree(\"/usr/local/lib/python3.10/dist-packages/ffmpy-0.3.0.dist-info\")\n",
        "\n",
        "    s = getoutput(\"nvidia-smi\")\n",
        "    with capture.capture_output() as cap:\n",
        "        if not \"T4\" in s:\n",
        "            !pip uninstall -y xformers\n",
        "            !pip install -q xformers==0.0.18 triton\n",
        "        del cap\n",
        "\n",
        "\n",
        "def read_config(filename):\n",
        "    if filename.endswith(\".json\"):\n",
        "        with open(filename, \"r\") as f:\n",
        "          config = json.load(f)\n",
        "    else:\n",
        "        with open(filename, 'r') as f:\n",
        "          config = f.read()\n",
        "    return config\n",
        "\n",
        "\n",
        "def write_config(filename, config):\n",
        "    if filename.endswith(\".json\"):\n",
        "        with open(filename, \"w\") as f:\n",
        "            json.dump(config, f, indent=4)\n",
        "    else:\n",
        "        with open(filename, 'w', encoding=\"utf-8\") as f:\n",
        "            f.write(config)\n",
        "\n",
        "\n",
        "def open_theme(filename):\n",
        "    themes_folder = os.path.join(webui_dir, \"extensions-builtin/sd_theme_editor/themes\")\n",
        "    themes_file = os.path.join(themes_folder, f\"{filename}.css\")\n",
        "    webui_style_path = os.path.join(webui_dir, \"style.css\")\n",
        "\n",
        "    style_config = read_config(webui_style_path)\n",
        "    style_css_contents = style_config.split(\"/*BREAKPOINT_CSS_CONTENT*/\")[1]\n",
        "\n",
        "    theme_config = read_config(themes_file)\n",
        "    style_data = \":host{\" + theme_config + \"}\" + \"/*BREAKPOINT_CSS_CONTENT*/\" + style_css_contents\n",
        "    write_config(webui_style_path, style_data)\n",
        "\n",
        "\n",
        "def change_config(filename):\n",
        "    config = read_config(filename)\n",
        "    if not \"stable-diffusion-webui\" in config[\"disabled_extensions\"]:\n",
        "        config[\"disabled_extensions\"].append(\"stable-diffusion-webui\")\n",
        "    config[\"outdir_txt2img_samples\"] = os.path.join(tmp_dir, \"outputs/txt2img-images\")\n",
        "    config[\"outdir_img2img_samples\"] = os.path.join(tmp_dir, \"outputs/img2img-images\")\n",
        "    config[\"outdir_extras_samples\"] = os.path.join(tmp_dir, \"outputs/extras-images\")\n",
        "    config[\"outdir_txt2img_grids\"] = os.path.join(tmp_dir, \"outputs/txt2img-grids\")\n",
        "    config[\"outdir_img2img_grids\"] = os.path.join(tmp_dir, \"outputs/img2img-grids\")\n",
        "    config[\"outdir_save\"] = os.path.join(tmp_dir, \"log/images\")\n",
        "    config[\"control_net_max_models_num\"] = control_net_max_models_num\n",
        "    config[\"control_net_models_path\"] = control_dir\n",
        "    config[\"control_net_allow_script_control\"] = True\n",
        "    config[\"additional_networks_extra_lora_path\"] = webui_lora_dir\n",
        "    config[\"CLIP_stop_at_last_layers\"] = 2\n",
        "    config[\"eta_noise_seed_delta\"] = 0\n",
        "    config[\"show_progress_every_n_steps\"] = 10\n",
        "    config[\"show_progressbar\"] = True\n",
        "    config[\"quicksettings\"] = \"sd_model_checkpoint, sd_vae, CLIP_stop_at_last_layers, use_old_karras_scheduler_sigmas, always_discard_next_to_last_sigma\"\n",
        "    write_config(filename, config)\n",
        "\n",
        "\n",
        "def change_ui_config(filename):\n",
        "    config = read_config(filename)\n",
        "    config[\"txt2img/Prompt/value\"] = default_prompt\n",
        "    config[\"txt2img/Negative prompt/value\"] = default_neg_prompt\n",
        "    config[\"txt2img/Sampling method/value\"] = default_sampler\n",
        "    config[\"txt2img/Sampling steps/value\"] = default_steps\n",
        "    config[\"txt2img/Width/value\"] = default_width\n",
        "    config[\"txt2img/Height/value\"] = default_height\n",
        "    config[\"txt2img/Upscaler/value\"] = \"Latent (nearest-exact)\"\n",
        "    config[\"txt2img/Denoising strength/value\"] = default_denoising_strength\n",
        "    config[\"txt2img/CFG Scale/value\"] = default_cfg_scale\n",
        "    config[\"img2img/Prompt/value\"] = default_prompt\n",
        "    config[\"img2img/Negative prompt/value\"] = default_neg_prompt\n",
        "    config[\"img2img/Sampling method/value\"] = default_sampler\n",
        "    config[\"img2img/Sampling steps/value\"] = default_steps\n",
        "    config[\"img2img/Width/value\"] = default_width\n",
        "    config[\"img2img/Height/value\"] = default_height\n",
        "    config[\"img2img/Denoising strength/value\"] = default_denoising_strength\n",
        "    config[\"img2img/CFG Scale/value\"] = default_cfg_scale\n",
        "    write_config(filename, config)\n",
        "\n",
        "\n",
        "def update_extensions():\n",
        "    start_time = time.time()\n",
        "    extensions_updated = []\n",
        "    with tqdm(\n",
        "        total=len(os.listdir(extensions_dir)),\n",
        "        desc=\"\u001b[1;32mUpdating extensions\",\n",
        "        mininterval=0,\n",
        "    ) as pbar:\n",
        "        for dir in os.listdir(extensions_dir):\n",
        "            if os.path.isdir(os.path.join(extensions_dir, dir)):\n",
        "                os.chdir(os.path.join(extensions_dir, dir))\n",
        "                try:\n",
        "                    with capture.capture_output() as cap:\n",
        "                        !git fetch origin\n",
        "                        !git pull\n",
        "                except Exception as e:\n",
        "                    print(f\"\u001b[1;32mAn error occurred while updating {dir}: {e}\")\n",
        "\n",
        "                output = cap.stdout.strip()\n",
        "                if \"Already up to date.\" not in output:\n",
        "                    extensions_updated.append(dir)\n",
        "                pbar.update(1)\n",
        "\n",
        "    print(\"\\n\")\n",
        "    for ext in extensions_updated:\n",
        "        print(f\"\u001b[1;32m- {ext} updated to new version\")\n",
        "\n",
        "    end_time = time.time()\n",
        "    elapsed_time = int(end_time - start_time)\n",
        "\n",
        "    if elapsed_time < 60:\n",
        "        print(f\"\\n\u001b[1;32mAll extensions are up to date. Took {elapsed_time} sec\")\n",
        "    else:\n",
        "        mins, secs = divmod(elapsed_time, 60)\n",
        "        print(f\"\\n\u001b[1;32mAll extensions are up to date. Took {mins} mins {secs} sec\")\n",
        "\n",
        "\n",
        "def main():\n",
        "    start_time = time.time()\n",
        "\n",
        "    print(\"\u001b[1;32mInstalling...\\n\")\n",
        "\n",
        "    if not os.path.exists(webui_dir):\n",
        "        desc = \"\u001b[1;32mUnpacking Webui\"\n",
        "        pre_download(desc)\n",
        "    else:\n",
        "        print(\"\u001b[1;32mAlready installed, skipping...\")\n",
        "\n",
        "    with capture.capture_output() as cap:\n",
        "        os.chdir(os.path.join(webui_dir, \"repositories/stable-diffusion-stability-ai\"))\n",
        "        !git apply {patches_dir}/stablediffusion-lowram.patch\n",
        "\n",
        "        !sed -i \"s@os.path.splitext(checkpoint_.*@os.path.splitext(checkpoint_file); map_location='cuda'@\" {webui_dir}/modules/sd_models.py\n",
        "        !sed -i 's@ui.create_ui().*@ui.create_ui();shared.demo.queue(concurrency_count=999999,status_update_rate=0.1)@' {webui_dir}/webui.py\n",
        "\n",
        "        !sed -i \"s@'cpu'@'cuda'@\" {webui_dir}/modules/extras.py\n",
        "        del cap\n",
        "      \n",
        "    end_time = time.time()\n",
        "    elapsed_time = int(end_time - start_time)\n",
        "\n",
        "    change_config(config_file)\n",
        "    change_ui_config(ui_config_file)\n",
        "    open_theme(theme)\n",
        "\n",
        "    if elapsed_time < 60:\n",
        "        print(f\"\u001b[1;32mFinished unpacking. Took {elapsed_time} sec\")\n",
        "    else:\n",
        "        mins, secs = divmod(elapsed_time, 60)\n",
        "        print(f\"\u001b[1;32mFinished unpacking. Took {mins} mins {secs} sec\")\n",
        "\n",
        "    update_extensions()\n",
        "\n",
        "    #@markdown > Get <b>your</b> `ngrok_token` [here](https://dashboard.ngrok.com/get-started/your-authtoken) \n",
        "    ngrok_token = \"\" #@param {type: 'string'}\n",
        "    ngrok_region = \"ap\" #@param [\"us\", \"eu\", \"au\", \"ap\", \"sa\", \"jp\", \"in\"]\n",
        "\n",
        "    with capture.capture_output() as cap:\n",
        "      for file in os.listdir(output_dir):\n",
        "        file_path = os.path.join(output_dir, file)\n",
        "        if file_path.endswith((\".safetensors\", \".pt\", \".ckpt\")):\n",
        "          !ln \"{file_path}\" {webui_lora_dir}\n",
        "\n",
        "      for file in os.listdir(pretrained_model):\n",
        "        file_path = os.path.join(pretrained_model, file)\n",
        "        if file_path.endswith((\".safetensors\", \".ckpt\")):\n",
        "          !ln \"{file_path}\" {webui_models_dir}\n",
        "\n",
        "      for file in os.listdir(vae_dir):\n",
        "        file_path = os.path.join(vae_dir, file)\n",
        "        if file_path.endswith(\".vae.pt\"):\n",
        "          !ln \"{file_path}\" {webui_vaes_dir}\n",
        "\n",
        "      del cap\n",
        "\n",
        "    os.chdir(webui_dir)\n",
        "\n",
        "    print(\"\u001b[1;32m\")\n",
        "\n",
        "    config = {\n",
        "        \"enable-insecure-extension-access\": True,\n",
        "        \"disable-safe-unpickle\": True,\n",
        "        \"multiple\": True if not ngrok_token else False,\n",
        "        \"ckpt-dir\": webui_models_dir,\n",
        "        \"vae-dir\": webui_vaes_dir,\n",
        "        \"share\": True if not ngrok_token else False,\n",
        "        \"no-half-vae\": True,\n",
        "        \"lowram\": True,\n",
        "        \"gradio-queue\": True,\n",
        "        \"no-hashing\": True,\n",
        "        \"disable-console-progressbars\": True,\n",
        "        \"ngrok\": ngrok_token if ngrok_token else None,\n",
        "        \"ngrok-region\": ngrok_region if ngrok_token else None,\n",
        "        \"xformers\": True,\n",
        "        \"opt-sub-quad-attention\": True,\n",
        "        \"opt-channelslast\": True,\n",
        "        \"theme\": \"dark\"\n",
        "    }\n",
        "\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "\n",
        "    final_args = f\"python launch.py {args}\"\n",
        "\n",
        "    os.chdir(webui_dir)\n",
        "    !{final_args}\n",
        "\n",
        "main()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "nyIl9BhNXKUq"
      },
      "source": [
        "# VII. Deployment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8YlP0cNQyujy"
      },
      "outputs": [],
      "source": [
        "# @title ## 7.1. Upload Config\n",
        "from huggingface_hub import login\n",
        "from huggingface_hub import HfApi\n",
        "from huggingface_hub.utils import validate_repo_id, HfHubHTTPError\n",
        "\n",
        "# @markdown Login to Huggingface Hub\n",
        "# @markdown > Get **your** huggingface `WRITE` token [here](https://huggingface.co/settings/tokens)\n",
        "write_token = \"\"  # @param {type:\"string\"}\n",
        "# @markdown Fill this if you want to upload to your organization, or just leave it empty.\n",
        "orgs_name = \"\"  # @param{type:\"string\"}\n",
        "# @markdown If your model/dataset repo does not exist, it will automatically create it.\n",
        "model_name = \"your-model-name\"  # @param{type:\"string\"}\n",
        "dataset_name = \"your-dataset-name\"  # @param{type:\"string\"}\n",
        "make_private = False  # @param{type:\"boolean\"}\n",
        "\n",
        "def authenticate(write_token):\n",
        "    login(write_token, add_to_git_credential=True)\n",
        "    api = HfApi()\n",
        "    return api.whoami(write_token), api\n",
        "\n",
        "\n",
        "def create_repo(api, user, orgs_name, repo_name, repo_type, make_private=False):\n",
        "    global model_repo\n",
        "    global datasets_repo\n",
        "    \n",
        "    if orgs_name == \"\":\n",
        "        repo_id = user[\"name\"] + \"/\" + repo_name.strip()\n",
        "    else:\n",
        "        repo_id = orgs_name + \"/\" + repo_name.strip()\n",
        "\n",
        "    try:\n",
        "        validate_repo_id(repo_id)\n",
        "        api.create_repo(repo_id=repo_id, repo_type=repo_type, private=make_private)\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' didn't exist, creating repo\")\n",
        "    except HfHubHTTPError as e:\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' exists, skipping create repo\")\n",
        "    \n",
        "    if repo_type == \"model\":\n",
        "        model_repo = repo_id\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' link: https://huggingface.co/{repo_id}\\n\")\n",
        "    else:\n",
        "        datasets_repo = repo_id\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' link: https://huggingface.co/datasets/{repo_id}\\n\")\n",
        "\n",
        "user, api = authenticate(write_token)\n",
        "\n",
        "if model_name:\n",
        "    create_repo(api, user, orgs_name, model_name, \"model\", make_private)\n",
        "if dataset_name:\n",
        "    create_repo(api, user, orgs_name, dataset_name, \"dataset\", make_private)\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Fuxghk8MnG6j"
      },
      "source": [
        "## 7.2. Upload with Huggingface Hub"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8HpOZOglBJTz"
      },
      "outputs": [],
      "source": [
        "# @title ### 7.2.1. Upload LoRA\n",
        "from huggingface_hub import HfApi\n",
        "from pathlib import Path\n",
        "\n",
        "%store -r\n",
        "\n",
        "api = HfApi()\n",
        "\n",
        "# @markdown This will be uploaded to model repo\n",
        "model_path = \"/content/LoRA/output\"  # @param {type :\"string\"}\n",
        "path_in_repo = \"\"  # @param {type :\"string\"}\n",
        "\n",
        "# @markdown Now you can save your config file for future use\n",
        "config_path = \"/content/LoRA/config\"  # @param {type :\"string\"}\n",
        "\n",
        "# @markdown Other Information\n",
        "commit_message = \"\"  # @param {type :\"string\"}\n",
        "\n",
        "if not commit_message:\n",
        "    commit_message = f\"feat: upload {project_name} lora model\"\n",
        "\n",
        "def upload_to_hf(model_path, is_folder, is_config):\n",
        "    path_obj = Path(model_path)\n",
        "    trained_model = path_obj.parts[-1]\n",
        "\n",
        "    if path_in_repo:\n",
        "        trained_model = path_in_repo\n",
        "\n",
        "    if is_config:\n",
        "        trained_model = f\"{project_name}_config\"\n",
        "\n",
        "    print(f\"Uploading {trained_model} to https://huggingface.co/{model_repo}\")\n",
        "    print(\"Please wait...\")\n",
        "\n",
        "    if is_folder:\n",
        "        api.upload_folder(\n",
        "            folder_path=model_path,\n",
        "            path_in_repo=trained_model,\n",
        "            repo_id=model_repo,\n",
        "            commit_message=commit_message,\n",
        "            ignore_patterns=\".ipynb_checkpoints\",\n",
        "        )\n",
        "        print(f\"Upload success, located at https://huggingface.co/{model_repo}/tree/main\\n\")\n",
        "    else:\n",
        "        api.upload_file(\n",
        "            path_or_fileobj=model_path,\n",
        "            path_in_repo=trained_model,\n",
        "            repo_id=model_repo,\n",
        "            commit_message=commit_message,\n",
        "        )\n",
        "        print(f\"Upload success, located at https://huggingface.co/{model_repo}/blob/main/{trained_model}\\n\")\n",
        "\n",
        "def upload():\n",
        "    is_model_file = model_path.endswith((\".ckpt\", \".safetensors\", \".pt\"))\n",
        "    upload_to_hf(model_path, not is_model_file, False)\n",
        "\n",
        "    if config_path:\n",
        "        upload_to_hf(config_path, True, True)\n",
        "\n",
        "upload()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "IW-hS9jnmf-E"
      },
      "outputs": [],
      "source": [
        "# @title ### 7.2.2. Upload Dataset\n",
        "from huggingface_hub import HfApi\n",
        "from pathlib import Path\n",
        "import shutil\n",
        "import zipfile\n",
        "import os\n",
        "\n",
        "api = HfApi()\n",
        "\n",
        "# @markdown This will be compressed to zip and  uploaded to datasets repo, leave it empty if not necessary\n",
        "train_data_path = \"/content/LoRA/train_data\"  # @param {type :\"string\"}\n",
        "\n",
        "# @markdown `Nerd stuff, only if you want to save training logs`\n",
        "logs_path = \"/content/LoRA/logs\"  # @param {type :\"string\"}\n",
        "\n",
        "tmp_dataset = f\"/content/LoRA/{project_name}_dataset\" if project_name else \"/content/LoRA/tmp_dataset\"\n",
        "tmp_train_data = f\"{tmp_dataset}/train_data\"\n",
        "dataset_zip = f\"{tmp_dataset}.zip\"\n",
        "\n",
        "# @markdown Other Information\n",
        "commit_message = \"\"  # @param {type :\"string\"}\n",
        "\n",
        "if not commit_message:\n",
        "    commit_message = f\"feat: upload {project_name} dataset and logs\"\n",
        "\n",
        "os.makedirs(tmp_dataset, exist_ok=True)\n",
        "os.makedirs(tmp_train_data, exist_ok=True)\n",
        "\n",
        "def upload_dataset(dataset_path, is_zip):\n",
        "    path_obj = Path(dataset_path)\n",
        "    dataset_name = path_obj.parts[-1]\n",
        "\n",
        "    print(f\"Uploading {dataset_name} to https://huggingface.co/datasets/{datasets_repo}\")\n",
        "    print(\"Please wait...\")\n",
        "\n",
        "    if is_zip:\n",
        "        api.upload_file(\n",
        "            path_or_fileobj=dataset_path,\n",
        "            path_in_repo=dataset_name,\n",
        "            repo_id=datasets_repo,\n",
        "            repo_type=\"dataset\",\n",
        "            commit_message=commit_message,\n",
        "        )\n",
        "        print(f\"Upload success, located at https://huggingface.co/datasets/{datasets_repo}/blob/main/{dataset_name}\\n\")\n",
        "    else:\n",
        "        api.upload_folder(\n",
        "            folder_path=dataset_path,\n",
        "            path_in_repo=dataset_name,\n",
        "            repo_id=datasets_repo,\n",
        "            repo_type=\"dataset\",\n",
        "            commit_message=commit_message,\n",
        "            ignore_patterns=\".ipynb_checkpoints\",\n",
        "        )\n",
        "        print(f\"Upload success, located at https://huggingface.co/datasets/{datasets_repo}/tree/main/{dataset_name}\\n\")\n",
        "\n",
        "def zip_file(folder_path):\n",
        "    zip_path = f\"{folder_path}.zip\"\n",
        "    with zipfile.ZipFile(zip_path, \"w\") as zip_file:\n",
        "        for root, dirs, files in os.walk(folder_path):\n",
        "            for file in files:\n",
        "                zip_file.write(os.path.join(root, file))\n",
        "\n",
        "def move(src_path, dst_path, move_metadata):\n",
        "    metadata_files = [\n",
        "        \"meta_cap.json\",\n",
        "        \"meta_cap_dd.json\",\n",
        "        \"meta_lat.json\",\n",
        "        \"meta_clean.json\",\n",
        "        \"meta_final.json\",\n",
        "    ]\n",
        "\n",
        "    if os.path.exists(src_path):\n",
        "        shutil.move(src_path, dst_path)\n",
        "\n",
        "    if move_metadata:\n",
        "        parent_meta_path = os.path.dirname(src_path)\n",
        "\n",
        "        for filename in os.listdir(parent_meta_path):\n",
        "            file_path = os.path.join(parent_meta_path, filename)\n",
        "            if filename in metadata_files:\n",
        "                shutil.move(file_path, dst_path)\n",
        "\n",
        "def upload():\n",
        "    if train_data_path:\n",
        "        move(train_data_path, tmp_train_data, False)\n",
        "        zip_file(tmp_dataset)\n",
        "        upload_dataset(dataset_zip, True)\n",
        "        os.remove(dataset_zip)\n",
        "    if logs_path:\n",
        "        upload_dataset(logs_path, False)\n",
        "\n",
        "upload()\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "CKZpg4keWS5c"
      },
      "source": [
        "## 7.3. Upload with GIT (Alternative)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "6nBlrOrytO9F"
      },
      "outputs": [],
      "source": [
        "# @title ### 7.3.1. Clone Repository\n",
        "\n",
        "clone_model = True  # @param {'type': 'boolean'}\n",
        "clone_dataset = True  # @param {'type': 'boolean'}\n",
        "\n",
        "def clone_repository(repo_url, local_path):\n",
        "    !git lfs install --skip-smudge\n",
        "    os.environ[\"GIT_LFS_SKIP_SMUDGE\"] = \"1\"\n",
        "    !git clone {repo_url} {local_path}\n",
        "\n",
        "if clone_model:\n",
        "    clone_repository(f\"https://huggingface.co/{model_repo}\", f\"/content/{model_name}\")\n",
        "\n",
        "if clone_dataset:\n",
        "    clone_repository(f\"https://huggingface.co/datasets/{datasets_repo}\", f\"/content/{dataset_name}\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "7bJev4PzOFFB"
      },
      "outputs": [],
      "source": [
        "# @title ### 7.3.2. Commit using Git\n",
        "import os\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "# @markdown Choose which repo you want to commit\n",
        "commit_model = True  # @param {'type': 'boolean'}\n",
        "commit_dataset = True  # @param {'type': 'boolean'}\n",
        "# @markdown Other Information\n",
        "commit_message = \"\"  # @param {type :\"string\"}\n",
        "\n",
        "if not commit_message:\n",
        "    commit_message = f\"feat: upload {project_name} lora model and dataset\"\n",
        "\n",
        "!git config --global user.email \"example@mail.com\"\n",
        "!git config --global user.name \"example\"\n",
        "\n",
        "def commit(repo_folder, commit_message):\n",
        "    os.chdir(os.path.join(root_dir, repo_folder))\n",
        "    !git lfs install\n",
        "    !huggingface-cli lfs-enable-largefiles .\n",
        "    !git add .\n",
        "    !git commit -m \"{commit_message}\"\n",
        "    !git push\n",
        "\n",
        "\n",
        "if commit_model:\n",
        "    commit(model_name, commit_message)\n",
        "\n",
        "if commit_dataset:\n",
        "    commit(dataset_name, commit_message)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}