Spanish_Lora_Trainer.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rmCPmqFL6hCQ"
      },
      "source": [
        "# ⭐ Entrenador 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": "vJ8clWTZEu-g"
      },
      "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": "dPQlB4djNm3C"
      },
      "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": "OglZzI_ujZq-"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import re\n",
        "import toml\n",
        "import shutil\n",
        "import zipfile\n",
        "from time import time\n",
        "from IPython.display import Markdown, display\n",
        "\n",
        "# These carry information from past executions\n",
        "if \"model_url\" in globals():\n",
        "  old_model_url = model_url\n",
        "else:\n",
        "  old_model_url = None\n",
        "if \"dependencies_installed\" not in globals():\n",
        "  dependencies_installed = False\n",
        "if \"model_file\" not in globals():\n",
        "  model_file = None\n",
        "\n",
        "# These may be set by other cells, some are legacy\n",
        "if \"custom_dataset\" not in globals():\n",
        "  custom_dataset = None\n",
        "if \"override_dataset_config_file\" not in globals():\n",
        "  override_dataset_config_file = None\n",
        "if \"override_config_file\" not in globals():\n",
        "  override_config_file = None\n",
        "if \"optimizer\" not in globals():\n",
        "  optimizer = \"AdamW8bit\"\n",
        "if \"optimizer_args\" not in globals():\n",
        "  optimizer_args = None\n",
        "if \"continue_from_lora\" not in globals():\n",
        "  continue_from_lora = \"\"\n",
        "if \"weighted_captions\" not in globals():\n",
        "  weighted_captions = False\n",
        "if \"adjust_tags\" not in globals():\n",
        "  adjust_tags = False\n",
        "if \"keep_tokens_weight\" not in globals():\n",
        "  keep_tokens_weight = 1.0\n",
        "\n",
        "COLAB = True # low ram\n",
        "XFORMERS = True\n",
        "SOURCE = \"https://github.com/kohya-ss/sd-scripts\"\n",
        "COMMIT = \"9a67e0df390033a89f17e70df5131393692c2a55\"\n",
        "BETTER_EPOCH_NAMES = True\n",
        "LOAD_TRUNCATED_IMAGES = True\n",
        "\n",
        "#@title ## 🚩 Empieza Aquí\n",
        "\n",
        "#@markdown ### ▶️ Base\n",
        "#@markdown El nombre de tu proyecto también es el nombre de la carpeta donde irán tus imágenes. 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",
        "#@markdown Decidir el modelo base de entrenamiento. Los modelos por defecto producen los resultados más limpios y consistentes. Puedes cambiarlo por un modelo propio si lo deseas.\n",
        "modelo_de_entrenamiento = \"Anime (animefull-final-pruned-fp16.safetensors)\" #@param [\"Anime (animefull-final-pruned-fp16.safetensors)\", \"AnyLora (AnyLoRA_noVae_fp16-pruned.ckpt)\", \"Stable Diffusion (sd-v1-5-pruned-noema-fp16.safetensors)\"]\n",
        "opcional_enlace_a_modelo_propio = \"\" #@param {type:\"string\"}\n",
        "modelo_propio_basado_en_sd2 = False #@param {type:\"boolean\"}\n",
        "custom_model_is_based_on_sd2 = modelo_propio_basado_en_sd2\n",
        "\n",
        "if opcional_enlace_a_modelo_propio:\n",
        "  model_url = opcional_enlace_a_modelo_propio\n",
        "elif \"AnyLora\" in modelo_de_entrenamiento:\n",
        "  model_url = \"https://huggingface.co/Lykon/AnyLoRA/resolve/main/AnyLoRA_noVae_fp16-pruned.ckpt\"\n",
        "elif \"Anime\" in modelo_de_entrenamiento:\n",
        "  model_url = \"https://huggingface.co/hollowstrawberry/stable-diffusion-guide/resolve/main/models/animefull-final-pruned-fp16.safetensors\"\n",
        "else:\n",
        "  model_url = \"https://huggingface.co/hollowstrawberry/stable-diffusion-guide/resolve/main/models/sd-v1-5-pruned-noema-fp16.safetensors\"\n",
        "\n",
        "#@markdown ### ▶️ Procesamiento <p>\n",
        "#@markdown La resolución de 512 es estándar en Stable Diffusion 1.5. No es necesario recortar o achicar, el proceso es automático.\n",
        "resolucion = 512 #@param {type:\"slider\", min:512, max:1024, step:128}\n",
        "resolution = resolucion\n",
        "#@markdown Esta opción va a voltear tus imágenes para así tener el doble y aprender mejor. <p>\n",
        "#@markdown **Desactiva esto si te importan los elementos asimétricos en tu Lora.**\n",
        "flip_aug = False #@param {type:\"boolean\"}\n",
        "#markdown Leave empty for no captions.\n",
        "caption_extension = \".txt\" #param {type:\"string\"}\n",
        "#@markdown Mezclar las tags de anime ayuda al aprendizaje. Una tag de activación va al inicio de cada archivo de texto y no se mezclará.\n",
        "mezclar_tags = True #@param {type:\"boolean\"}\n",
        "shuffle_caption = mezclar_tags\n",
        "tags_de_activacion = \"1\" #@param [0,1,2,3]\n",
        "keep_tokens = int(tags_de_activacion)\n",
        "\n",
        "#@markdown ### ▶️ Pasos <p>\n",
        "#@markdown Tus imágenes se repetirán este número de veces durante el entrenamiento. Recomiendo que el valor total sea entre 200 y 400.\n",
        "num_repeats = 10 #@param {type:\"number\"}\n",
        "#@markdown Cuánto tiempo deseas entrenar. Un buen punto de partida puede ser alrededor de 10 epochs o alrededor de 2000 pasos. <p>\n",
        "#@markdown Un epoch es una cantidad de pasos igual a: tu cantidad de imágenes multipliccada por sus repeticiones, y dividido en el batch size.\n",
        "unidad_preferida = \"Epochs\" #@param [\"Epochs\", \"Pasos\"]\n",
        "cuantos = 10 #@param {type:\"number\"}\n",
        "max_train_epochs = cuantos if unidad_preferida == \"Epochs\" else None\n",
        "max_train_steps = cuantos if unidad_preferida == \"Pasos\" else None\n",
        "#@markdown Guardar más epochs te permitirá comparar mejor el progreso de tu Lora.\n",
        "guardar_cada_cuantos_epochs = 1 #@param {type:\"number\"}\n",
        "save_every_n_epochs = guardar_cada_cuantos_epochs\n",
        "guardar_solo_ultimos_epochs = 10 #@param {type:\"number\"}\n",
        "keep_only_last_n_epochs = guardar_solo_ultimos_epochs\n",
        "if not save_every_n_epochs:\n",
        "  save_every_n_epochs = max_train_epochs\n",
        "if not keep_only_last_n_epochs:\n",
        "  keep_only_last_n_epochs = max_train_epochs\n",
        "#@markdown Un batch size mayor hace el entrenamiento más rápido, pero puede empeorar el aprendizaje. Se recomienda 2 o 3.\n",
        "batch_size = 2 #@param {type:\"slider\", min:1, max:8, step:1}\n",
        "train_batch_size = batch_size\n",
        "\n",
        "#@markdown ### ▶️ Aprendizaje\n",
        "#@markdown La tasa de aprendizaje es lo más importante. Si deseas entrenar más lento con muchas imágenes, o si tienes un dim y alpha altos, usa un unet de 2e-4 o menor. <p>\n",
        "#@markdown El text encoder ayuda a tu Lora a aprender conceptos un poco mejor. Se recomienda la mitad o un quinto del unet. Puedes dejarlo en 0 para algunos estilos.\n",
        "aprendizaje_unet = 5e-4 #@param {type:\"number\"}\n",
        "unet_lr = aprendizaje_unet\n",
        "aprendizaje_text_encoder = 1e-4 #@param {type:\"number\"}\n",
        "text_encoder_lr = aprendizaje_text_encoder\n",
        "#@markdown El scheduler es el algoritmo matemático que guiará el entrenamiento. Para personajes recomiendo `cosine_with_restarts` con un valor de 3. Si no estás seguro ponlo en `constant` e ignora el valor.\n",
        "scheduler = \"cosine_with_restarts\" #@param [\"constant\", \"cosine\", \"cosine_with_restarts\", \"constant_with_warmup\", \"linear\", \"polynomial\"]\n",
        "lr_scheduler = scheduler\n",
        "valor_de_scheduler = 3 #@param {type:\"number\"}\n",
        "lr_scheduler_number = valor_de_scheduler\n",
        "lr_scheduler_num_cycles = lr_scheduler_number if lr_scheduler == \"cosine_with_restarts\" else 0\n",
        "lr_scheduler_power = lr_scheduler_number if lr_scheduler == \"polynomial\" else 0\n",
        "#@markdown Pasos de calentamiento durante el entrenamiento para un inicio eficiente. Recomiendo dejarlo en 5%.\n",
        "calentamiento = 0.05 #@param {type:\"slider\", min:0.0, max:0.5, step:0.01}\n",
        "lr_warmup_ratio = calentamiento\n",
        "lr_warmup_steps = 0\n",
        "#@markdown Nueva función que hace el aprendizaje mucho más eficiente. Puede que tus Loras estén listos en la mitad de epochs. Se usará un valor de 5.0 como en la [investigación](https://arxiv.org/abs/2303.09556).\n",
        "min_snr_gamma = True #@param {type:\"boolean\"}\n",
        "min_snr_gamma_value = 5.0 if min_snr_gamma else None\n",
        "\n",
        "#@markdown ### ▶️ Estructura\n",
        "#@markdown LoRA es el clásico y útil para muchos usos. LoCon es bueno con estilos ya que aprende con más capas.\n",
        "lora_type = \"LoRA\" #@param [\"LoRA\", \"LoCon\"]\n",
        "\n",
        "#@markdown Aquí hay algunos valores recomendados para las opciones de abajo:\n",
        "\n",
        "#@markdown | type | network_dim | network_alpha | conv_dim | conv_alpha |\n",
        "#@markdown | :---: | :---: | :---: | :---: | :---: |\n",
        "#@markdown | LoRA | 16 | 8 |   |   |\n",
        "#@markdown | LoCon | 16 | 8 | 8 | 4 |\n",
        "\n",
        "#@markdown Un dim mayor equivale a un Lora más grande, pero no siempre es mejor. Se recomienda de 8 a 32, con un alpha igual a la mitad del dim.\n",
        "network_dim = 16 #@param {type:\"slider\", min:1, max:128, step:1}\n",
        "network_alpha = 8 #@param {type:\"slider\", min:1, max:128, step:1}\n",
        "#@markdown Los siguientes sólo aplican a las capas adicionales de LoCon.\n",
        "conv_dim = 8 #@param {type:\"slider\", min:1, max:64, step:1}\n",
        "conv_alpha = 4 #@param {type:\"slider\", min:1, max:64, step:1}\n",
        "\n",
        "network_module = \"networks.lora\"\n",
        "network_args = None\n",
        "if lora_type.lower() == \"locon\":\n",
        "  network_args = [f\"conv_dim={conv_dim}\", f\"conv_alpha={conv_alpha}\"]\n",
        "\n",
        "#@markdown ### ▶️ Listo\n",
        "#@markdown Ahora puedes correr esta celda apretando el botón circular a la izquierda. ¡Buena suerte!\n",
        "\n",
        "\n",
        "# 👩‍💻 Cool code goes here\n",
        "\n",
        "if optimizer.lower() == \"prodigy\" or \"dadapt\" in optimizer.lower():\n",
        "  if override_values_for_dadapt_and_prodigy:\n",
        "    unet_lr = 0.5\n",
        "    text_encoder_lr = 0.5\n",
        "    lr_scheduler = \"constant_with_warmup\"\n",
        "    lr_warmup_ratio = 0.05\n",
        "    network_alpha = network_dim\n",
        "\n",
        "  if not optimizer_args:\n",
        "    optimizer_args = [\"decouple=True\",\"weight_decay=0.01\",\"betas=[0.9,0.999]\"]\n",
        "    if optimizer == \"Prodigy\":\n",
        "      optimizer_args.extend([\"d_coef=2\",\"use_bias_correction=True\"])\n",
        "      if lr_warmup_ratio > 0:\n",
        "        optimizer_args.append(\"safeguard_warmup=True\")\n",
        "      else:\n",
        "        optimizer_args.append(\"safeguard_warmup=False\")\n",
        "\n",
        "root_dir = \"/content\" if COLAB else \"~/Loras\"\n",
        "deps_dir = os.path.join(root_dir, \"deps\")\n",
        "repo_dir = os.path.join(root_dir, \"kohya-trainer\")\n",
        "\n",
        "if \"/Loras\" in folder_structure:\n",
        "  main_dir      = os.path.join(root_dir, \"drive/MyDrive/Loras\") if COLAB else root_dir\n",
        "  log_folder    = os.path.join(main_dir, \"_logs\")\n",
        "  config_folder = os.path.join(main_dir, project_name)\n",
        "  images_folder = os.path.join(main_dir, project_name, \"dataset\")\n",
        "  output_folder = os.path.join(main_dir, project_name, \"output\")\n",
        "else:\n",
        "  main_dir      = os.path.join(root_dir, \"drive/MyDrive/lora_training\") if COLAB else root_dir\n",
        "  images_folder = os.path.join(main_dir, \"datasets\", project_name)\n",
        "  output_folder = os.path.join(main_dir, \"output\", project_name)\n",
        "  config_folder = os.path.join(main_dir, \"config\", project_name)\n",
        "  log_folder    = os.path.join(main_dir, \"log\")\n",
        "\n",
        "config_file = os.path.join(config_folder, \"training_config.toml\")\n",
        "dataset_config_file = os.path.join(config_folder, \"dataset_config.toml\")\n",
        "accelerate_config_file = os.path.join(repo_dir, \"accelerate_config/config.yaml\")\n",
        "\n",
        "def install_dependencies():\n",
        "  os.chdir(root_dir)\n",
        "  !git clone {SOURCE} {repo_dir}\n",
        "  os.chdir(repo_dir)\n",
        "  if COMMIT:\n",
        "    !git reset --hard {COMMIT}\n",
        "  !wget https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/train_network_wrapper.py -q -O train_network_wrapper.py\n",
        "  !wget https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/dracula.py -q -O dracula.py\n",
        "\n",
        "  !apt -y update -qq\n",
        "  !apt -y install aria2 -qq\n",
        "  !pip install accelerate==0.15.0 diffusers==0.10.2 transformers==4.26.0 bitsandbytes==0.41.3.post2 opencv-python==4.8.0.76 tensorflow torchvision==0.16.0 torchtext==0.16.0 torchaudio==2.1.0 jax==0.4.23 jaxlib==0.4.23\n",
        "  !pip install toml==0.10.2 ftfy==6.1.1 einops==0.6.0 timm==0.6.12 fairscale==0.4.13 albumentations==1.3.1 voluptuous==0.13.1 requests==2.31.0 pytorch-lightning==1.9.0 \n",
        "  !pip install safetensors lion_pytorch==0.0.6 dadaptation==3.1 prodigyopt==1.0 pygments\n",
        "  !pip install .\n",
        "  if XFORMERS:\n",
        "    !pip install xformers==0.0.22.post7\n",
        "\n",
        "  # patch kohya for minor stuff\n",
        "  if COLAB:\n",
        "    !sed -i \"s@cpu@cuda@\" library/model_util.py # low ram\n",
        "  if LOAD_TRUNCATED_IMAGES:\n",
        "    !sed -i 's/from PIL import Image/from PIL import Image, ImageFile\\nImageFile.LOAD_TRUNCATED_IMAGES=True/g' library/train_util.py # fix truncated jpegs error\n",
        "  if BETTER_EPOCH_NAMES:\n",
        "    !sed -i 's/{:06d}/{:02d}/g' library/train_util.py # make epoch names shorter\n",
        "    !sed -i 's/\".\" + args.save_model_as)/\"-{:02d}.\".format(num_train_epochs) + args.save_model_as)/g' train_network.py # name of the last epoch will match the rest\n",
        "\n",
        "  from accelerate.utils import write_basic_config\n",
        "  if not os.path.exists(accelerate_config_file):\n",
        "    write_basic_config(save_location=accelerate_config_file)\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",
        "def validate_dataset():\n",
        "  global lr_warmup_steps, lr_warmup_ratio, caption_extension, keep_tokens, keep_tokens_weight, weighted_captions, adjust_tags\n",
        "  supported_types = (\".png\", \".jpg\", \".jpeg\", \".webp\", \".bmp\")\n",
        "\n",
        "  print(\"\\n💿 Revisando archivos...\")\n",
        "  if not project_name.strip() or any(c in project_name for c in \" .()\\\"'\\\\/\"):\n",
        "    print(\"💥 Error: Por favor elije un nombre de proyecto válido.\")\n",
        "    return\n",
        "\n",
        "  if custom_dataset:\n",
        "    try:\n",
        "      datconf = toml.loads(custom_dataset)\n",
        "      datasets = [d for d in datconf[\"datasets\"][0][\"subsets\"]]\n",
        "    except:\n",
        "      print(f\"💥 Error: ¡Tu configuración de datos propia es inválida o tiene errores! Por favor revisa el ejemplo original.\")\n",
        "      return\n",
        "    reg = [d.get(\"image_dir\") for d in datasets if d.get(\"is_reg\", False)]\n",
        "    datasets_dict = {d[\"image_dir\"]: d[\"num_repeats\"] for d in datasets}\n",
        "    folders = datasets_dict.keys()\n",
        "    files = [f for folder in folders for f in os.listdir(folder)]\n",
        "    images_repeats = {folder: (len([f for f in os.listdir(folder) if f.lower().endswith(supported_types)]), datasets_dict[folder]) for folder in folders}\n",
        "  else:\n",
        "    reg = []\n",
        "    folders = [images_folder]\n",
        "    files = os.listdir(images_folder)\n",
        "    images_repeats = {images_folder: (len([f for f in files if f.lower().endswith(supported_types)]), num_repeats)}\n",
        "\n",
        "  for folder in folders:\n",
        "    if not os.path.exists(folder):\n",
        "      print(f\"💥 Error: La carpeta {folder.replace('/content/drive/', '')} no existe.\")\n",
        "      return\n",
        "  for folder, (img, rep) in images_repeats.items():\n",
        "    if not img:\n",
        "      print(f\"💥 Error: La carpeta {folder.replace('/content/drive/', '')} está vacía.\")\n",
        "      return\n",
        "  for f in files:\n",
        "    if not f.lower().endswith((\".txt\", \".npz\")) and not f.lower().endswith(supported_types):\n",
        "      print(f\"💥 Error: Archivo inválido encontrado: \\\"{f}\\\". Abortando.\")\n",
        "      return\n",
        "\n",
        "  if not [txt for txt in files if txt.lower().endswith(\".txt\")]:\n",
        "    caption_extension = \"\"\n",
        "  if continue_from_lora and not (continue_from_lora.endswith(\".safetensors\") and os.path.exists(continue_from_lora)):\n",
        "    print(f\"💥 Error: Archivo de continuar_lora inválido. Ejemplo: /content/drive/MyDrive/Loras/ejemplo.safetensors\")\n",
        "    return\n",
        "\n",
        "  pre_steps_per_epoch = sum(img*rep for (img, rep) in images_repeats.values())\n",
        "  steps_per_epoch = pre_steps_per_epoch/train_batch_size\n",
        "  total_steps = max_train_steps or int(max_train_epochs*steps_per_epoch)\n",
        "  estimated_epochs = int(total_steps/steps_per_epoch)\n",
        "  lr_warmup_steps = int(total_steps*lr_warmup_ratio)\n",
        "\n",
        "  for folder, (img, rep) in images_repeats.items():\n",
        "    print(\"📁\"+folder.replace(\"/content/drive/\", \"\") + (\" (Regularización)\" if folder in reg else \"\"))\n",
        "  print(f\"📈 Se encontraron {img} imágenes con {rep} repeticiones, equivalente a {img*rep} pasos.\")\n",
        "  print(f\"📉 Divide {pre_steps_per_epoch} pasos en {train_batch_size} batch size para obtener {steps_per_epoch} pasos por epoch.\")\n",
        "  if max_train_epochs:\n",
        "    print(f\"🔮 Habrá {max_train_epochs} epochs, para un total de alrededor de {total_steps} pasos totales.\")\n",
        "  else:\n",
        "    print(f\"🔮 Habrá {total_steps} pasos, divididos en {estimated_epochs} epochs y un poco más.\")\n",
        "\n",
        "  if total_steps > 10000:\n",
        "    print(\"💥 Error: Tus pasos totales on muy altos. Probablemente cometiste un error. Abortando...\") \n",
        "    return\n",
        "\n",
        "  if adjust_tags:\n",
        "    print(f\"\\n📎 Weighted tags: {'ON' if weighted_captions else 'OFF'}\")\n",
        "    if weighted_captions:\n",
        "      print(f\"📎 Will use {keep_tokens_weight} weight on {keep_tokens} activation tag(s)\")\n",
        "    print(\"📎 Adjusting tags...\")\n",
        "    adjust_weighted_tags(folders, keep_tokens, keep_tokens_weight, weighted_captions)\n",
        "\n",
        "  return True\n",
        "\n",
        "def adjust_weighted_tags(folders, keep_tokens: int, keep_tokens_weight: float, weighted_captions: bool):\n",
        "  weighted_tag = re.compile(r\"\\((.+?):[.\\d]+\\)(,|$)\")\n",
        "  for folder in folders:\n",
        "    for txt in [f for f in os.listdir(folder) if f.lower().endswith(\".txt\")]:\n",
        "      with open(os.path.join(folder, txt), 'r') as f:\n",
        "        content = f.read()\n",
        "      # reset previous changes\n",
        "      content = content.replace('\\\\', '')\n",
        "      content = weighted_tag.sub(r'\\1\\2', content)\n",
        "      if weighted_captions:\n",
        "        # re-apply changes\n",
        "        content = content.replace(r'(', r'\\(').replace(r')', r'\\)').replace(r':', r'\\:')\n",
        "        if keep_tokens_weight > 1:\n",
        "          tags = [s.strip() for s in content.split(\",\")]\n",
        "          for i in range(min(keep_tokens, len(tags))):\n",
        "            tags[i] = f'({tags[i]}:{keep_tokens_weight})'\n",
        "          content = \", \".join(tags)\n",
        "      with open(os.path.join(folder, txt), 'w') as f:\n",
        "        f.write(content)\n",
        "\n",
        "def create_config():\n",
        "  global dataset_config_file, config_file, model_file\n",
        "\n",
        "  if override_config_file:\n",
        "    config_file = override_config_file\n",
        "    print(f\"\\n⭕ Usando configuración propia {config_file}\")\n",
        "  else:\n",
        "    config_dict = {\n",
        "      \"additional_network_arguments\": {\n",
        "        \"unet_lr\": unet_lr,\n",
        "        \"text_encoder_lr\": text_encoder_lr,\n",
        "        \"network_dim\": network_dim,\n",
        "        \"network_alpha\": network_alpha,\n",
        "        \"network_module\": network_module,\n",
        "        \"network_args\": network_args,\n",
        "        \"network_train_unet_only\": True if text_encoder_lr == 0 else None,\n",
        "        \"network_weights\": continue_from_lora if continue_from_lora else None\n",
        "      },\n",
        "      \"optimizer_arguments\": {\n",
        "        \"learning_rate\": unet_lr,\n",
        "        \"lr_scheduler\": lr_scheduler,\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",
        "        \"lr_warmup_steps\": lr_warmup_steps if lr_scheduler != \"constant\" else None,\n",
        "        \"optimizer_type\": optimizer,\n",
        "        \"optimizer_args\": optimizer_args if optimizer_args else None,\n",
        "      },\n",
        "      \"training_arguments\": {\n",
        "        \"max_train_steps\": max_train_steps,\n",
        "        \"max_train_epochs\": max_train_epochs,\n",
        "        \"save_every_n_epochs\": save_every_n_epochs,\n",
        "        \"save_last_n_epochs\": keep_only_last_n_epochs,\n",
        "        \"train_batch_size\": train_batch_size,\n",
        "        \"noise_offset\": None,\n",
        "        \"clip_skip\": 2,\n",
        "        \"min_snr_gamma\": min_snr_gamma_value,\n",
        "        \"weighted_captions\": weighted_captions,\n",
        "        \"seed\": 42,\n",
        "        \"max_token_length\": 225,\n",
        "        \"xformers\": XFORMERS,\n",
        "        \"lowram\": COLAB,\n",
        "        \"max_data_loader_n_workers\": 8,\n",
        "        \"persistent_data_loader_workers\": True,\n",
        "        \"save_precision\": \"fp16\",\n",
        "        \"mixed_precision\": \"fp16\",\n",
        "        \"output_dir\": output_folder,\n",
        "        \"logging_dir\": log_folder,\n",
        "        \"output_name\": project_name,\n",
        "        \"log_prefix\": project_name,\n",
        "      },\n",
        "      \"model_arguments\": {\n",
        "        \"pretrained_model_name_or_path\": model_file,\n",
        "        \"v2\": custom_model_is_based_on_sd2,\n",
        "        \"v_parameterization\": True if custom_model_is_based_on_sd2 else None,\n",
        "      },\n",
        "      \"saving_arguments\": {\n",
        "        \"save_model_as\": \"safetensors\",\n",
        "      },\n",
        "      \"dreambooth_arguments\": {\n",
        "        \"prior_loss_weight\": 1.0,\n",
        "      },\n",
        "      \"dataset_arguments\": {\n",
        "        \"cache_latents\": True,\n",
        "      },\n",
        "    }\n",
        "\n",
        "    for key in config_dict:\n",
        "      if isinstance(config_dict[key], dict):\n",
        "        config_dict[key] = {k: v for k, v in config_dict[key].items() if v is not None}\n",
        "\n",
        "    with open(config_file, \"w\") as f:\n",
        "      f.write(toml.dumps(config_dict))\n",
        "    print(f\"\\n📄 Configuración guardada en {config_file}\")\n",
        "\n",
        "  if override_dataset_config_file:\n",
        "    dataset_config_file = override_dataset_config_file\n",
        "    print(f\"⭕ Usando configuración de datos propia {dataset_config_file}\")\n",
        "  else:\n",
        "    dataset_config_dict = {\n",
        "      \"general\": {\n",
        "        \"resolution\": resolution,\n",
        "        \"shuffle_caption\": shuffle_caption,\n",
        "        \"keep_tokens\": keep_tokens,\n",
        "        \"flip_aug\": flip_aug,\n",
        "        \"caption_extension\": caption_extension,\n",
        "        \"enable_bucket\": True,\n",
        "        \"bucket_reso_steps\": 64,\n",
        "        \"bucket_no_upscale\": False,\n",
        "        \"min_bucket_reso\": 320 if resolution > 640 else 256,\n",
        "        \"max_bucket_reso\": 1280 if resolution > 640 else 1024,\n",
        "      },\n",
        "      \"datasets\": toml.loads(custom_dataset)[\"datasets\"] if custom_dataset else [\n",
        "        {\n",
        "          \"subsets\": [\n",
        "            {\n",
        "              \"num_repeats\": num_repeats,\n",
        "              \"image_dir\": images_folder,\n",
        "              \"class_tokens\": None if caption_extension else project_name\n",
        "            }\n",
        "          ]\n",
        "        }\n",
        "      ]\n",
        "    }\n",
        "\n",
        "    for key in dataset_config_dict:\n",
        "      if isinstance(dataset_config_dict[key], dict):\n",
        "        dataset_config_dict[key] = {k: v for k, v in dataset_config_dict[key].items() if v is not None}\n",
        "\n",
        "    with open(dataset_config_file, \"w\") as f:\n",
        "      f.write(toml.dumps(dataset_config_dict))\n",
        "    print(f\"📄 Configuración de datos guardada en {dataset_config_file}\")\n",
        "\n",
        "def download_model():\n",
        "  global old_model_url, model_url, model_file\n",
        "  real_model_url = model_url.strip()\n",
        "\n",
        "  if real_model_url.lower().endswith((\".ckpt\", \".safetensors\")):\n",
        "    model_file = f\"/content{real_model_url[real_model_url.rfind('/'):]}\"\n",
        "  else:\n",
        "    model_file = \"/content/downloaded_model.safetensors\"\n",
        "    if os.path.exists(model_file):\n",
        "      !rm \"{model_file}\"\n",
        "\n",
        "  if m := re.search(r\"(?:https?://)?(?:www\\.)?huggingface\\.co/[^/]+/[^/]+/blob\", model_url):\n",
        "    real_model_url = real_model_url.replace(\"blob\", \"resolve\")\n",
        "  elif m := re.search(r\"(?:https?://)?(?:www\\.)?civitai\\.com/models/([0-9]+)\", model_url):\n",
        "    real_model_url = f\"https://civitai.com/api/download/models/{m.group(1)}\"\n",
        "\n",
        "  !aria2c \"{real_model_url}\" --console-log-level=warn -c -s 16 -x 16 -k 10M -d / -o \"{model_file}\"\n",
        "\n",
        "  if model_file.lower().endswith(\".safetensors\"):\n",
        "    from safetensors.torch import load_file as load_safetensors\n",
        "    try:\n",
        "      test = load_safetensors(model_file)\n",
        "      del test\n",
        "    except:\n",
        "      #if \"HeaderTooLarge\" in str(e):\n",
        "      new_model_file = os.path.splitext(model_file)[0]+\".ckpt\"\n",
        "      !mv \"{model_file}\" \"{new_model_file}\"\n",
        "      model_file = new_model_file\n",
        "      print(f\"Modelo renombrado a {os.path.splitext(model_file)[0]}.ckpt\")\n",
        "\n",
        "  if model_file.lower().endswith(\".ckpt\"):\n",
        "    from torch import load as load_ckpt\n",
        "    try:\n",
        "      test = load_ckpt(model_file)\n",
        "      del test\n",
        "    except:\n",
        "      return False\n",
        "\n",
        "  return True\n",
        "\n",
        "def main():\n",
        "  global dependencies_installed\n",
        "\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",
        "  for dir in (main_dir, deps_dir, repo_dir, log_folder, images_folder, output_folder, config_folder):\n",
        "    os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "  if not validate_dataset():\n",
        "    return\n",
        "\n",
        "  if not dependencies_installed:\n",
        "    print(\"\\n🏭 Instalando...\\n\")\n",
        "    t0 = time()\n",
        "    install_dependencies()\n",
        "    t1 = time()\n",
        "    dependencies_installed = True\n",
        "    print(f\"\\n✅ Instalación completada en {int(t1-t0)} segundos.\")\n",
        "  else:\n",
        "    print(\"\\n✅ Ya se ha realizado la instalación.\")\n",
        "\n",
        "  if old_model_url != model_url or not model_file or not os.path.exists(model_file):\n",
        "    print(\"\\n🔄 Descargando modelo...\")\n",
        "    if not download_model():\n",
        "      print(\"\\n💥 Error: El modelo que elegiste es inválido o está corrupto, o no se pudo encontrar. Recomiendo usar un enlace de huggingface o civitai.\")\n",
        "      return\n",
        "    print()\n",
        "  else:\n",
        "    print(\"\\n🔄 El modelo ya ha sido descargado.\\n\")\n",
        "\n",
        "  create_config()\n",
        "\n",
        "  print(\"\\n⭐ Iniciando entrenador...\\n\")\n",
        "  os.chdir(repo_dir)\n",
        "\n",
        "  !accelerate launch --config_file={accelerate_config_file} --num_cpu_threads_per_process=1 train_network_wrapper.py --dataset_config={dataset_config_file} --config_file={config_file}\n",
        "\n",
        "main()\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mBMUJ7BuvNcn"
      },
      "source": [
        "## *️⃣ Extras\n",
        "\n",
        "You can run these before starting the training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "sy9jU2yrdYar"
      },
      "outputs": [],
      "source": [
        "#@markdown ### 🔮 Optimizador\n",
        "#@markdown Si corres esta celda cambiarás el optimizador usado en el entrenamiento. Sino, el por defecto es `AdamW8bit`, el cual es recomendado.<p>\n",
        "#@markdown * Dadapt y Prodigy manejan la tasa de aprendizaje de forma automática, y son muy buenos con datasets pequeños. Puedes usarlos sin cambiar nada más aquí.<p>\n",
        "#@markdown Con Dadapt o Prodigy, los siguientes valores serán sobreescritos:<p>\n",
        "#@markdown `learning_rate=0.5`, `network_alpha=network_dim`, `lr_scheduler=\"constant_with_warmup\"`, `lr_warmup_ratio=0.05`<p>\n",
        "#@markdown Con Dadapt o Prodigy, si `optimizer_args` está vacío su valor será `decouple=True, weight_decay=0.01, betas=[0.9,0.999]`<p>\n",
        "#@markdown Y además con Prodigy: `d_coef=2, use_bias_correction=True, safeguard_warmup=True`<p>\n",
        "optimizer = \"Prodigy\" #@param [\"AdamW8bit\", \"Prodigy\", \"DAdaptation\", \"DadaptAdam\", \"DadaptLion\", \"AdamW\", \"Lion\", \"SGDNesterov\", \"SGDNesterov8bit\", \"AdaFactor\"]\n",
        "optimizer_args = \"\" #@param {type:\"string\"}\n",
        "splitter = \", \" if \", \" in optimizer_args else \",\"\n",
        "optimizer_args = [a.strip() for a in optimizer_args.split(splitter) if a]\n",
        "override_values_for_dadapt_and_prodigy = True #@param {type:\"boolean\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wd4916Eu1tb9"
      },
      "source": [
        "### 📚 Múltiples carpetas\n",
        "**Para usuarios avanzados:** Antes de iniciar el entrenamiento, puedes editar y correr la celda aquí abajo, la cual tiene un ejemplo para definir tus propias carpetas de imágenes con diferentes repeticiones.\n",
        "\n",
        "(El número de repeticiones de la celda principal será ignorado, y también la carpeta principal con el nombre del proyecto)\n",
        "\n",
        "Puedes hacer que una carpeta contenga imágenes de regularización con la frase `is_reg = true`\n",
        "También puedes poner distintos `keep_tokens`, `flip_aug`, etc."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y037lagnJWmn"
      },
      "outputs": [],
      "source": [
        "custom_dataset = \"\"\"\n",
        "[[datasets]]\n",
        "\n",
        "[[datasets.subsets]]\n",
        "image_dir = \"/content/drive/MyDrive/Loras/ejemplo/dataset/imagenes_buenas\"\n",
        "num_repeats = 3\n",
        "\n",
        "[[datasets.subsets]]\n",
        "image_dir = \"/content/drive/MyDrive/Loras/ejemplo/dataset/imagenes_normales\"\n",
        "num_repeats = 1\n",
        "\n",
        "\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W84Jxf-U2TIU"
      },
      "outputs": [],
      "source": [
        "custom_dataset = None"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-Yq5mNvcCy2l"
      },
      "outputs": [],
      "source": [
        "#@markdown ### 🤓 Otros\n",
        "#@markdown Estas opciones son innecesarias, pero pueden servir a algunas personas.\n",
        "\n",
        "#@markdown Weighted captions is a new feature that allows you to use (parentheses) to give more weight to certain tags in your dataset, same as in your webui prompts. <p>\n",
        "#@markdown Normal parentheses in your tags such as `(series names)` will need to be escaped like `\\(series names\\)`\n",
        "weighted_captions = False #@param {type:\"boolean\"}\n",
        "\n",
        "#markdown By enabling `adjust_tags`, you will let this colab modify your tags before running to automatically adjust to `weighted_captions` being on or off. <p>\n",
        "#markdown Then, you may increase `activation_tag_weight` to improve how effective your activation tag is.\n",
        "adjust_tags = False #param {type:\"boolean\"}\n",
        "activation_tag_weight = \"1.0\" #param [\"1.0\",\"1.1\",\"1.2\"]\n",
        "keep_tokens_weight = float(activation_tag_weight)\n",
        "\n",
        "#@markdown Here you can write a path in your Google Drive to load an existing Lora file to continue training on.<p>\n",
        "#@markdown **Warning:** It's not the same as one long training session. The epochs start from scratch, and it may have worse results.\n",
        "continue_from_lora = \"\" #@param {type:\"string\"}\n",
        "if continue_from_lora and not continue_from_lora.startswith(\"/content/drive/MyDrive\"):\n",
        "  import os\n",
        "  continue_from_lora = os.path.join(\"/content/drive/MyDrive\", continue_from_lora)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "WDjkp4scvPgE"
      },
      "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/mi_dataset.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(\"📂 Conectando a 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(\"✅ Listo\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "aKWlpsG0jrX3"
      },
      "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",
        "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(\"📂 Conectando a 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": "markdown",
      "metadata": {},
      "source": [
        "# 📈 Graficar resultados de entrenamiento\n",
        "Puedes hacer esto tras el entrenamiento. No es necesario a menos que sepas lo que haces.  \n",
        "Puede que la primera celda falle en cargar todos tus datos. Sigue intentando la segunda celda hasta que terminen de cargar."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir={log_folder}/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "from tensorboard import notebook\n",
        "notebook.display(port=6006, height=800) "
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}