Kohya Trainer XL Runpod.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d61cd191-c05c-4fdf-abf9-d8fc2e899679",
   "metadata": {},
   "source": [
    "# KOHYA TRAINER XL "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c57fd336-eac2-4917-ad00-b0ae7ca92467",
   "metadata": {},
   "source": [
    "# Install Kohya Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47275565-07d5-42a2-aee0-1500ef509fcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "\n",
    "# root_dir\n",
    "root_dir          = \"/workspace\"\n",
    "drive_dir         = os.path.join(root_dir, \"drive\", \"MyDrive\")\n",
    "repo_dir          = os.path.join(root_dir, \"kohya-trainer\")\n",
    "training_dir      = os.path.join(root_dir, \"fine_tune\")\n",
    "pretrained_model  = os.path.join(root_dir, \"pretrained_model\")\n",
    "vae_dir           = os.path.join(root_dir, \"vae\")\n",
    "lora_dir          = os.path.join(root_dir, \"network_weight\")\n",
    "config_dir        = os.path.join(training_dir, \"config\")\n",
    "output_dir        = os.path.join(training_dir, \"outputs\")\n",
    "tools_dir         = os.path.join(repo_dir, \"tools\")\n",
    "finetune_dir      = os.path.join(repo_dir, \"finetune\")\n",
    "accelerate_config = os.path.join(repo_dir, \"accelerate_config\", \"config.yaml\")\n",
    "\n",
    "repo_url          = \"https://github.com/qaneel/kohya-trainer\"\n",
    "branch            = \"main\" \n",
    "\n",
    "def clone_repo(url, dir, branch):\n",
    "    if not os.path.exists(dir):\n",
    "       !git clone -b {branch} {url} {dir}\n",
    "\n",
    "def install_dependencies():\n",
    "    !apt update -yqq\n",
    "    !apt install aria2 -yqq\n",
    "    !pip install -q --upgrade -r requirements.txt\n",
    "    !pip install xformers\n",
    "\n",
    "    from accelerate.utils import write_basic_config\n",
    "\n",
    "    if not os.path.exists(accelerate_config):\n",
    "        write_basic_config(save_location=accelerate_config)\n",
    "\n",
    "def prepare_environment():\n",
    "    os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
    "    os.environ[\"SAFETENSORS_FAST_GPU\"] = \"1\"\n",
    "    os.environ[\"PYTHONWARNINGS\"] = \"ignore\"\n",
    "\n",
    "def main():\n",
    "    os.chdir(root_dir)\n",
    "    clone_repo(repo_url, repo_dir, branch)\n",
    "    os.chdir(repo_dir)\n",
    "    for dir in [training_dir, config_dir, pretrained_model, vae_dir, output_dir]:\n",
    "        os.makedirs(dir, exist_ok=True)\n",
    "    install_dependencies()\n",
    "    prepare_environment()\n",
    "\n",
    "main()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29066e0b-e268-4b58-9ea4-3893bbe516fe",
   "metadata": {},
   "source": [
    "# Download SDXL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d1c8656-4323-4813-a785-6cee93bac903",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import requests\n",
    "import subprocess\n",
    "from urllib.parse import urlparse, unquote\n",
    "from pathlib import Path\n",
    "\n",
    "os.chdir(root_dir)\n",
    "\n",
    "HUGGINGFACE_TOKEN = \"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\"\n",
    "SDXL_MODEL_URL    = \"https://huggingface.co/stabilityai/stable-diffusion-xl-base-0.9/resolve/main/sd_xl_base_0.9.safetensors\"\n",
    "SDXL_VAE_URL      = \"https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors\"\n",
    "\n",
    "def get_supported_extensions():\n",
    "    return tuple([\".ckpt\", \".safetensors\", \".pt\", \".pth\"])\n",
    "\n",
    "def get_filename(url, bearer_token, quiet=True):\n",
    "    headers = {\"Authorization\": f\"Bearer {bearer_token}\"}\n",
    "    response = requests.get(url, headers=headers, stream=True)\n",
    "    response.raise_for_status()\n",
    "\n",
    "    if 'content-disposition' in response.headers:\n",
    "        content_disposition = response.headers['content-disposition']\n",
    "        filename = re.findall('filename=\"?([^\"]+)\"?', content_disposition)[0]\n",
    "    else:\n",
    "        url_path = urlparse(url).path\n",
    "        filename = unquote(os.path.basename(url_path))\n",
    "\n",
    "    return filename\n",
    "\n",
    "def parse_args(config):\n",
    "    args = []\n",
    "\n",
    "    for k, v in config.items():\n",
    "        if k.startswith(\"_\"):\n",
    "            args.append(f\"{v}\")\n",
    "        elif isinstance(v, str) and v is not None:\n",
    "            args.append(f'--{k}={v}')\n",
    "        elif isinstance(v, bool) and v:\n",
    "            args.append(f\"--{k}\")\n",
    "        elif isinstance(v, float) and not isinstance(v, bool):\n",
    "            args.append(f\"--{k}={v}\")\n",
    "        elif isinstance(v, int) and not isinstance(v, bool):\n",
    "            args.append(f\"--{k}={v}\")\n",
    "\n",
    "    return args\n",
    "\n",
    "def aria2_download(dir, filename, url, token):\n",
    "    user_header = f\"Authorization: Bearer {token}\"\n",
    "\n",
    "    aria2_config = {\n",
    "        \"console-log-level\"         : \"error\",\n",
    "        \"summary-interval\"          : 10,\n",
    "        \"header\"                    : user_header if \"huggingface.co\" in url else None,\n",
    "        \"continue\"                  : True,\n",
    "        \"max-connection-per-server\" : 16,\n",
    "        \"min-split-size\"            : \"1M\",\n",
    "        \"split\"                     : 16,\n",
    "        \"dir\"                       : dir,\n",
    "        \"out\"                       : filename,\n",
    "        \"_url\"                      : url,\n",
    "    }\n",
    "    aria2_args = parse_args(aria2_config)\n",
    "    subprocess.run([\"aria2c\", *aria2_args])\n",
    "\n",
    "def download(url, dst, token):\n",
    "    filename = get_filename(url, token, quiet=False)\n",
    "    filepath = os.path.join(dst, filename)\n",
    "\n",
    "    if url.startswith(\"/workspace\"):\n",
    "        return url\n",
    "    elif \"huggingface.co\" in url:\n",
    "        if \"/blob/\" in url:\n",
    "            url = url.replace(\"/blob/\", \"/resolve/\")\n",
    "                \n",
    "        aria2_download(dst, filename, url, token)\n",
    "\n",
    "    return filepath\n",
    "\n",
    "def main():\n",
    "    global model_path, vae_path\n",
    "\n",
    "    model_path = vae_path = None\n",
    "\n",
    "    download_targets = {\n",
    "        \"model\" : (SDXL_MODEL_URL, pretrained_model),\n",
    "        \"vae\"   : (SDXL_VAE_URL, vae_dir),\n",
    "    }\n",
    "    selected_files = {}\n",
    "\n",
    "    for target, (url, dst) in download_targets.items():\n",
    "        if url:\n",
    "            downloader = download(url, dst, HUGGINGFACE_TOKEN)\n",
    "            selected_files[target] = downloader\n",
    "\n",
    "            if target == \"model\":\n",
    "                model_path = selected_files[\"model\"] if not downloader else downloader\n",
    "            elif target == \"vae\":\n",
    "                vae_path = selected_files[\"vae\"] if not downloader else downloader\n",
    "\n",
    "    for category, path in {\n",
    "        \"model\": model_path,\n",
    "        \"vae\": vae_path,\n",
    "    }.items():\n",
    "        if path is not None and os.path.exists(path):\n",
    "            print(f\"Selected {category}: {path}\")\n",
    "\n",
    "main()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "525cdc14-d546-4cf4-a854-0877b65bb75f",
   "metadata": {},
   "source": [
    "# Directory Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63254a56-77e9-45d5-ac58-745a70c2459e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "train_data_dir = \"/workspace/fine_tune/train_data\"\n",
    "\n",
    "os.makedirs(train_data_dir, exist_ok=True)\n",
    "print(f\"Your train data directory : {train_data_dir}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5062b8e7-a53c-4072-983e-45d41d8cf771",
   "metadata": {},
   "source": [
    "# Data Gathering"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9eeba204-3abe-4e5e-af6f-3cac048e2407",
   "metadata": {},
   "source": [
    "## Unzip Dataset\n",
    "If your dataset is in a `zip` file and has been uploaded to a location, use this section to extract it. The dataset will be downloaded and automatically extracted to `train_data_dir` if `unzip_to` is empty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "240535e0-b9c0-4e17-b278-c8df0dc56ca5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import zipfile\n",
    "import shutil\n",
    "from pathlib import Path\n",
    "\n",
    "zipfile_url  = \"https://huggingface.co/datasets/Linaqruf/hitokomoru-lora-dataset/resolve/main/hitokomoru_dataset.zip\"\n",
    "unzip_to     = \"\"\n",
    "\n",
    "if unzip_to:\n",
    "    os.makedirs(unzip_to, exist_ok=True)\n",
    "else:\n",
    "    unzip_to = train_data_dir\n",
    "\n",
    "def extract_dataset(zip_file, output_path):\n",
    "    with zipfile.ZipFile(zip_file, \"r\") as zip_ref:\n",
    "        zip_ref.extractall(output_path)\n",
    "        \n",
    "def remove_files(train_dir, files_to_move):\n",
    "    for filename in os.listdir(train_dir):\n",
    "        file_path = os.path.join(train_dir, filename)\n",
    "        if filename in files_to_move:\n",
    "            if not os.path.exists(file_path):\n",
    "                shutil.move(file_path, training_dir)\n",
    "            else:\n",
    "                os.remove(file_path)\n",
    "\n",
    "zip_file = download(zipfile_url, root_dir, HUGGINGFACE_TOKEN)\n",
    "extract_dataset(zip_file, unzip_to)\n",
    "os.remove(zip_file)\n",
    "\n",
    "files_to_move = (\n",
    "    \"meta_cap.json\",\n",
    "    \"meta_cap_dd.json\",\n",
    "    \"meta_lat.json\",\n",
    "    \"meta_clean.json\",\n",
    ")\n",
    "\n",
    "remove_files(train_data_dir, files_to_move)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6e29913-5b62-455b-bd53-d293ee7c9840",
   "metadata": {},
   "source": [
    "# Bucketing and Latents Caching\n",
    "This code will create buckets based on the `bucket_resolution` provided for multi-aspect ratio training, and then convert all images within the `train_data_dir` to latents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e521a3f-7b1b-4214-8806-2df6a12f98a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# @title ## **3.4. Bucketing and Latents Caching**\n",
    "import time\n",
    "\n",
    "bucketing_json    = os.path.join(training_dir, \"meta_lat.json\")\n",
    "metadata_json     = os.path.join(training_dir, \"meta_clean.json\")\n",
    "bucket_resolution = 1024\n",
    "mixed_precision   = \"no\" # choose between [\"no\", \"fp16\", \"bf16\"]\n",
    "flip_aug          = False \n",
    "\n",
    "# Use `clean_caption` option to clean such as duplicate tags, `women` to `girl`, etc\n",
    "clean_caption     = True \n",
    "# Use the `recursive` option to process subfolders as well\n",
    "recursive         = True\n",
    "\n",
    "metadata_config = {\n",
    "    \"_train_data_dir\": train_data_dir,\n",
    "    \"_out_json\": metadata_json,\n",
    "    \"recursive\": recursive,\n",
    "    \"full_path\": recursive,\n",
    "    \"clean_caption\": clean_caption\n",
    "}\n",
    "\n",
    "bucketing_config = {\n",
    "    \"_train_data_dir\": train_data_dir,\n",
    "    \"_in_json\": metadata_json,\n",
    "    \"_out_json\": bucketing_json,\n",
    "    \"_model_name_or_path\": model_path,\n",
    "    \"recursive\": recursive,\n",
    "    \"full_path\": recursive,\n",
    "    \"flip_aug\": flip_aug,\n",
    "    \"batch_size\": 4,\n",
    "    \"max_data_loader_n_workers\": 2,\n",
    "    \"max_resolution\": f\"{bucket_resolution}, {bucket_resolution}\",\n",
    "    \"mixed_precision\": mixed_precision,\n",
    "}\n",
    "\n",
    "def generate_args(config):\n",
    "    args = \"\"\n",
    "    for k, v in config.items():\n",
    "        if k.startswith(\"_\"):\n",
    "            args += f'\"{v}\" '\n",
    "        elif isinstance(v, str):\n",
    "            args += f'--{k}=\"{v}\" '\n",
    "        elif isinstance(v, bool) and v:\n",
    "            args += f\"--{k} \"\n",
    "        elif isinstance(v, float) and not isinstance(v, bool):\n",
    "            args += f\"--{k}={v} \"\n",
    "        elif isinstance(v, int) and not isinstance(v, bool):\n",
    "            args += f\"--{k}={v} \"\n",
    "    return args.strip()\n",
    "\n",
    "merge_metadata_args = generate_args(metadata_config)\n",
    "prepare_buckets_args = generate_args(bucketing_config)\n",
    "\n",
    "merge_metadata_command = f\"python merge_all_to_metadata.py {merge_metadata_args}\"\n",
    "prepare_buckets_command = f\"python prepare_buckets_latents.py {prepare_buckets_args}\"\n",
    "\n",
    "os.chdir(finetune_dir)\n",
    "!{merge_metadata_command}\n",
    "time.sleep(1)\n",
    "!{prepare_buckets_command}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11483fb8-d0d1-4669-aa5f-b324708f614a",
   "metadata": {},
   "source": [
    "# Optimizer Config\n",
    "\n",
    "1. For `optimizer_type`, use `Adafactor` optimizer. `RMSprop 8bit` or `Adagrad 8bit` may work. `AdamW 8bit` doesn't seem to work.\n",
    "2. Choose between [\"AdamW\", \"AdamW8bit\", \"Lion8bit\", \"Lion\", \"SGDNesterov\", \"SGDNesterov8bit\", \"DAdaptation(DAdaptAdamPreprint)\", \"DAdaptAdaGrad\", \"DAdaptAdam\", \"DAdaptAdan\", \"DAdaptAdanIP\", \"DAdaptLion\", \"DAdaptSGD\", \"AdaFactor\"]\n",
    "3. Specify `optimizer_args` to add `additional` args for optimizer, e.g: `[\"weight_decay=0.6\"]`\n",
    "4. It's not recommended to Train Text Encoder for SDXL\n",
    "5. `lr_scheduler` provides several methods to adjust the learning rate based on the number of epochs.\n",
    "6. Choose between [\"linear\", \"cosine\", \"cosine_with_restarts\", \"polynomial\", \"constant\", \"constant_with_warmup\", \"adafactor\"]\n",
    "7. Specify `lr_scheduler_num` with `num_cycles` value for `cosine_with_restarts` or `power` value for `polynomial`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba8f1071-b7a1-4c9a-8550-99f8fc4bc01e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import toml\n",
    "import ast\n",
    "\n",
    "optimizer_type = \"AdaFactor\"  \n",
    "optimizer_args = \"[ \\\"scale_parameter=False\\\", \\\"relative_step=False\\\", \\\"warmup_init=False\\\" ]\"\n",
    "learning_rate = 4e-7\n",
    "train_text_encoder = False\n",
    "lr_scheduler = \"constant_with_warmup\" \n",
    "lr_warmup_steps = 100\n",
    "lr_scheduler_num = 0\n",
    "\n",
    "if isinstance(optimizer_args, str):\n",
    "    optimizer_args = optimizer_args.strip()\n",
    "    if optimizer_args.startswith('[') and optimizer_args.endswith(']'):\n",
    "        try:\n",
    "            optimizer_args = ast.literal_eval(optimizer_args)\n",
    "        except (SyntaxError, ValueError) as e:\n",
    "            print(f\"Error parsing optimizer_args: {e}\\n\")\n",
    "            optimizer_args = []\n",
    "    elif len(optimizer_args) > 0:\n",
    "        print(f\"WARNING! '{optimizer_args}' is not a valid list! Put args like this: [\\\"args=1\\\", \\\"args=2\\\"]\\n\")\n",
    "        optimizer_args = []\n",
    "    else:\n",
    "        optimizer_args = []\n",
    "else:\n",
    "    optimizer_args = []\n",
    "\n",
    "optimizer_config = {\n",
    "    \"optimizer_arguments\": {\n",
    "        \"optimizer_type\"          : optimizer_type,\n",
    "        \"learning_rate\"           : learning_rate,\n",
    "        \"train_text_encoder\"      : train_text_encoder,\n",
    "        \"max_grad_norm\"           : 1.0,\n",
    "        \"optimizer_args\"          : optimizer_args,\n",
    "        \"lr_scheduler\"            : lr_scheduler,\n",
    "        \"lr_warmup_steps\"         : lr_warmup_steps,\n",
    "        \"lr_scheduler_num_cycles\" : lr_scheduler_num if lr_scheduler == \"cosine_with_restarts\" else None,\n",
    "        \"lr_scheduler_power\"      : lr_scheduler_num if lr_scheduler == \"polynomial\" else None,\n",
    "        \"lr_scheduler_type\"       : None,\n",
    "        \"lr_scheduler_args\"       : None,\n",
    "    },\n",
    "}\n",
    "\n",
    "print(toml.dumps(optimizer_config))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "620ce942-f3f8-4d4f-b2b4-7c807cf35851",
   "metadata": {},
   "source": [
    "# Advanced Training Config\n",
    "1. Specify `optimizer_state_path` to resume training with Optimizer State\n",
    "2. You can't use both `noise_offset` and `multires_noise` at the same time\n",
    "3. Uncomment if necessary\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13fced50-aaeb-4101-936f-0689b0bf8e96",
   "metadata": {},
   "outputs": [],
   "source": [
    "import toml\n",
    "\n",
    "optimizer_state_path      = \"\" \n",
    "# noise_offset          = 0.1 \n",
    "# multires_noise_iterations = 6 \n",
    "# multires_noise_discount = 0.3\n",
    "min_snr_gamma             = -1 \n",
    "\n",
    "advanced_training_config = {\n",
    "    \"advanced_training_config\": {\n",
    "        \"resume\"                    : optimizer_state_path,\n",
    "        # \"noise_offset\"              : noise_offset, \n",
    "        # \"multires_noise_iterations\" : multires_noise_iterations, \n",
    "        # \"multires_noise_discount\"   : multires_noise_discount, \n",
    "        \"min_snr_gamma\"             : min_snr_gamma if not min_snr_gamma == -1 else None,\n",
    "\n",
    "    }\n",
    "}\n",
    "\n",
    "print(toml.dumps(advanced_training_config))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c5dce50-df19-4e60-b326-ea9841599e14",
   "metadata": {},
   "source": [
    "# Deployment Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31a22e72-2fcf-4131-a012-1297aee350f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import toml\n",
    "\n",
    "huggingface_repo_id = \"sdxl_finetune\"\n",
    "huggingface_write_token = \"\"\n",
    "huggingface_path_in_repo = \"\"\n",
    "huggingface_repo_visibility = \"private\" # private or public\n",
    "async_upload = True\n",
    "\n",
    "deployment_config = {\n",
    "    \"save_to_hub_config\": {\n",
    "        \"huggingface_repo_id\"         : huggingface_repo_id,\n",
    "        \"huggingface_repo_type\"       : \"model\", \n",
    "        \"huggingface_path_in_repo\"    : huggingface_path_in_repo, \n",
    "        \"huggingface_token\"           : huggingface_write_token,\n",
    "        \"async_upload\"                : async_upload, \n",
    "        \"huggingface_repo_visibility\" : huggingface_repo_visibility,\n",
    "    }\n",
    "}\n",
    "print(toml.dumps(deployment_config))\n",
    "\n",
    "if huggingface_write_token == \"\":\n",
    "    del deployment_config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87792429-20e3-464a-9ce4-95abcc6af155",
   "metadata": {},
   "source": [
    "# Training Config \n",
    "1. Get your `wandb_api_key` here: https://wandb.ai/settings\n",
    "2. `cache_text_encoder_outputs` is the recommended parameter for SDXL training but if you enable it, `shuffle_caption` won't work\n",
    "3. `min_timestep` and `max_timestep` can be used to train U-Net with different timesteps. The default values are 0 and 1000.\n",
    "4. Sampler List: [\"ddim\", \"pndm\", \"lms\", \"euler\", \"euler_a\", \"heun\", \"dpm_2\", \"dpm_2_a\", \"dpmsolver\",\"dpmsolver++\", \"dpmsingle\", \"k_lms\", \"k_euler\", \"k_euler_a\", \"k_dpm_2\", \"k_dpm_2_a\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79fa0e7f-9254-475f-95eb-96bc5c5cefd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import toml\n",
    "import os\n",
    "import random\n",
    "from subprocess import getoutput\n",
    "\n",
    "# PROJECT CONFIG\n",
    "project_name            = \"sdxl_finetune\"\n",
    "wandb_api_key           = \"\" \n",
    "in_json                 = \"/workspace/fine_tune/meta_lat.json\"\n",
    "\n",
    "# SDXL CONFIG\n",
    "grad_checkpointing      = True\n",
    "no_half_vae             = True \n",
    "cache_text_encoder_outputs = True\n",
    "min_timestep            = 0 \n",
    "max_timestep            = 1000\n",
    "\n",
    "# DATASET CONFIG\n",
    "num_repeats             = 1\n",
    "resolution              = 1024\n",
    "keep_tokens             = 0\n",
    "\n",
    "# GENERAL CONFIG\n",
    "max_train_steps         = 2500\n",
    "train_batch_size        = 4\n",
    "mixed_precision         = \"fp16\"\n",
    "seed                    = -1\n",
    "\n",
    "# SAVE OUTPUT AS\n",
    "save_precision          = \"fp16\" \n",
    "save_every_n_steps      = 1000\n",
    "save_optimizer_state    = False \n",
    "save_model_as           = \"safetensors\"\n",
    "\n",
    "# SAMPLE PROMPT\n",
    "prompt                  = \"masterpiece, best quality, 1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, hoop earrings, jewelry, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt\"\n",
    "custom_negative_prompt  = \"lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry\"\n",
    "sample_interval         = 100 \n",
    "sampler                 = \"euler_a\" \n",
    "logging_dir             = os.path.join(training_dir, \"logs\")\n",
    "\n",
    "os.chdir(repo_dir)\n",
    "\n",
    "prompt_config = {\n",
    "    \"prompt\": {\n",
    "        \"negative_prompt\" : negative_prompt if not custom_negative_prompt else custom_negative_prompt,\n",
    "        \"width\"           : resolution,\n",
    "        \"height\"          : resolution,\n",
    "        \"scale\"           : 7,\n",
    "        \"sample_steps\"    : 28,\n",
    "        \"subset\"          : [\n",
    "            {\n",
    "                \"prompt\" : prompt,\n",
    "            }\n",
    "        ],\n",
    "    }\n",
    "}\n",
    "\n",
    "train_config = {\n",
    "    \"sdxl_arguments\": {\n",
    "        \"cache_text_encoder_outputs\" : cache_text_encoder_outputs,\n",
    "        \"no_half_vae\"                : no_half_vae,\n",
    "        \"min_timestep\"               : min_timestep,\n",
    "        \"max_timestep\"               : max_timestep,\n",
    "        \"shuffle_caption\"            : True if not cache_text_encoder_outputs else False,\n",
    "    },\n",
    "    \"model_arguments\": {\n",
    "        \"pretrained_model_name_or_path\" : model_path,\n",
    "        \"vae\"                           : vae_path,\n",
    "    },\n",
    "    \"dataset_arguments\": {\n",
    "        \"debug_dataset\"                 : False,\n",
    "        \"in_json\"                       : in_json,\n",
    "        \"train_data_dir\"                : train_data_dir,\n",
    "        \"dataset_repeats\"               : num_repeats,\n",
    "        \"keep_tokens\"                   : keep_tokens,\n",
    "        \"resolution\"                    : str(resolution) + ',' + str(resolution),\n",
    "        \"caption_dropout_rate\"          : 0,\n",
    "        \"caption_tag_dropout_rate\"      : 0,\n",
    "        \"caption_dropout_every_n_epochs\": 0,\n",
    "        \"color_aug\"                     : False,\n",
    "        \"face_crop_aug_range\"           : None,\n",
    "        \"token_warmup_min\"              : 1,\n",
    "        \"token_warmup_step\"             : 0,\n",
    "    },\n",
    "    \"training_arguments\": {\n",
    "        \"output_dir\"                    : output_dir,\n",
    "        \"output_name\"                   : project_name if project_name else \"last\",\n",
    "        \"save_precision\"                : save_precision,\n",
    "        \"save_every_n_steps\"            : save_every_n_steps,\n",
    "        \"save_n_epoch_ratio\"            : None,\n",
    "        \"save_last_n_epochs\"            : None,\n",
    "        \"save_state\"                    : None,\n",
    "        \"save_last_n_epochs_state\"      : None,\n",
    "        \"resume\"                        : None,\n",
    "        \"train_batch_size\"              : train_batch_size,\n",
    "        \"max_token_length\"              : 225,\n",
    "        \"mem_eff_attn\"                  : False,\n",
    "        \"xformers\"                      : True,\n",
    "        \"max_train_steps\"               : max_train_steps,\n",
    "        \"max_data_loader_n_workers\"     : 8,\n",
    "        \"persistent_data_loader_workers\": True,\n",
    "        \"seed\"                          : seed if seed > 0 else None,\n",
    "        \"gradient_checkpointing\"        : grad_checkpointing,\n",
    "        \"gradient_accumulation_steps\"   : 1,\n",
    "        \"mixed_precision\"               : mixed_precision,\n",
    "    },\n",
    "    \"logging_arguments\": {\n",
    "        \"log_with\"          : \"wandb\" if wandb_api_key else \"tensorboard\",\n",
    "        \"log_tracker_name\"  : project_name if wandb_api_key and not project_name == \"last\" else None,\n",
    "        \"logging_dir\"       : logging_dir,\n",
    "        \"log_prefix\"        : project_name if not wandb_api_key else None,\n",
    "    },\n",
    "    \"sample_prompt_arguments\": {\n",
    "        \"sample_every_n_steps\"    : sample_interval,\n",
    "        \"sample_every_n_epochs\"   : None,\n",
    "        \"sample_sampler\"          : sampler,\n",
    "    },\n",
    "    \"saving_arguments\": {\n",
    "        \"save_model_as\": \"safetensors\"\n",
    "    },\n",
    "}\n",
    "\n",
    "def write_file(filename, contents):\n",
    "    with open(filename, \"w\") as f:\n",
    "        f.write(contents)\n",
    "\n",
    "def eliminate_none_variable(config):\n",
    "    for key in config:\n",
    "        if isinstance(config[key], dict):\n",
    "            for sub_key in config[key]:\n",
    "                if config[key][sub_key] == \"\":\n",
    "                    config[key][sub_key] = None\n",
    "        elif config[key] == \"\":\n",
    "            config[key] = None\n",
    "\n",
    "    return config\n",
    "\n",
    "try:\n",
    "    train_config.update(optimizer_config)\n",
    "except NameError:\n",
    "    raise NameError(\"'optimizer_config' dictionary is missing. Please run  'Optimizer Config' cell.\")\n",
    "\n",
    "advanced_training_warning = False\n",
    "try:\n",
    "    train_config.update(advanced_training_config)\n",
    "except NameError:\n",
    "    advanced_training_warning = True\n",
    "    pass\n",
    "\n",
    "deployment_config_warning = False\n",
    "try:\n",
    "    train_config.update(deployment_config)\n",
    "except NameError:\n",
    "    deployment_config_warning = True\n",
    "    pass\n",
    "\n",
    "config_path         = os.path.join(config_dir, \"config_file.toml\")\n",
    "prompt_path         = os.path.join(config_dir, \"sample_prompt.toml\")\n",
    "\n",
    "config_str          = toml.dumps(eliminate_none_variable(train_config))\n",
    "prompt_str          = toml.dumps(eliminate_none_variable(prompt_config))\n",
    "\n",
    "write_file(config_path, config_str)\n",
    "write_file(prompt_path, prompt_str)\n",
    "\n",
    "print(config_str)\n",
    "\n",
    "if advanced_training_warning:\n",
    "    import textwrap\n",
    "    error_message = \"WARNING: This is not an error message, but the [advanced_training_config] dictionary is missing. Please run the 'Advanced Training Config' cell if you intend to use it, or continue to the next step.\"\n",
    "    wrapped_message = textwrap.fill(error_message, width=80)\n",
    "    print('\\033[38;2;204;102;102m' + wrapped_message + '\\033[0m\\n')\n",
    "    pass\n",
    "    \n",
    "if deployment_config_warning:\n",
    "    import textwrap\n",
    "    error_message = \"WARNING: This is not an error message, but the [deployment_config] dictionary is missing. Please run the 'Deployment Training Config' cell if you intend to use it, or continue to the next step.\"\n",
    "    wrapped_message = textwrap.fill(error_message, width=80)\n",
    "    print('\\033[38;2;204;102;102m' + wrapped_message + '\\033[0m\\n')\n",
    "    pass\n",
    "\n",
    "print(prompt_str)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0842452e-f924-4600-8d6b-517298f632c5",
   "metadata": {},
   "source": [
    "# Start Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c2740ef-14a6-4d93-8fc2-2a51cbb7446c",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import toml\n",
    "\n",
    "sample_prompt   = \"/workspace/fine_tune/config/sample_prompt.toml\"\n",
    "config_file     = \"/workspace/fine_tune/config/config_file.toml\"\n",
    "\n",
    "def read_file(filename):\n",
    "    with open(filename, \"r\") as f:\n",
    "        contents = f.read()\n",
    "    return contents\n",
    "\n",
    "def train(config):\n",
    "    args = \"\"\n",
    "    for k, v in config.items():\n",
    "        if k.startswith(\"_\"):\n",
    "            args += f'\"{v}\" '\n",
    "        elif isinstance(v, str):\n",
    "            args += f'--{k}=\"{v}\" '\n",
    "        elif isinstance(v, bool) and v:\n",
    "            args += f\"--{k} \"\n",
    "        elif isinstance(v, float) and not isinstance(v, bool):\n",
    "            args += f\"--{k}={v} \"\n",
    "        elif isinstance(v, int) and not isinstance(v, bool):\n",
    "            args += f\"--{k}={v} \"\n",
    "\n",
    "    return args\n",
    "\n",
    "accelerate_conf = {\n",
    "    \"config_file\" : accelerate_config,\n",
    "    \"num_cpu_threads_per_process\" : 1,\n",
    "}\n",
    "\n",
    "train_conf = {\n",
    "    \"sample_prompts\"  : sample_prompt if os.path.exists(sample_prompt) else None,\n",
    "    \"config_file\"     : config_file,\n",
    "    \"wandb_api_key\"   : wandb_api_key if wandb_api_key else None,\n",
    "}\n",
    "\n",
    "accelerate_args = train(accelerate_conf)\n",
    "train_args = train(train_conf)\n",
    "\n",
    "final_args = f\"accelerate launch {accelerate_args} sdxl_train.py {train_args}\"\n",
    "\n",
    "os.chdir(repo_dir)\n",
    "!{final_args}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9736e6bc-f923-4e7a-b37a-0748659c2756",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3960b795-2a3e-4464-abe4-500e89a693ce",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "from PIL import Image, ImageOps\n",
    "from IPython.display import display\n",
    "\n",
    "ckpt_path = \"/workspace/pretrained_model/sd_xl_base_0.9.safetensors\" \n",
    "prompt = \"1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, hoop earrings, jewelry, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt\" \n",
    "negative_prompt = \"lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry\" \n",
    "output_path = \"/workspace/tmp/\" \n",
    "resolution = \"1024,1024\"\n",
    "optimization = \"scaled dot-product attention\" # [\"xformers\", \"scaled dot-product attention\"]\n",
    "conditional_resolution = \"1024,1024\"\n",
    "steps = 28 \n",
    "sampler = \"euler_a\" \n",
    "scale = 7 \n",
    "seed = -1\n",
    "images_per_prompt = 1\n",
    "batch_size = 1 \n",
    "clip_skip = 2\n",
    "\n",
    "os.makedirs(output_path, exist_ok=True)\n",
    "\n",
    "separators = [\"*\", \"x\", \",\"]\n",
    "\n",
    "for separator in separators:\n",
    "    if separator in resolution:\n",
    "        width, height = [value.strip() for value in resolution.split(separator)]\n",
    "        original_width, original_height = [value.strip() for value in conditional_resolution.split(separator)]\n",
    "        break\n",
    "\n",
    "config = {\n",
    "    \"prompt\": prompt + \" --n \" + negative_prompt,\n",
    "    \"images_per_prompt\": images_per_prompt,\n",
    "    \"outdir\": output_path,\n",
    "    \"W\": width,\n",
    "    \"H\": height,\n",
    "    \"original_width\": original_width,\n",
    "    \"original_height\": original_height,\n",
    "    \"batch_size\": batch_size,\n",
    "    \"vae_batch_size\": 1,\n",
    "    \"no_half_vae\": True,\n",
    "    \"steps\": steps,\n",
    "    \"sampler\": sampler,\n",
    "    \"scale\": scale,\n",
    "    \"ckpt\": ckpt_path,\n",
    "    \"vae\": vae_path,\n",
    "    \"seed\": seed if seed > 0 else None,\n",
    "    \"fp16\": True,\n",
    "    \"sdpa\": True if optimization == \"scaled dot-product attention\" else False,\n",
    "    \"xformers\": True if optimization == \"xformers\" else False,\n",
    "    \"opt_channels_last\": True,\n",
    "    \"clip_skip\": clip_skip,\n",
    "    \"max_embeddings_multiples\": 3,\n",
    "}\n",
    "\n",
    "def display_results(count):\n",
    "    samples = os.listdir(output_path)\n",
    "    samples.sort(reverse=True)\n",
    "    samples = samples[:count]\n",
    "\n",
    "    for sample in samples:\n",
    "        if sample.endswith((\".png\", \".jpg\")):\n",
    "            image_path = os.path.join(output_path, sample)\n",
    "\n",
    "            if os.path.exists(image_path):\n",
    "                img = Image.open(image_path)\n",
    "                img = img.resize((512, 512))  # Resize the image to 512x512 pixels\n",
    "                display(img)\n",
    "                \n",
    "args = \"\"\n",
    "for k, v in config.items():\n",
    "    if k.startswith(\"_\"):\n",
    "        args += f'\"{v}\" '\n",
    "    elif isinstance(v, str):\n",
    "        args += f'--{k}=\"{v}\" '\n",
    "    elif isinstance(v, bool) and v:\n",
    "        args += f\"--{k} \"\n",
    "    elif isinstance(v, float) and not isinstance(v, bool):\n",
    "        args += f\"--{k}={v} \"\n",
    "    elif isinstance(v, int) and not isinstance(v, bool):\n",
    "        args += f\"--{k}={v} \"\n",
    "\n",
    "final_args = f\"python sdxl_gen_img.py {args}\"\n",
    "\n",
    "os.chdir(repo_dir)\n",
    "!{final_args}\n",
    "\n",
    "display_results(batch_size)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}