Dataset_Maker.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rmCPmqFL6hCQ"
      },
      "source": [
        "# 📊 Dataset Maker by Hollowstrawberry\n",
        "\n",
        "This is based on the work of [Kohya-ss](https://github.com/kohya-ss/sd-scripts) and [Linaqruf](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-dreambooth.ipynb). Thank you!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UZ0f5TOg_iIl"
      },
      "source": [
        "### ⭕ Disclaimer\n",
        "The purpose of this document is to research bleeding-edge technologies in the field of machine learning inference.  \n",
        "Please read and follow the [Google Colab guidelines](https://research.google.com/colaboratory/faq.html) and its [Terms of Service](https://research.google.com/colaboratory/tos_v3.html)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-rdgF2AWLS2h"
      },
      "source": [
        "| |GitHub|🇬🇧 English|🇪🇸 Spanish|\n",
        "|:--|:-:|:-:|:-:|\n",
        "| 🏠 **Homepage** | [![GitHub](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/github.svg)](https://github.com/hollowstrawberry/kohya-colab) | | |\n",
        "| 📊 **Dataset Maker** | [![GitHub](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/github.svg)](https://github.com/hollowstrawberry/kohya-colab/blob/main/Dataset_Maker.ipynb) | [![Open in Colab](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/colab-badge.svg)](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Dataset_Maker.ipynb) | [![Abrir en Colab](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/colab-badge-spanish.svg)](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Spanish_Dataset_Maker.ipynb) |\n",
        "| ⭐ **Lora Trainer** | [![GitHub](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/github.svg)](https://github.com/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer.ipynb) | [![Open in Colab](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/colab-badge.svg)](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer.ipynb) | [![Abrir en Colab](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/colab-badge-spanish.svg)](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Spanish_Lora_Trainer.ipynb) |\n",
        "| 🌟 **XL Lora Trainer** | [![GitHub](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/github.svg)](https://github.com/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer_XL.ipynb) | [![Open in Colab](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/colab-badge.svg)](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer_XL.ipynb) |  |"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "cBa7KdewQ4BU"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from IPython import get_ipython\n",
        "from IPython.display import display, Markdown\n",
        "\n",
        "COLAB = True\n",
        "\n",
        "if COLAB:\n",
        "  from google.colab.output import clear as clear_output\n",
        "else:\n",
        "  from IPython.display import clear_output\n",
        "\n",
        "#@title ## 🚩 Start Here\n",
        "\n",
        "#@markdown ### 1️⃣ Setup\n",
        "#@markdown This cell will load some requirements and create the necessary folders in your Google Drive. <p>\n",
        "#@markdown Your project name can't contain spaces but it can contain a single / to make a subfolder in your dataset.\n",
        "project_name = \"\" #@param {type:\"string\"}\n",
        "project_name = project_name.strip()\n",
        "#@markdown The folder structure doesn't matter and is purely for comfort. Make sure to always pick the same one. I like organizing by project.\n",
        "folder_structure = \"Organize by project (MyDrive/Loras/project_name/dataset)\" #@param [\"Organize by category (MyDrive/lora_training/datasets/project_name)\", \"Organize by project (MyDrive/Loras/project_name/dataset)\"]\n",
        "\n",
        "if not project_name or any(c in project_name for c in \" .()\\\"'\\\\\") or project_name.count(\"/\") > 1:\n",
        "  print(\"Please write a valid project_name.\")\n",
        "else:\n",
        "  if COLAB and not os.path.exists('/content/drive'):\n",
        "    from google.colab import drive\n",
        "    print(\"📂 Connecting to Google Drive...\")\n",
        "    drive.mount('/content/drive')\n",
        "\n",
        "  project_base = project_name if \"/\" not in project_name else project_name[:project_name.rfind(\"/\")]\n",
        "  project_subfolder = project_name if \"/\" not in project_name else project_name[project_name.rfind(\"/\")+1:]\n",
        "\n",
        "  root_dir = \"/content\" if COLAB else \"~/Loras\"\n",
        "  deps_dir = os.path.join(root_dir, \"deps\")\n",
        "\n",
        "  if \"/Loras\" in folder_structure:\n",
        "    main_dir      = os.path.join(root_dir, \"drive/MyDrive/Loras\") if COLAB else root_dir\n",
        "    config_folder = os.path.join(main_dir, project_base)\n",
        "    images_folder = os.path.join(main_dir, project_base, \"dataset\")\n",
        "    if \"/\" in project_name:\n",
        "      images_folder = os.path.join(images_folder, project_subfolder)\n",
        "  else:\n",
        "    main_dir      = os.path.join(root_dir, \"drive/MyDrive/lora_training\") if COLAB else root_dir\n",
        "    config_folder = os.path.join(main_dir, \"config\", project_name)\n",
        "    images_folder = os.path.join(main_dir, \"datasets\", project_name)\n",
        "\n",
        "  for dir in [main_dir, deps_dir, images_folder, config_folder]:\n",
        "    os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "  print(f\"✅ Project {project_name} is ready!\")\n",
        "  step1_installed_flag = True\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "afu5dCKTV31E"
      },
      "outputs": [],
      "source": [
        "if \"step1_installed_flag\" not in globals():\n",
        "  raise Exception(\"Please run step 1 first!\")\n",
        "\n",
        "import json\n",
        "import time\n",
        "from urllib.request import urlopen, Request\n",
        "\n",
        "#@markdown ### 2️⃣ Scrape images from Gelbooru\n",
        "\n",
        "#@markdown We will grab images from the popular anime gallery [Gelbooru](https://gelbooru.com/). Images are sorted by tags, including poses, scenes, character traits, character names, artists, etc. <p>\n",
        "#@markdown * If you instead want to use your own images, upload them to your Google Drive's `Loras/project_name/dataset` folder.\n",
        "#@markdown * If you instead want to download screencaps of anime episodes, try [this other colab by another person](https://colab.research.google.com/drive/1oBSntB40BKzNmKceXUlkXzujzdQw-Ci7). It's more complicated though.\n",
        "\n",
        "#@markdown Up to 1000 images may be downloaded by this step in just one minute. Don't abuse it. <p>\n",
        "#@markdown Your target tags should include the relevant tags for your character/concept/artstyle, and exclude undesired tags (for example, explicit images may affect learning).\n",
        "#@markdown Separate words with underscores, separate tags with spaces, and use - to exclude a tag. You can also include a minimum score: `score:>10`\n",
        "tags = \"1girl -sex -bdsm -loli -greyscale -monochrome hatsune_miku\" #@param {type:\"string\"}\n",
        "##@markdown If an image is bigger than this resolution a smaller version will be downloaded instead.\n",
        "max_resolution = 3072 #param {type:\"slider\", min:1024, max:8196, step:1024}\n",
        "##@markdown Posts with a parent post are often minor variations of the same image.\n",
        "include_posts_with_parent = True #param {type:\"boolean\"}\n",
        "\n",
        "tags = tags.replace(\" \", \"+\")\\\n",
        "           .replace(\"(\", \"%28\")\\\n",
        "           .replace(\")\", \"%29\")\\\n",
        "           .replace(\":\", \"%3a\")\\\n",
        "           .replace(\"&\", \"%26\")\\\n",
        "\n",
        "url = \"https://gelbooru.com/index.php?page=dapi&json=1&s=post&q=index&limit=100&tags={}\".format(tags)\n",
        "user_agent = \"Mozilla/5.0 AppleWebKit/537.36 (KHTML, like Gecko; compatible; Googlebot/2.1; +http://www.google.com/bot.html) Chrome/93.0.4577.83 Safari/537.36\"\n",
        "limit = 100 # hardcoded by gelbooru\n",
        "total_limit = 1000 # you can edit this if you want but I wouldn't recommend it\n",
        "supported_types = (\".png\", \".jpg\", \".jpeg\")\n",
        "\n",
        "def ubuntu_deps():\n",
        "  print(\"🏭 Installing dependencies...\\n\")\n",
        "  !apt -y install aria2\n",
        "  return not get_ipython().__dict__['user_ns']['_exit_code']\n",
        "\n",
        "if \"step2_installed_flag\" not in globals():\n",
        "  if ubuntu_deps():\n",
        "    clear_output()\n",
        "    step2_installed_flag = True\n",
        "  else:\n",
        "    print(\"❌ Error installing dependencies, attempting to continue anyway...\")\n",
        "\n",
        "def get_json(url):\n",
        "  with urlopen(Request(url, headers={\"User-Agent\": user_agent})) as page:\n",
        "    return json.load(page)\n",
        "\n",
        "def filter_images(data):\n",
        "  return [p[\"file_url\"] if p[\"width\"]*p[\"height\"] <= max_resolution**2 else p[\"sample_url\"]\n",
        "          for p in data[\"post\"]\n",
        "          if (p[\"parent_id\"] == 0 or include_posts_with_parent)\n",
        "          and p[\"file_url\"].lower().endswith(supported_types)]\n",
        "\n",
        "def download_images():\n",
        "  data = get_json(url)\n",
        "  count = data[\"@attributes\"][\"count\"]\n",
        "\n",
        "  if count == 0:\n",
        "    print(\"📷 No results found\")\n",
        "    return\n",
        "\n",
        "  print(f\"🎯 Found {count} results\")\n",
        "  test_url = \"https://gelbooru.com/index.php?page=post&s=list&tags={}\".format(tags)\n",
        "  display(Markdown(f\"[Click here to open in browser!]({test_url})\"))\n",
        "  print (f\"🔽 Will download to {images_folder.replace('/content/drive/', '')} (A confirmation box should appear below, otherwise run this cell again)\")\n",
        "  inp = input(\"❓ Enter the word 'yes' if you want to proceed with the download: \")\n",
        "\n",
        "  if inp.lower().strip() != 'yes':\n",
        "    print(\"❌ Download cancelled\")\n",
        "    return\n",
        "\n",
        "  print(\"📩 Grabbing image list...\")\n",
        "\n",
        "  image_urls = set()\n",
        "  image_urls = image_urls.union(filter_images(data))\n",
        "  for i in range(total_limit // limit):\n",
        "    count -= limit\n",
        "    if count <= 0:\n",
        "      break\n",
        "    time.sleep(0.1)\n",
        "    image_urls = image_urls.union(filter_images(get_json(url+f\"&pid={i+1}\")))\n",
        "\n",
        "  scrape_file = os.path.join(config_folder, f\"scrape_{project_subfolder}.txt\")\n",
        "  with open(scrape_file, \"w\") as f:\n",
        "    f.write(\"\\n\".join(image_urls))\n",
        "\n",
        "  print(f\"🌐 Saved links to {scrape_file}\\n\\n🔁 Downloading images...\\n\")\n",
        "  old_img_count = len([f for f in os.listdir(images_folder) if f.lower().endswith(supported_types)])\n",
        "\n",
        "  os.chdir(images_folder)\n",
        "  !aria2c --console-log-level=warn -c -x 16 -k 1M -s 16 -i {scrape_file}\n",
        "\n",
        "  new_img_count = len([f for f in os.listdir(images_folder) if f.lower().endswith(supported_types)])\n",
        "  print(f\"\\n✅ Downloaded {new_img_count - old_img_count} images.\")\n",
        "\n",
        "download_images()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "b218DEEMpwzB"
      },
      "outputs": [],
      "source": [
        "if \"step1_installed_flag\" not in globals():\n",
        "  raise Exception(\"Please run step 1 first!\")\n",
        "\n",
        "#@markdown ### 3️⃣ Curate your images\n",
        "#@markdown We will find duplicate images with the FiftyOne AI, and mark them with `delete`. <p>\n",
        "#@markdown Then, an interactive area will appear below this cell that lets you visualize all your images and manually mark with `delete` to the ones you don't like. <p>\n",
        "#@markdown If the interactive area appears blank for over a minute, try enabling cookies and removing tracking protection for the Google Colab website, as they may break it.\n",
        "#@markdown Regardless, you can save your changes by sending Enter in the input box above the interactive area.<p>\n",
        "#@markdown This is how similar 2 images must be to be marked for deletion. I recommend 0.97 to 0.99:\n",
        "similarity_threshold = 0.985 #@param {type:\"number\"}\n",
        "\n",
        "\n",
        "os.chdir(root_dir)\n",
        "model_name = \"clip-vit-base32-torch\"\n",
        "supported_types = (\".png\", \".jpg\", \".jpeg\")\n",
        "img_count = len(os.listdir(images_folder))\n",
        "batch_size = min(250, img_count)\n",
        "\n",
        "if \"step3_installed_flag\" not in globals():\n",
        "  print(\"🏭 Installing dependencies...\\n\")\n",
        "  !pip -q install fiftyone ftfy\n",
        "  !pip -q install fiftyone-db-ubuntu2204\n",
        "  if not get_ipython().__dict__['user_ns']['_exit_code']:\n",
        "    clear_output()\n",
        "    step3_installed_flag = True\n",
        "  else:\n",
        "    print(\"❌ Error installing dependencies, attempting to continue anyway...\")\n",
        "\n",
        "import numpy as np\n",
        "import fiftyone as fo\n",
        "import fiftyone.zoo as foz\n",
        "from fiftyone import ViewField as F\n",
        "from sklearn.metrics.pairwise import cosine_similarity\n",
        "\n",
        "non_images = [f for f in os.listdir(images_folder) if not f.lower().endswith(supported_types)]\n",
        "if non_images:\n",
        "  print(f\"💥 Error: Found non-image file {non_images[0]} - This program doesn't allow it. Sorry! Use the Extras at the bottom to clean the folder.\")\n",
        "elif img_count == 0:\n",
        "  print(f\"💥 Error: No images found in {images_folder}\")\n",
        "else:\n",
        "  print(\"\\n💿 Analyzing dataset...\\n\")\n",
        "  dataset = fo.Dataset.from_dir(images_folder, dataset_type=fo.types.ImageDirectory)\n",
        "  model = foz.load_zoo_model(model_name)\n",
        "  embeddings = dataset.compute_embeddings(model, batch_size=batch_size)\n",
        "\n",
        "  batch_embeddings = np.array_split(embeddings, batch_size)\n",
        "  similarity_matrices = []\n",
        "  max_size_x = max(array.shape[0] for array in batch_embeddings)\n",
        "  max_size_y = max(array.shape[1] for array in batch_embeddings)\n",
        "\n",
        "  for i, batch_embedding in enumerate(batch_embeddings):\n",
        "    similarity = cosine_similarity(batch_embedding)\n",
        "    #Pad 0 for np.concatenate\n",
        "    padded_array = np.zeros((max_size_x, max_size_y))\n",
        "    padded_array[0:similarity.shape[0], 0:similarity.shape[1]] = similarity\n",
        "    similarity_matrices.append(padded_array)\n",
        "\n",
        "  similarity_matrix = np.concatenate(similarity_matrices, axis=0)\n",
        "  similarity_matrix = similarity_matrix[0:embeddings.shape[0], 0:embeddings.shape[0]]\n",
        "\n",
        "  similarity_matrix = cosine_similarity(embeddings)\n",
        "  similarity_matrix -= np.identity(len(similarity_matrix))\n",
        "\n",
        "  dataset.match(F(\"max_similarity\") > similarity_threshold)\n",
        "  dataset.tags = [\"delete\", \"has_duplicates\"]\n",
        "\n",
        "  id_map = [s.id for s in dataset.select_fields([\"id\"])]\n",
        "  samples_to_remove = set()\n",
        "  samples_to_keep = set()\n",
        "\n",
        "  for idx, sample in enumerate(dataset):\n",
        "    if sample.id not in samples_to_remove:\n",
        "      # Keep the first instance of two duplicates\n",
        "      samples_to_keep.add(sample.id)\n",
        "\n",
        "      dup_idxs = np.where(similarity_matrix[idx] > similarity_threshold)[0]\n",
        "      for dup in dup_idxs:\n",
        "          # We kept the first instance so remove all other duplicates\n",
        "          samples_to_remove.add(id_map[dup])\n",
        "\n",
        "      if len(dup_idxs) > 0:\n",
        "          sample.tags.append(\"has_duplicates\")\n",
        "          sample.save()\n",
        "    else:\n",
        "      sample.tags.append(\"delete\")\n",
        "      sample.save()\n",
        "\n",
        "  clear_output()\n",
        "\n",
        "  sidebar_groups = fo.DatasetAppConfig.default_sidebar_groups(dataset)\n",
        "  for group in sidebar_groups[1:]:\n",
        "    group.expanded = False\n",
        "  dataset.app_config.sidebar_groups = sidebar_groups\n",
        "  dataset.save()\n",
        "  session = fo.launch_app(dataset)\n",
        "\n",
        "  print(\"❗ Wait a minute for the session to load. If it doesn't, read above.\")\n",
        "  print(\"❗ When it's ready, you'll see a grid of your images.\")\n",
        "  print(\"❗ On the left side enable \\\"sample tags\\\" to visualize the images marked for deletion.\")\n",
        "  print(\"❗ You can mark your own images with the \\\"delete\\\" label by selecting them and pressing the tag icon at the top.\")\n",
        "  input(\"⭕ When you're done, enter something here to save your changes: \")\n",
        "\n",
        "  print(\"💾 Saving...\")\n",
        "\n",
        "  marked = [s for s in dataset if \"delete\" in s.tags]\n",
        "  dataset.delete_samples(marked)\n",
        "  previous_folder = images_folder[:images_folder.rfind(\"/\")]\n",
        "  dataset.export(export_dir=os.path.join(images_folder, project_subfolder), dataset_type=fo.types.ImageDirectory)\n",
        "\n",
        "  temp_suffix = \"_temp\"\n",
        "  !mv {images_folder} {images_folder}{temp_suffix}\n",
        "  !mv {images_folder}{temp_suffix}/{project_subfolder} {images_folder}\n",
        "  !rm -r {images_folder}{temp_suffix}\n",
        "\n",
        "  session.refresh()\n",
        "  fo.close_app()\n",
        "  clear_output()\n",
        "\n",
        "  print(f\"\\n✅ Removed {len(marked)} images from dataset. You now have {len(os.listdir(images_folder))} images.\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "sl4FD7Mz-uea"
      },
      "outputs": [],
      "source": [
        "if \"step1_installed_flag\" not in globals():\n",
        "  raise Exception(\"Please run step 1 first!\")\n",
        "\n",
        "#@markdown ### 4️⃣ Tag your images\n",
        "#@markdown We will be using AI to automatically tag your images, specifically [Waifu Diffusion](https://huggingface.co/SmilingWolf/wd-v1-4-swinv2-tagger-v2) in the case of anime and [BLIP](https://huggingface.co/spaces/Salesforce/BLIP) in the case of photos.\n",
        "#@markdown Giving tags/captions to your images allows for much better training. This process should take a couple minutes. <p>\n",
        "method = \"Anime tags\" #@param [\"Anime tags\", \"Photo captions\"]\n",
        "#@markdown **Anime:** The threshold is the minimum level of confidence the tagger must have in order to include a tag. Lower threshold = More tags. Recommended 0.35 to 0.5\n",
        "tag_threshold = 0.35 #@param {type:\"slider\", min:0.0, max:1.0, step:0.01}\n",
        "blacklist_tags = \"bangs, breasts, multicolored hair, two-tone hair, gradient hair, virtual youtuber, parody, style parody, official alternate costume, official alternate hairstyle, official alternate hair length, alternate costume, alternate hairstyle, alternate hair length, alternate hair color\" #@param {type:\"string\"}\n",
        "#@markdown **Photos:** The minimum and maximum length of tokens/words in each caption.\n",
        "caption_min = 10 #@param {type:\"number\"}\n",
        "caption_max = 75 #@param {type:\"number\"}\n",
        "\n",
        "%env PYTHONPATH=/env/python\n",
        "os.chdir(root_dir)\n",
        "kohya = \"/content/kohya-trainer\"\n",
        "if not os.path.exists(kohya):\n",
        "  !git clone https://github.com/kohya-ss/sd-scripts {kohya}\n",
        "  os.chdir(kohya)\n",
        "  !git reset --hard 9a67e0df390033a89f17e70df5131393692c2a55\n",
        "  os.chdir(root_dir)\n",
        "\n",
        "if \"tags\" in method:\n",
        "  if \"step4a_installed_flag\" not in globals() or \"step4b_installed_flag\" in globals():\n",
        "    print(\"\\n🏭 Installing dependencies...\\n\")\n",
        "    !pip install accelerate==0.25.0 diffusers[torch]==0.25.0 einops==0.6.0 tensorflow==2.15.0 \\\n",
        "      keras==2.15.0 transformers safetensors huggingface-hub==0.22.0 torchvision albumentations \\\n",
        "      jax==0.4.23 jaxlib==0.4.23\n",
        "    if not get_ipython().__dict__['user_ns']['_exit_code']:\n",
        "      clear_output()\n",
        "      step4a_installed_flag = True\n",
        "    else:\n",
        "      print(\"❌ Error installing dependencies, trying to continue anyway...\")\n",
        "\n",
        "  print(\"\\n🚶‍♂️ Launching program...\\n\")\n",
        "\n",
        "  os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'\n",
        "  %env PYTHONPATH={kohya}\n",
        "  !python {kohya}/finetune/tag_images_by_wd14_tagger.py \\\n",
        "    {images_folder} \\\n",
        "    --repo_id=SmilingWolf/wd-v1-4-swinv2-tagger-v2 \\\n",
        "    --model_dir={root_dir} \\\n",
        "    --thresh={tag_threshold} \\\n",
        "    --batch_size=8 \\\n",
        "    --caption_extension=.txt \\\n",
        "    --force_download\n",
        "\n",
        "  if not get_ipython().__dict__['user_ns']['_exit_code']:\n",
        "    print(\"removing underscores and blacklist...\")\n",
        "    blacklisted_tags = [t.strip() for t in blacklist_tags.split(\",\")]\n",
        "    from collections import Counter\n",
        "    top_tags = Counter()\n",
        "    for txt in [f for f in os.listdir(images_folder) if f.lower().endswith(\".txt\")]:\n",
        "      with open(os.path.join(images_folder, txt), 'r') as f:\n",
        "        tags = [t.strip() for t in f.read().split(\",\")]\n",
        "        tags = [t.replace(\"_\", \" \") if len(t) > 3 else t for t in tags]\n",
        "        tags = [t for t in tags if t not in blacklisted_tags]\n",
        "      top_tags.update(tags)\n",
        "      with open(os.path.join(images_folder, txt), 'w') as f:\n",
        "        f.write(\", \".join(tags))\n",
        "\n",
        "    %env PYTHONPATH=/env/python\n",
        "    clear_output()\n",
        "    print(f\"📊 Tagging complete. Here are the top 50 tags in your dataset:\")\n",
        "    print(\"\\n\".join(f\"{k} ({v})\" for k, v in top_tags.most_common(50)))\n",
        "\n",
        "\n",
        "else: # Photos\n",
        "  if \"step4b_installed_flag\" not in globals() or \"step4a_installed_flag\" in globals():\n",
        "    print(\"\\n🏭 Installing dependencies...\\n\")\n",
        "    !pip install timm==0.6.12 fairscale==0.4.13 transformers==4.26.0 requests \\\n",
        "      accelerate==0.25.0 diffusers[torch]==0.25.0 einops==0.6.0 safetensors \\\n",
        "      jax==0.4.23 jaxlib==0.4.23 huggingface-hub==0.22.0\n",
        "    if not get_ipython().__dict__['user_ns']['_exit_code']:\n",
        "      clear_output()\n",
        "      step4b_installed_flag = True\n",
        "    else:\n",
        "      print(\"❌ Error installing dependencies, trying to continue anyway...\")\n",
        "\n",
        "  print(\"\\n🚶‍♂️ Launching program...\\n\")\n",
        "\n",
        "  os.chdir(kohya)\n",
        "  %env PYTHONPATH={kohya}\n",
        "  !python {kohya}/finetune/make_captions.py \\\n",
        "    {images_folder} \\\n",
        "    --beam_search \\\n",
        "    --max_data_loader_n_workers=2 \\\n",
        "    --batch_size=8 \\\n",
        "    --min_length={caption_min} \\\n",
        "    --max_length={caption_max} \\\n",
        "    --caption_extension=.txt\n",
        "\n",
        "  if not get_ipython().__dict__['user_ns']['_exit_code']:\n",
        "    import random\n",
        "    captions = [f for f in os.listdir(images_folder) if f.lower().endswith(\".txt\")]\n",
        "    sample = []\n",
        "    for txt in random.sample(captions, min(10, len(captions))):\n",
        "      with open(os.path.join(images_folder, txt), 'r') as f:\n",
        "        sample.append(f.read())\n",
        "\n",
        "    os.chdir(root_dir)\n",
        "    %env PYTHONPATH=/env/python\n",
        "    clear_output()\n",
        "    print(f\"📊 Captioning complete. Here are {len(sample)} example captions from your dataset:\")\n",
        "    print(\"\".join(sample))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "WBFik7accyDz"
      },
      "outputs": [],
      "source": [
        "if \"step1_installed_flag\" not in globals():\n",
        "  raise Exception(\"Please run step 1 first!\")\n",
        "\n",
        "#@markdown ### 5️⃣ Curate your tags\n",
        "#@markdown Modify your dataset's tags. You can run this cell multiple times with different parameters. <p>\n",
        "\n",
        "#@markdown Put an activation tag at the start of every text file. This is useful to make learning better and activate your Lora easier. Set `keep_tokens` to 1 when training.<p>\n",
        "#@markdown Common tags that are removed such as hair color, etc. will be \"absorbed\" by your activation tag.\n",
        "global_activation_tag = \"hatsune miku\" #@param {type:\"string\"}\n",
        "remove_tags = \"\" #@param {type:\"string\"}\n",
        "#@markdown &nbsp;\n",
        "\n",
        "#@markdown In this advanced section, you can search text files containing matching tags, and replace them with less/more/different tags. If you select the checkbox below, any extra tags will be put at the start of the file, letting you assign different activation tags to different parts of your dataset. Still, you may want a more advanced tool for this.\n",
        "search_tags = \"\" #@param {type:\"string\"}\n",
        "replace_with = \"\" #@param {type:\"string\"}\n",
        "search_mode = \"OR\" #@param [\"OR\", \"AND\"]\n",
        "new_becomes_activation_tag = False #@param {type:\"boolean\"}\n",
        "#@markdown These may be useful sometimes. Will remove existing activation tags, be careful.\n",
        "sort_alphabetically = False #@param {type:\"boolean\"}\n",
        "remove_duplicates = False #@param {type:\"boolean\"}\n",
        "\n",
        "def split_tags(tagstr):\n",
        "  return [s.strip() for s in tagstr.split(\",\") if s.strip()]\n",
        "\n",
        "activation_tag_list = split_tags(global_activation_tag)\n",
        "remove_tags_list = split_tags(remove_tags)\n",
        "search_tags_list = split_tags(search_tags)\n",
        "replace_with_list = split_tags(replace_with)\n",
        "replace_new_list = [t for t in replace_with_list if t not in search_tags_list]\n",
        "\n",
        "replace_with_list = [t for t in replace_with_list if t not in replace_new_list]\n",
        "replace_new_list.reverse()\n",
        "activation_tag_list.reverse()\n",
        "\n",
        "remove_count = 0\n",
        "replace_count = 0\n",
        "\n",
        "for txt in [f for f in os.listdir(images_folder) if f.lower().endswith(\".txt\")]:\n",
        "\n",
        "  with open(os.path.join(images_folder, txt), 'r') as f:\n",
        "    tags = [s.strip() for s in f.read().split(\",\")]\n",
        "\n",
        "  if remove_duplicates:\n",
        "    tags = list(set(tags))\n",
        "  if sort_alphabetically:\n",
        "    tags.sort()\n",
        "\n",
        "  for rem in remove_tags_list:\n",
        "    if rem in tags:\n",
        "      remove_count += 1\n",
        "      tags.remove(rem)\n",
        "\n",
        "  if \"AND\" in search_mode and all(r in tags for r in search_tags_list) \\\n",
        "      or \"OR\" in search_mode and any(r in tags for r in search_tags_list):\n",
        "    replace_count += 1\n",
        "    for rem in search_tags_list:\n",
        "      if rem in tags:\n",
        "        tags.remove(rem)\n",
        "    for add in replace_with_list:\n",
        "      if add not in tags:\n",
        "        tags.append(add)\n",
        "    for new in replace_new_list:\n",
        "      if new_becomes_activation_tag:\n",
        "        if new in tags:\n",
        "          tags.remove(new)\n",
        "        tags.insert(0, new)\n",
        "      else:\n",
        "        if new not in tags:\n",
        "          tags.append(new)\n",
        "\n",
        "  for act in activation_tag_list:\n",
        "    if act in tags:\n",
        "      tags.remove(act)\n",
        "    tags.insert(0, act)\n",
        "\n",
        "  with open(os.path.join(images_folder, txt), 'w') as f:\n",
        "    f.write(\", \".join(tags))\n",
        "\n",
        "if global_activation_tag:\n",
        "  print(f\"\\n📎 Applied new activation tag(s): {', '.join(activation_tag_list)}\")\n",
        "if remove_tags:\n",
        "  print(f\"\\n🚮 Removed {remove_count} tags.\")\n",
        "if search_tags:\n",
        "  print(f\"\\n💫 Replaced in {replace_count} files.\")\n",
        "print(\"\\n✅ Done! Check your updated tags in the Extras below.\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HuJB7BGAyZCw"
      },
      "outputs": [],
      "source": [
        "#@markdown ### 6️⃣ Ready\n",
        "#@markdown You should be ready to [train your Lora](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer.ipynb)!\n",
        "\n",
        "from IPython.display import Markdown, display\n",
        "display(Markdown(f\"### 🦀 [Click here to open the Lora trainer](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer.ipynb)\"))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gDB9GXRONfiU"
      },
      "source": [
        "## *️⃣ Extras"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "xEsqOglcc6hA"
      },
      "outputs": [],
      "source": [
        "if \"step1_installed_flag\" not in globals():\n",
        "  raise Exception(\"Please run step 1 first!\")\n",
        "\n",
        "#@markdown ### 📈 Analyze Tags\n",
        "#@markdown Perhaps you need another look at your dataset.\n",
        "show_top_tags = 50 #@param {type:\"number\"}\n",
        "\n",
        "from collections import Counter\n",
        "top_tags = Counter()\n",
        "\n",
        "for txt in [f for f in os.listdir(images_folder) if f.lower().endswith(\".txt\")]:\n",
        "  with open(os.path.join(images_folder, txt), 'r') as f:\n",
        "    top_tags.update([s.strip() for s in f.read().split(\",\")])\n",
        "\n",
        "top_tags = Counter(top_tags)\n",
        "print(f\"📊 Top {show_top_tags} tags:\")\n",
        "for k, v in top_tags.most_common(show_top_tags):\n",
        "  print(f\"{k} ({v})\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "x56xQYwuOz2V"
      },
      "outputs": [],
      "source": [
        "#@markdown ### 📂 Unzip dataset\n",
        "#@markdown It's much slower to upload individual files to your Drive, so you may want to upload a zip if you have your dataset in your computer.\n",
        "zip = \"/content/drive/MyDrive/Loras/example.zip\" #@param {type:\"string\"}\n",
        "extract_to = \"/content/drive/MyDrive/Loras/example/dataset\" #@param {type:\"string\"}\n",
        "\n",
        "import os, zipfile\n",
        "\n",
        "if not os.path.exists('/content/drive'):\n",
        "  from google.colab import drive\n",
        "  print(\"📂 Connecting to Google Drive...\")\n",
        "  drive.mount('/content/drive')\n",
        "\n",
        "os.makedirs(extract_to, exist_ok=True)\n",
        "\n",
        "with zipfile.ZipFile(zip, 'r') as f:\n",
        "  f.extractall(extract_to)\n",
        "\n",
        "print(\"✅ Done\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "dLetTcLVOvAE"
      },
      "outputs": [],
      "source": [
        "#@markdown ### 🔢 Count datasets\n",
        "#@markdown Google Drive makes it impossible to count the files in a folder, so this will show you the file counts in all folders and subfolders.\n",
        "folder = \"/content/drive/MyDrive/Loras\" #@param {type:\"string\"}\n",
        "\n",
        "import os\n",
        "from google.colab import drive\n",
        "\n",
        "if not os.path.exists('/content/drive'):\n",
        "    print(\"📂 Connecting to Google Drive...\\n\")\n",
        "    drive.mount('/content/drive')\n",
        "\n",
        "tree = {}\n",
        "exclude = (\"_logs\", \"/output\")\n",
        "for i, (root, dirs, files) in enumerate(os.walk(folder, topdown=True)):\n",
        "  dirs[:] = [d for d in dirs if all(ex not in d for ex in exclude)]\n",
        "  images = len([f for f in files if f.lower().endswith((\".png\", \".jpg\", \".jpeg\"))])\n",
        "  captions = len([f for f in files if f.lower().endswith(\".txt\")])\n",
        "  others = len(files) - images - captions\n",
        "  path = root[folder.rfind(\"/\")+1:]\n",
        "  tree[path] = None if not images else f\"{images:>4} images | {captions:>4} captions |\"\n",
        "  if tree[path] and others:\n",
        "    tree[path] += f\" {others:>4} other files\"\n",
        "\n",
        "pad = max(len(k) for k in tree)\n",
        "print(\"\\n\".join(f\"📁{k.ljust(pad)} | {v}\" for k, v in tree.items() if v))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "dt8PmSM5_iI0"
      },
      "outputs": [],
      "source": [
        "if \"step1_installed_flag\" not in globals():\n",
        "  raise Exception(\"Please run step 1 first!\")\n",
        "\n",
        "from PIL import Image\n",
        "import os\n",
        "Image.MAX_IMAGE_PIXELS = None\n",
        "\n",
        "#@markdown ### 🖼️ Reduce dataset filesize\n",
        "#@markdown This will convert all images in the project folder to jpeg, reducing filesize without affecting quality too much. This can also solve some errors.\n",
        "location = images_folder\n",
        "\n",
        "for dir in [d[0] for d in os.walk(location)]:\n",
        "    os.chdir(dir)\n",
        "    converted = False\n",
        "    for file_name in list(os.listdir(\".\")):\n",
        "        try:\n",
        "            # Convert png to jpeg\n",
        "            if file_name.endswith(\".png\"):\n",
        "                if not converted:\n",
        "                    print(f\"Converting {dir}\")\n",
        "                    converted = True\n",
        "                im = Image.open(file_name)\n",
        "                im = im.convert(\"RGB\")\n",
        "                new_file_name = os.path.splitext(file_name)[0] + \".jpeg\"\n",
        "                im.save(new_file_name, quality=95)\n",
        "                os.remove(file_name)\n",
        "                file_name = new_file_name\n",
        "            # Resize large jpegs\n",
        "            if file_name.endswith((\".jpeg\", \".jpg\")) and os.path.getsize(file_name) > 2000000:\n",
        "                if not converted:\n",
        "                    print(f\"Converting {dir}\")\n",
        "                    converted = True\n",
        "                im = Image.open(file_name)\n",
        "                im = im.resize((int(im.width/2), int(im.height/2)))\n",
        "                im.save(file_name, quality=95)\n",
        "            # Rename jpg to jpeg\n",
        "            if file_name.endswith(\".jpg\"):\n",
        "                if not converted:\n",
        "                    print(f\"Converting {dir}\")\n",
        "                new_file_name = os.path.splitext(file_name)[0] + \".jpeg\"\n",
        "                os.rename(file_name, new_file_name)\n",
        "        except Exception as e:\n",
        "            print(f\"An error occurred while processing {file_name}: {e}\")\n",
        "    if converted:\n",
        "        print(f\"Converted {dir}\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "y6PKW-LIr214"
      },
      "outputs": [],
      "source": [
        "if \"step1_installed_flag\" not in globals():\n",
        "  raise Exception(\"Please run step 1 first!\")\n",
        "\n",
        "#@markdown ### 🚮 Clean folder\n",
        "#@markdown Careful! Deletes all non-image files in the project folder.\n",
        "\n",
        "!find {images_folder} -type f ! \\( -name '*.png' -o -name '*.jpg' -o -name '*.jpeg' \\) -delete\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}