Spanish_Dataset_Maker.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rmCPmqFL6hCQ"
      },
      "source": [
        "# 📊 Preparador de Lora de Hollowstrawberry\n",
        "\n",
        "Basado en el trabajo de [Kohya_ss](https://github.com/kohya-ss/sd-scripts) y [Linaqruf](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-dreambooth.ipynb#scrollTo=-Z4w3lfFKLjr). ¡Gracias!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hDoUi4JM_lz9"
      },
      "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": "0BJTp5PVN_q-"
      },
      "source": [
        "| |GitHub|🇬🇧 English|🇪🇸 Spanish|\n",
        "|:--|:-:|:-:|:-:|\n",
        "| 🏠 **Origen** | [![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 ## 🚩 Empezar aquí\n",
        "\n",
        "#@markdown ### 1️⃣  Inicio\n",
        "#@markdown Esta celda cargará algunos requerimientos y creará las carpetas correspondientes en tu Google Drive. <p>\n",
        "#@markdown Tu nombre de proyecto será la carpeta donde trabajaremos. No se permiten espacios.\n",
        "nombre_proyecto = \"\" #@param {type:\"string\"}\n",
        "project_name = nombre_proyecto.strip()\n",
        "#@markdown La estructura de carpetas no importa y es por comodidad. Asegúrate de siempre elegir la misma. Me gusta organizar por proyecto.\n",
        "estructura_de_carpetas = \"Organizar por proyecto (MyDrive/Loras/nombre_proyecto/dataset)\" #@param [\"Organizar por categoría (MyDrive/lora_training/datasets/nombre_proyecto)\", \"Organizar por proyecto (MyDrive/Loras/nombre_proyecto/dataset)\"]\n",
        "folder_structure = estructura_de_carpetas\n",
        "\n",
        "if not project_name or any(c in project_name for c in \" .()\\\"'\\\\\") or project_name.count(\"/\") > 1:\n",
        "  print(\"Por favor elige un nombre válido.\")\n",
        "else:\n",
        "  if COLAB and not os.path.exists('/content/drive'):\n",
        "    from google.colab import drive\n",
        "    print(\"📂 Conectando a 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\"✅ ¡Proyecto {project_name} listo!\")\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(\"Por favor usa el paso 1 primero.\")\n",
        "\n",
        "import json\n",
        "import time\n",
        "from urllib.request import urlopen, Request\n",
        "\n",
        "#@markdown ### 2️⃣ Obtener imágenes\n",
        "\n",
        "#@markdown Obtendremos imágenes de la galería de anime llamada [Gelbooru](https://gelbooru.com/). Las imágenes se organizan por miles de tags que describen todo acerca de una imagen. <p>\n",
        "#@markdown * Si quieres encontrar y usar tus propias imágenes, ponlas dentro de la carpeta `Loras/nombre_proyecto/dataset` en tu Google Drive.\n",
        "#@markdown * Si quieres descargar capturas de episodios de anime, existe [este otro colab de otra persona](https://colab.research.google.com/drive/1oBSntB40BKzNmKceXUlkXzujzdQw-Ci7) aunque aquel es más complicado.\n",
        "\n",
        "#@markdown Hasta 1000 imágenes se descargarán en un minuto, no debes abusar de ello. <p>\n",
        "#@markdown Tus tags deben ser relevantes para lo que desees entrenar, y excluir elementos no deseados (el contenido explícito puede hacer más difícil el entrenamiento).\n",
        "#@markdown Las palabras van separadas por guionbajos, las tags van separadas por espacios, y usa - para excluir esa tag. También puedes incluir una puntuación mínima: `score:>10`\n",
        "tags = \"1girl -sex -bdsm -loli -greyscale -monochrome hatsune_miku\" #@param {type:\"string\"}\n",
        "#markdown Si una imagen supera esta resolución, se descargará una versión más pequeña.\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",
        "\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(\"🏭 Instalando...\\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 en la instalación, intentando continuar...\")\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 se encontraron resultados.\")\n",
        "    return\n",
        "\n",
        "  print(f\"🎯 Se encontraron {count} resultados\")\n",
        "  test_url = \"https://gelbooru.com/index.php?page=post&s=list&tags={}\".format(tags)\n",
        "  display(Markdown(f\"[¡Click aquí para verlos en tu navegador!]({test_url})\"))\n",
        "  print (f\"🔽 Se descargará en {images_folder.replace('/content/drive/', '')} (Aparecerá una confirmación aquí abajo, sino vuelve a correr esta celda)\")\n",
        "  inp = input(\"❓ Escribe \\\"si\\\" para continuar con la descarga: \")\n",
        "\n",
        "  if inp.lower().strip() not in (\"si\", \"sí\"):\n",
        "    print(\"❌ Download cancelled\")\n",
        "    return\n",
        "\n",
        "  print(\"📩 Obteniendo lista de imágenes...\")\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\"🌐 Enlaces guardados a {scrape_file}\\n\\n🔁 Descargando imágenes ...\\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✅ Se han descargado {new_img_count - old_img_count} imágenes.\")\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(\"Por favor usa el paso 1 primero.\")\n",
        "\n",
        "#@markdown ### 3️⃣ Filtra tus imágenes\n",
        "#@markdown Vamos a analizar y encontrar imágenes duplicadas usando la IA de FiftyOne, y las marcaremos con `eliminar`. <p>\n",
        "#@markdown Cuando termine este proceso de varios minutos, aparecerá una zona interactiva bajo esta celda que te permitirá ver todas tus imágenes y marcar las que no te gusten con `eliminar` ttambién. <p>\n",
        "#@markdown Si el área interactiva no carga después de un minuto, intenta activar las cookies o desactivar la protección del navegador para la página de Google Colab, ya que estos pueden interferir. <p>\n",
        "#@markdown En cualquier caso, cuando estés satisfecho puedes mandar Enter en la casilla de texto arriba de la zona interactiva para finalizar y guardar los cambios.\n",
        "#@markdown <p>&nbsp;<p>Qué tan similares deben ser dos imágenes para ser consideradas duplicadas. Recomiendo 0.97 a 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(\"🏭 Instalando...\\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 en la instalación, intentando continuar...\")\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: El archivo {non_images[0]} no es una imagen. Este programa no puede correr con otros archivos de por medio, lo lamento. Puedes usar los Extras para limpiar la carpeta.\")\n",
        "elif img_count == 0:\n",
        "  print(f\"💥 Error: No hay imágenes en {images_folder}\")\n",
        "else:\n",
        "  print(\"\\n💿 Analizando imágenes...\")\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 = [\"eliminar\", \"tiene_duplicados\"]\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(\"tiene_duplicados\")\n",
        "        sample.save()\n",
        "    else:\n",
        "      sample.tags.append(\"eliminar\")\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(\"❗ Espera un minuto mientras carga la zona interactiva. Si no carga lee arriba.\")\n",
        "  print(\"❗ Cuando esté listo verás una cuadrícula con tus imágenes.\")\n",
        "  print(\"❗ Al lado izquierdo selecciona la etiqueta \\\"eliminar\\\" para ver cuáles imágenes serán borradas.\")\n",
        "  print(\"❗ Puedes marcar tus propias imágenes no deseadas con \\\"eliminar\\\" al seleccionarlas y luego apretar el ícono de etiqueta en la parte de arriba.\")\n",
        "  input(\"⭕ Cuando termines, manda Enter aquí para guardar los cambios: \")\n",
        "\n",
        "  print(\"💾 Guardando...\")\n",
        "\n",
        "  marked = [s for s in dataset if \"eliminar\" in s.tags]\n",
        "  dataset.remove_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\"✅ Se han eliminado {len(marked)} imágenes no deseadas, quedan {len(os.listdir(images_folder))} imágenes.\")\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(\"Por favor usa el paso 1 primero.\")\n",
        "\n",
        "#@markdown ### 4️⃣ Generar descripciones\n",
        "#@markdown Usaremos inteligencia artificial para describir tus imágenes, específicamente [Waifu Diffusion](https://huggingface.co/SmilingWolf/wd-v1-4-swinv2-tagger-v2) en el caso de anime (etiquetas/tags) y [BLIP](https://huggingface.co/spaces/Salesforce/BLIP) en el caso de fotografías (subtítulos/captions).\n",
        "#@markdown Estas descripciones que van junto a tus imágenes mejoran notablemente la calidad de tu Lora a la hora de entrenar. <p>\n",
        "metodo = \"Anime\" #@param [\"Anime\", \"Fotografía\"]\n",
        "method = metodo\n",
        "#@markdown **Anime:** El umbral es el nivel de certeza al que debe llegar la IA para asignar cada tag. Menor umbral = Más tags. Recomiendo 0.35 a 0.5\n",
        "umbral = 0.35 #@param {type:\"slider\", min:0.0, max:1.0, step:0.01}\n",
        "tag_threshold = umbral\n",
        "tags_no_permitidas = \"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",
        "blacklist_tags = tags_no_permitidas\n",
        "#@markdown **Fotografías:** El mínmimo y máximo largo de cada subtítulo (medido en tokens/palabras).\n",
        "largo_minimo = 10 #@param {type:\"number\"}\n",
        "caption_min = largo_minimo\n",
        "largo_maximo = 75 #@param {type:\"number\"}\n",
        "caption_max = largo_maximo\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 \"Anime\" in method:\n",
        "  if \"step4a_installed_flag\" not in globals() or \"step4b_installed_flag\" in globals():\n",
        "    print(\"🏭 Instalando...\\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 flax==0.7.5\n",
        "    if not get_ipython().__dict__['user_ns']['_exit_code']:\n",
        "      clear_output()\n",
        "      step4a_installed_flag = True\n",
        "    else:\n",
        "      print(\"❌ Error en la instalación, intentando continuar...\")\n",
        "\n",
        "  print(\"🚶‍♂️ Iniciando programa...\\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\"📊 Finalizado. Aquí están las 50 tags más comunes en tus imágenes:\")\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(\"🏭 Instalando...\\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 flax==0.7.5 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 en la instalación, intentando continuar...\")\n",
        "\n",
        "  print(\"🚶‍♂️ Iniciando programa...\\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\"📊 Finalizado. Aquí hay {len(sample)} ejemplos de subtítulos en tus imágenes:\")\n",
        "    print(\"\".join(sample))\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "WBFik7accyDz"
      },
      "outputs": [],
      "source": [
        "if \"step1_installed_flag\" not in globals():\n",
        "  raise Exception(\"Por favor usa el paso 1 primero.\")\n",
        "\n",
        "#@markdown ### 5️⃣ Editar las tags\n",
        "#@markdown Paso opcional para tags de anime. Puedes correr este paso cuantas veces quieras. <p>\n",
        "\n",
        "#@markdown Añadir una palabra de activación a tu Lora, útil para mejorar el entrenamiento y usarlo en tus prompts. En el entrenamiento debes poner `keep_tokens` igual a 1.<p>\n",
        "#@markdown Si quitas tags comunes como el color de pelo/ojos éstas serán \"absorbidas\" por tu palabra de activación.\n",
        "palabra_de_activacion = \"hatsune miku\" #@param {type:\"string\"}\n",
        "global_activation_tag = palabra_de_activacion\n",
        "quitar_tags = \"\" #@param {type:\"string\"}\n",
        "remove_tags = quitar_tags\n",
        "#@markdown <p>&nbsp;<p> En esta zona avanzada puedes realizar reemplazos o combinaciones de tags para así mejorar su calidad. Puedes reemplazar 1 tag por varias, o varias por 1, o una por otra, etc. También puedes añadir palabras de activación específicas.\n",
        "buscar_tags = \"\" #@param {type:\"string\"}\n",
        "search_tags = buscar_tags\n",
        "reemplazar_con = \"\" #@param {type:\"string\"}\n",
        "replace_with = reemplazar_con\n",
        "modo_de_busqueda = \"OR (puede tener cualquiera)\" #@param [\"OR (puede tener cualquiera)\", \"AND (debe tener todo)\"]\n",
        "search_mode = modo_de_busqueda\n",
        "tag_nueva_se_convierte_en_palabra_de_activacion = False #@param {type:\"boolean\"}\n",
        "new_becomes_activation_tag = tag_nueva_se_convierte_en_palabra_de_activacion\n",
        "#@markdown Estas pueden ser útiles a veces. Ten cuidado, pueden quitar las palabras de activación previas.\n",
        "ordenar_alfabeticamente = False #@param {type:\"boolean\"}\n",
        "sort_alphabetically = ordenar_alfabeticamente\n",
        "quitar_duplicados = False #@param {type:\"boolean\"}\n",
        "remove_duplicates = quitar_duplicados\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 remove_tags:\n",
        "  print(f\"\\n🚮 Se han quitado {remove_count} tags.\")\n",
        "if search_tags:\n",
        "  print(f\"\\n💫 Se han hecho {replace_count} reemplazos.\")\n",
        "print(\"\\n✅ ¡Listo!\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HuJB7BGAyZCw"
      },
      "outputs": [],
      "source": [
        "#@markdown ### 6️⃣  Listo\n",
        "#@markdown Ahora debes estar listo para [entrenar tu Lora](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Spanish_Lora_Trainer.ipynb).\n",
        "\n",
        "from IPython.display import Markdown, display\n",
        "display(Markdown(f\"🦀 [Click aquí para abrir el colab de entrenamiento](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Spanish_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(\"Por favor usa el paso 1 primero.\")\n",
        "\n",
        "#@markdown ### 📈 Analizar tags\n",
        "#@markdown Volver a ver las tags más comunes en tus imágenes.\n",
        "ver_top = 50 #@param {type:\"number\"}\n",
        "show_top_tags = ver_top\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\"📊 Tus {show_top_tags} tags más comunes:\")\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": "hcAHDi3zPjtg"
      },
      "outputs": [],
      "source": [
        "#@markdown ### 📂 Extraer datos\n",
        "#@markdown Es lento subir muchos archivos pequeños, si quieres puedes subir un zip y extraerlo aquí.\n",
        "zip = \"/content/drive/MyDrive/Loras/warrior.zip\" #@param {type:\"string\"}\n",
        "extract_to = \"/content/drive/MyDrive/Loras/ejemplo/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": "7AdGcZRyPmtm"
      },
      "outputs": [],
      "source": [
        "#@markdown ### 🔢 Contar archivos\n",
        "#@markdown Google Drive hace imposible contar los archivos en una carpeta, por lo que aquí puedes ver la cantidad de archivos en carpetas y subcarpetas.\n",
        "carpeta = \"/content/drive/MyDrive/Loras\" #@param {type:\"string\"}\n",
        "folder = carpeta\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": "FnN-rD4m_l0L"
      },
      "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 ### 🖼️ Reducir tamaño de imágenes\n",
        "#@markdown Esto convertirá todas las imágenes en la carpeta del proyecto a jpeg, reduciendo el tamaño sin afectar mucho la calidad. Esto también puede solucionar algunos errores.\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": "l0tzRu6xBqj9"
      },
      "outputs": [],
      "source": [
        "if \"step1_installed_flag\" not in globals():\n",
        "  raise Exception(\"Por favor usa el paso 1 primero.\")\n",
        "\n",
        "#@markdown ### 🚮 Limpiar carpeta\n",
        "#@markdown Cuidado, borra todos los archivos que no sean imágenes de la carpeta del proyecto.\n",
        "\n",
        "!find {images_folder} -type f ! \\( -name '*.png' -o -name '*.jpg' -o -name '*.jpeg' \\) -delete"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}