DFL_Colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "DFL_Colab.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "IYtWMzOvLQ3s",
        "BDg_jiQ9adQe",
        "JuVn21kt40Gw",
        "hqwOlJG4MdLC",
        "tUNVcbujhm00",
        "WTuyUxgdLA13",
        "avAcSL_uvtq_"
      ],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0cKdTCuv4tXh"
      },
      "source": [
        "# Welcome to DFL-Colab!\n",
        "\n",
        "This is an adapted version of the DFL for Google Colab.\n",
        "\n",
        "\n",
        "# Overview\n",
        "*   Extractor works in full functionality.\n",
        "*   Training can work without preview.\n",
        "*   Merger works in full functionality.\n",
        "*   You can import/export workspace with your Google Drive.\n",
        "*   Import/export and another manipulations with workspace you can do in \"Manage workspace\" block\n",
        "*   Google Colab machine active for 12 hours. DFL-Colab makes a backup of your workspace in training mode.\n",
        "*   Google does not like long-term heavy calculations. Therefore, for training more than two sessions in a row, use two Google accounts. It is recommended to split your training over 2 accounts, but you can use one Google Drive account to store your workspace.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IYtWMzOvLQ3s"
      },
      "source": [
        "## Prevent random disconnects\n",
        "\n",
        "This cell runs JS code to automatic reconnect to runtime."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jtClEMAMLVHw"
      },
      "source": [
        "import IPython\n",
        "from google.colab import output\n",
        "\n",
        "display(IPython.display.Javascript('''\n",
        " function ClickConnect(){\n",
        "   btn = document.querySelector(\"colab-connect-button\")\n",
        "   if (btn != null){\n",
        "     console.log(\"Click colab-connect-button\"); \n",
        "     btn.click() \n",
        "     }\n",
        "   \n",
        "   btn = document.getElementById('ok')\n",
        "   if (btn != null){\n",
        "     console.log(\"Click reconnect\"); \n",
        "     btn.click() \n",
        "     }\n",
        "  }\n",
        "  \n",
        "setInterval(ClickConnect,60000)\n",
        "'''))\n",
        "\n",
        "print(\"Done.\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BDg_jiQ9adQe"
      },
      "source": [
        "## Check GPU\n",
        "\n",
        "*   Google Colab can provide you with one of Tesla graphics cards: K80, T4, P4 or P100\n",
        "*   Here you can check the model of GPU before using DeepFaceLab\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WJe71S6gbzt3"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JuVn21kt40Gw"
      },
      "source": [
        "## Install or update DeepFaceLab\n",
        "\n",
        "* Install or update DeepFAceLab directly from Github\n",
        "* Requirements install is automatically\n",
        "* Automatically sets timer to prevent random disconnects\n",
        "* \"Download FFHQ\" option means to download high quality FFHQ dataset instead of CelebA. FFHQ takes up more memory, so it will take longer to download than CelebA. It is recommended to enable this option if you are doing pretrain."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JG-f2WqT4fLK",
        "cellView": "form"
      },
      "source": [
        "#@title Install or update DeepFaceLab from Github\n",
        "\n",
        "Mode = \"install\" #@param [\"install\", \"update\"]\n",
        "Download_FFHQ = False #@param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "pretrain_link = \"https://github.com/chervonij/DFL-Colab/releases/download/\"\n",
        "pretrain_link = pretrain_link+\"pretrain_GenericFFHQ/pretrain_FFHQ.zip\" if Download_FFHQ else pretrain_link+\"pretrain-CelebA/pretrain_CelebA.zip\"\n",
        "\n",
        "from pathlib import Path\n",
        "if (Mode == \"install\"):\n",
        "  !git clone https://github.com/iperov/DeepFaceLab.git\n",
        "  %cd \"/content/DeepFaceLab\"\n",
        "  #!git checkout 9ad9728b4021d1dff62905cce03e2157d0c0868d\n",
        "  %cd \"/content\"\n",
        "else:\n",
        "  %cd /content/DeepFaceLab\n",
        "  !git pull\n",
        "\n",
        "!pip install h5py==3.8.0\n",
        "!pip install opencv-python==4.7.0.72\n",
        "!pip install ffmpeg-python==0.2.0\n",
        "!pip install colorama==0.4.6\n",
        "!pip install tf2onnx==1.14.0\n",
        "\n",
        "if not Path(\"/content/pretrain\").exists():\n",
        "  print(\"Downloading Pretrain faceset ... \")\n",
        "  !wget -q --no-check-certificate -r $pretrain_link -O /content/pretrain_faceset.zip\n",
        "  !mkdir /content/pretrain\n",
        "  !unzip -q /content/pretrain_faceset.zip -d /content/pretrain/\n",
        "  !rm /content/pretrain_faceset.zip\n",
        "\n",
        "if not Path(\"/content/pretrain_Q96\").exists():\n",
        "  print(\"Downloading Q96 pretrained model ...\")\n",
        "  !wget -q --no-check-certificate -r 'https://github.com/chervonij/DFL-Colab/releases/download/Q96_model_pretrained/Q96_model_pretrained.zip' -O /content/pretrain_Q96.zip\n",
        "  !mkdir /content/pretrain_Q96\n",
        "  !unzip -q /content/pretrain_Q96.zip -d /content/pretrain_Q96/\n",
        "  !rm /content/pretrain_Q96.zip\n",
        "\n",
        "if not Path(\"/content/workspace\").exists():\n",
        "  !mkdir /content/workspace; mkdir /content/workspace/data_src; mkdir /content/workspace/data_src/aligned; mkdir /content/workspace/data_dst; mkdir /content/workspace/data_dst/aligned; mkdir /content/workspace/model  \n",
        "\n",
        "import IPython\n",
        "from google.colab import output\n",
        "\n",
        "display(IPython.display.Javascript('''\n",
        " function ClickConnect(){\n",
        "   btn = document.querySelector(\"colab-connect-button\")\n",
        "   if (btn != null){\n",
        "     console.log(\"Click colab-connect-button\"); \n",
        "     btn.click() \n",
        "     }\n",
        "   \n",
        "   btn = document.getElementById('ok')\n",
        "   if (btn != null){\n",
        "     console.log(\"Click reconnect\"); \n",
        "     btn.click() \n",
        "     }\n",
        "  }\n",
        "  \n",
        "setInterval(ClickConnect,60000)\n",
        "'''))\n",
        "\n",
        "print(\"\\nDone!\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hqwOlJG4MdLC"
      },
      "source": [
        "## Manage workspace\n",
        "\n",
        "\n",
        "\n",
        "*   You can import/export workspace or individual data, like model files with Google Drive\n",
        "*   Also, you can use HFS (HTTP Fileserver) for directly import/export you workspace from your computer\n",
        "*   You can clear all workspace or delete part of it\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z4w_sUzgOQmL",
        "cellView": "form"
      },
      "source": [
        "#@title Import from Drive\n",
        "\n",
        "Mode = \"workspace\" #@param [\"workspace\", \"data_src\", \"data_dst\", \"data_src aligned\", \"data_dst aligned\", \"models\"]\n",
        "Archive_name = \"workspace.zip\" #@param {type:\"string\"}\n",
        "\n",
        "#Mount Google Drive as folder\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "def zip_and_copy(path, mode):\n",
        "  unzip_cmd=\" -q \"+Archive_name\n",
        "  \n",
        "  %cd $path\n",
        "  copy_cmd = \"/content/drive/My\\ Drive/\"+Archive_name+\" \"+path\n",
        "  !cp $copy_cmd\n",
        "  !unzip $unzip_cmd    \n",
        "  !rm $Archive_name\n",
        "\n",
        "if Mode == \"workspace\":\n",
        "  zip_and_copy(\"/content\", \"workspace\")\n",
        "elif Mode == \"data_src\":\n",
        "  zip_and_copy(\"/content/workspace\", \"data_src\")\n",
        "elif Mode == \"data_dst\":\n",
        "  zip_and_copy(\"/content/workspace\", \"data_dst\")\n",
        "elif Mode == \"data_src aligned\":\n",
        "  zip_and_copy(\"/content/workspace/data_src\", \"aligned\")\n",
        "elif Mode == \"data_dst aligned\":\n",
        "  zip_and_copy(\"/content/workspace/data_dst\", \"aligned\")\n",
        "elif Mode == \"models\":\n",
        "  zip_and_copy(\"/content/workspace\", \"model\")\n",
        "  \n",
        "print(\"Done!\")\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0Y3WfuwoNXqC",
        "cellView": "form"
      },
      "source": [
        "#@title Export to Drive { form-width: \"30%\" }\n",
        "Mode = \"workspace\" #@param [\"workspace\", \"data_src\", \"data_dst\", \"data_src aligned\", \"data_dst aligned\", \"merged\", \"merged_mask\", \"models\", \"result video\", \"result_mask video\"]\n",
        "Archive_name = \"workspace.zip\" #@param {type:\"string\"}\n",
        "\n",
        "#Mount Google Drive as folder\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "def zip_and_copy(path, mode):\n",
        "  zip_cmd=\"-0 -r -q \"+Archive_name+\" \"\n",
        "  \n",
        "  %cd $path\n",
        "  zip_cmd+=mode\n",
        "  !zip $zip_cmd\n",
        "  copy_cmd = \" \"+Archive_name+\"  /content/drive/My\\ Drive/\"\n",
        "  !cp $copy_cmd\n",
        "  !rm $Archive_name\n",
        "\n",
        "if Mode == \"workspace\":\n",
        "  zip_and_copy(\"/content\", \"workspace\")\n",
        "elif Mode == \"data_src\":\n",
        "  zip_and_copy(\"/content/workspace\", \"data_src\")\n",
        "elif Mode == \"data_dst\":\n",
        "  zip_and_copy(\"/content/workspace\", \"data_dst\")\n",
        "elif Mode == \"data_src aligned\":\n",
        "  zip_and_copy(\"/content/workspace/data_src\", \"aligned\")\n",
        "elif Mode == \"data_dst aligned\":\n",
        "  zip_and_copy(\"/content/workspace/data_dst\", \"aligned\")\n",
        "elif Mode == \"merged\":\n",
        "  zip_and_copy(\"/content/workspace/data_dst\", \"merged\")\n",
        "elif Mode == \"merged_mask\":\n",
        "  zip_and_copy(\"/content/workspace/data_dst\", \"merged_mask\")\n",
        "elif Mode == \"models\":\n",
        "  zip_and_copy(\"/content/workspace\", \"model\")\n",
        "elif Mode == \"result video\":\n",
        "  !cp /content/workspace/result.mp4 /content/drive/My\\ Drive/\n",
        "elif Mode == \"result_mask video\":\n",
        "  !cp /content/workspace/result_mask.mp4 /content/drive/My\\ Drive/\n",
        "  \n",
        "print(\"Done!\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0hIvJtxwTGcb"
      },
      "source": [
        "#@title Import from URL{ form-width: \"30%\", display-mode: \"form\" }\n",
        "URL = \"http://\" #@param {type:\"string\"}\n",
        "Mode = \"unzip to content\" #@param [\"unzip to content\", \"unzip to content/workspace\", \"unzip to content/workspace/data_src\", \"unzip to content/workspace/data_src/aligned\", \"unzip to content/workspace/data_dst\", \"unzip to content/workspace/data_dst/aligned\", \"unzip to content/workspace/model\", \"download to content/workspace\"]\n",
        "\n",
        "import urllib.request\n",
        "from pathlib import Path\n",
        "\n",
        "def unzip(zip_path, dest_path):\n",
        "\n",
        "    \n",
        "  unzip_cmd = \" unzip -q \" + zip_path + \" -d \"+dest_path\n",
        "  !$unzip_cmd  \n",
        "  rm_cmd = \"rm \"+dest_path + url_path.name\n",
        "  !$rm_cmd\n",
        "  print(\"Unziped!\")\n",
        "  \n",
        "\n",
        "if Mode == \"unzip to content\":\n",
        "  dest_path = \"/content/\"\n",
        "elif Mode == \"unzip to content/workspace\":\n",
        "  dest_path = \"/content/workspace/\"\n",
        "elif Mode == \"unzip to content/workspace/data_src\":\n",
        "  dest_path = \"/content/workspace/data_src/\"\n",
        "elif Mode == \"unzip to content/workspace/data_src/aligned\":\n",
        "  dest_path = \"/content/workspace/data_src/aligned/\"\n",
        "elif Mode == \"unzip to content/workspace/data_dst\":\n",
        "  dest_path = \"/content/workspace/data_dst/\"\n",
        "elif Mode == \"unzip to content/workspace/data_dst/aligned\":\n",
        "  dest_path = \"/content/workspace/data_dst/aligned/\"\n",
        "elif Mode == \"unzip to content/workspace/model\":\n",
        "  dest_path = \"/content/workspace/model/\"\n",
        "elif Mode == \"download to content/workspace\":\n",
        "  dest_path = \"/content/workspace/\"\n",
        "\n",
        "if not Path(\"/content/workspace\").exists():\n",
        "  cmd = \"mkdir /content/workspace; mkdir /content/workspace/data_src; mkdir /content/workspace/data_src/aligned; mkdir /content/workspace/data_dst; mkdir /content/workspace/data_dst/aligned; mkdir /content/workspace/model\"\n",
        "  !$cmd\n",
        "\n",
        "url_path = Path(URL)\n",
        "urllib.request.urlretrieve ( URL, dest_path + url_path.name )\n",
        "\n",
        "if (url_path.suffix == \".zip\") and (Mode!=\"download to content/workspace\"):\n",
        "  unzip(dest_path + url_path.name, dest_path)\n",
        "\n",
        "  \n",
        "print(\"Done!\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7V1sc7rxNKLO",
        "cellView": "form"
      },
      "source": [
        "#@title Export to URL\n",
        "URL = \"http://\" #@param {type:\"string\"}\n",
        "Mode = \"upload workspace\" #@param [\"upload workspace\", \"upload data_src\", \"upload data_dst\", \"upload data_src aligned\", \"upload data_dst aligned\", \"upload merged\", \"upload model\", \"upload result video\"]\n",
        "\n",
        "cmd_zip = \"zip -0 -r -q \"\n",
        "\n",
        "def run_cmd(zip_path, curl_url):\n",
        "  cmd_zip = \"zip -0 -r -q \"+zip_path\n",
        "  cmd_curl = \"curl --silent -F \"+curl_url+\" -D out.txt > /dev/null\"\n",
        "  !$cmd_zip\n",
        "  !$cmd_curl\n",
        "\n",
        "\n",
        "if Mode == \"upload workspace\":\n",
        "  %cd \"/content\"\n",
        "  run_cmd(\"workspace.zip workspace/\",\"'data=@/content/workspace.zip' \"+URL)\n",
        "elif Mode == \"upload data_src\":\n",
        "  %cd \"/content/workspace\"\n",
        "  run_cmd(\"data_src.zip data_src/\", \"'data=@/content/workspace/data_src.zip' \"+URL)\n",
        "elif Mode == \"upload data_dst\":\n",
        "  %cd \"/content/workspace\"\n",
        "  run_cmd(\"data_dst.zip data_dst/\", \"'data=@/content/workspace/data_dst.zip' \"+URL)\n",
        "elif Mode == \"upload data_src aligned\":\n",
        "  %cd \"/content/workspace\"\n",
        "  run_cmd(\"data_src_aligned.zip data_src/aligned\", \"'data=@/content/workspace/data_src_aligned.zip' \"+URL )\n",
        "elif Mode == \"upload data_dst aligned\":\n",
        "  %cd \"/content/workspace\"\n",
        "  run_cmd(\"data_dst_aligned.zip data_dst/aligned/\", \"'data=@/content/workspace/data_dst_aligned.zip' \"+URL)\n",
        "elif Mode == \"upload merged\":\n",
        "  %cd \"/content/workspace/data_dst\"\n",
        "  run_cmd(\"merged.zip merged/\",\"'data=@/content/workspace/data_dst/merged.zip' \"+URL )\n",
        "elif Mode == \"upload model\":\n",
        "  %cd \"/content/workspace\"\n",
        "  run_cmd(\"model.zip model/\", \"'data=@/content/workspace/model.zip' \"+URL)\n",
        "elif Mode == \"upload result video\":\n",
        "  %cd \"/content/workspace\"\n",
        "  run_cmd(\"result.zip result.mp4\", \"'data=@/content/workspace/result.zip' \"+URL)\n",
        "  \n",
        "  \n",
        "!rm *.zip\n",
        "\n",
        "%cd \"/content\"\n",
        "print(\"Done!\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ta6ue_UGMkki",
        "cellView": "form"
      },
      "source": [
        "#@title Delete and recreate\n",
        "Mode = \"Delete and recreate workspace\" #@param [\"Delete and recreate workspace\", \"Delete models\", \"Delete data_src\", \"Delete data_src aligned\", \"Delete data_src video\", \"Delete data_dst\", \"Delete data_dst aligned\", \"Delete merged frames\"]\n",
        "\n",
        "%cd \"/content\" \n",
        "\n",
        "if Mode == \"Delete and recreate workspace\":\n",
        "  cmd = \"rm -r /content/workspace ; mkdir /content/workspace; mkdir /content/workspace/data_src; mkdir /content/workspace/data_src/aligned; mkdir /content/workspace/data_dst; mkdir /content/workspace/data_dst/aligned; mkdir /content/workspace/model\"  \n",
        "elif Mode == \"Delete models\":\n",
        "  cmd = \"rm -r /content/workspace/model/*\"\n",
        "elif Mode == \"Delete data_src\":\n",
        "  cmd = \"rm /content/workspace/data_src/*.png || rm -r /content/workspace/data_src/*.jpg\"\n",
        "elif Mode == \"Delete data_src aligned\":\n",
        "  cmd = \"rm -r /content/workspace/data_src/aligned/*\"\n",
        "elif Mode == \"Delete data_src video\":\n",
        "  cmd = \"rm -r /content/workspace/data_src.*\"\n",
        "elif Mode == \"Delete data_dst\":\n",
        "  cmd = \"rm /content/workspace/data_dst/*.png || rm /content/workspace/data_dst/*.jpg\"\n",
        "elif Mode == \"Delete data_dst aligned\":\n",
        "  cmd = \"rm -r /content/workspace/data_dst/aligned/*\"\n",
        "elif Mode == \"Delete merged frames\":\n",
        "  cmd = \"rm -r /content/workspace/data_dst/merged; rm -r /content/workspace/data_dst/merged_mask\"\n",
        "  \n",
        "!$cmd\n",
        "print(\"Done!\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tUNVcbujhm00"
      },
      "source": [
        "## Extract, sorting and faceset tools\n",
        "* Extract frames for SRC or DST video.\n",
        "* Denoise SRC or DST video. \"Factor\" param set intesity of denoising\n",
        "* Detect and align faces. If you need, you can get frames with debug landmarks.\n",
        "* Export workspace to Google Drive after extract and sort it manually (In \"Manage Workspace\" block)\n",
        "* You can enhance your facesets with DFL FacesetEnhancer.\n",
        "* Resize faceset to your model resolution. Since Colab doesn't have a powerful CPU, resizing samples during training increases iteration time. Faceset resize reduces iteration time by about 2x times. Don't forget to keep save original faceset on your PC.\n",
        "* Pack or unpack facesets with DFL packing tool.\n",
        "* Apply or remove trained XSeg model to the extracted faces.\n",
        "* Recommended for use, Generic XSeg model for auto segmentation.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qwJEbz5Nhot0",
        "cellView": "form"
      },
      "source": [
        "#@title Extract frames\n",
        "Video = \"data_src\" #@param [\"data_src\", \"data_dst\"]\n",
        "\n",
        "%cd \"/content\"\n",
        "\n",
        "cmd = \"DeepFaceLab/main.py videoed extract-video\"\n",
        "\n",
        "if Video == \"data_dst\":\n",
        "  cmd+= \" --input-file workspace/data_dst.* --output-dir workspace/data_dst/\"\n",
        "else:\n",
        "  cmd+= \" --input-file workspace/data_src.* --output-dir workspace/data_src/\"\n",
        "  \n",
        "!python $cmd"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bFmPo0s2lTil",
        "cellView": "form"
      },
      "source": [
        "#@title Denoise frames\n",
        "Data = \"data_src\" #@param [\"data_src\", \"data_dst\"]\n",
        "Factor = 1 #@param {type:\"slider\", min:1, max:20, step:1}\n",
        "\n",
        "cmd = \"DeepFaceLab/main.py videoed denoise-image-sequence --input-dir workspace/\"+Data+\" --factor \"+str(Factor)\n",
        "\n",
        "%cd \"/content\"\n",
        "!python $cmd"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nmq0Sj2bmq7d",
        "cellView": "form"
      },
      "source": [
        "#@title Detect faces\n",
        "Data = \"data_src\" #@param [\"data_src\", \"data_dst\"]\n",
        "Detector = \"S3FD\" #@param [\"S3FD\", \"S3FD (whole face)\"]\n",
        "Debug = False #@param {type:\"boolean\"}\n",
        "\n",
        "detect_type = \"s3fd\"\n",
        "dbg = \" --output-debug\" if Debug else \" --no-output-debug\"\n",
        "\n",
        "folder = \"workspace/\"+Data\n",
        "folder_aligned = folder+\"/aligned\"\n",
        "\n",
        "cmd = \"DeepFaceLab/main.py extract --input-dir \"+folder+\" --output-dir \"+folder_aligned\n",
        "cmd+=\" --detector \"+detect_type+\" --force-gpu-idxs 0\"+dbg\n",
        "\n",
        "if \"whole face\" in Detector:\n",
        "  cmd+=\" --face-type whole_face\" \n",
        "%cd \"/content\"\n",
        "!python $cmd"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TRNxUFE6p6Eu",
        "cellView": "form"
      },
      "source": [
        "#@title Sort aligned\n",
        "Data = \"data_src\" #@param [\"data_src\", \"data_dst\"]\n",
        "sort_type = \"hist\" #@param [\"blur\", \"motion-blur\", \"face-yaw\", \"face-pitch\", \"face-source-rect-size\", \"hist\", \"hist-dissim\", \"brightness\", \"hue\", \"black\", \"origname\", \"oneface\", \"final-by-blur\", \"final-by-size\", \"absdiff\"]\n",
        "\n",
        "cmd = \"DeepFaceLab/main.py sort --input-dir workspace/\"+Data+\"/aligned --by \"+sort_type\n",
        "\n",
        "%cd \"/content\"\n",
        "!python $cmd"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O5MbnVDyXkP7",
        "cellView": "form"
      },
      "source": [
        "#@title Faceset Enhancer\n",
        "Data = \"data_src\" #@param [\"data_src\", \"data_dst\"]\n",
        "\n",
        "data_path = \"/content/workspace/\"+Data+\"/aligned\"\n",
        "cmd = \"/content/DeepFaceLab/main.py facesettool enhance --input-dir \"+data_path\n",
        "!python $cmd"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "Hyg5SREuMx8Q"
      },
      "source": [
        "#@title Resize faceset\n",
        "Data = \"data_src\" #@param [\"data_src\", \"data_dst\"]\n",
        "\n",
        "cmd = \"/content/DeepFaceLab/main.py facesettool resize --input-dir /content/workspace/\" + \\\n",
        "      f\"{Data}/aligned\"\n",
        "\n",
        "!python $cmd"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ypLfPUNHZNEp",
        "cellView": "form"
      },
      "source": [
        "#@title Pack/Unpack aligned faceset\n",
        "\n",
        "Folder = \"data_src\" #@param [\"data_src\", \"data_dst\"]\n",
        "Mode = \"unpack\" #@param [\"pack\", \"unpack\"]\n",
        "\n",
        "cmd = \"/content/DeepFaceLab/main.py util --input-dir /content/workspace/\" + \\\n",
        "      f\"{Folder}/aligned --{Mode}-faceset\"\n",
        "\n",
        "!python $cmd"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-VVvtoBMGnrA",
        "cellView": "form"
      },
      "source": [
        "#@title Apply or remove XSeg mask to the faces\n",
        "Mode = \"Apply mask\" #@param [\"Apply mask\", \"Remove mask\"]\n",
        "Data = \"data_src\" #@param [\"data_src\", \"data_dst\"]\n",
        "GenericXSeg = True #@param {type:\"boolean\"}\n",
        "\n",
        "from pathlib import Path\n",
        "mode_arg = 'apply' if Mode == \"Apply mask\" else 'remove'\n",
        "\n",
        "if GenericXSeg and not Path('/content/GenericXSeg').exists():\n",
        "  print('Downloading Generic XSeg model ... ')\n",
        "  xseg_link = 'https://github.com/chervonij/DFL-Colab/releases/download/GenericXSeg/GenericXSeg.zip'\n",
        "  !mkdir /content/GenericXSeg\n",
        "  !wget -q --no-check-certificate -r $xseg_link -O /content/GenericXSeg.zip\n",
        "  !unzip -q /content/GenericXSeg.zip -d /content/GenericXSeg/\n",
        "  !rm /content/GenericXSeg.zip\n",
        "\n",
        "main_path = '/content/DeepFaceLab/main.py'\n",
        "data_path = f'/content/workspace/{Data}/aligned'\n",
        "model_path = '/content/workspace/model' if not GenericXSeg else '/content/GenericXSeg'\n",
        "\n",
        "cmd = f'{main_path} xseg {mode_arg} --input-dir {data_path} '\n",
        "cmd += f'--model-dir {model_path}' if mode_arg == 'apply' else ''\n",
        "\n",
        "!python $cmd"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WTuyUxgdLA13"
      },
      "source": [
        "## Train model\n",
        "\n",
        "* Choose your model type, but SAEHD is recommend for everyone\n",
        "* Set model options on output field\n",
        "* You can see preview manually, if go to model folder in filemanager and double click on preview.jpg file\n",
        "* Your workspace will be archived and upload to mounted Drive after 11 hours from start session\n",
        "* If you select \"Backup_every_hour\" option, your workspace will be backed up every hour.\n",
        "* Also, you can export your workspace manually in \"Manage workspace\" block\n",
        "* \"Silent_Start\" option provides to automatically start with best GPU and last used model. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z0Kya-PJLDhv",
        "cellView": "form"
      },
      "source": [
        "#@title Training\n",
        "Model = \"SAEHD\" #@param [\"SAEHD\", \"AMP\", \"Quick96\", \"XSeg\"]\n",
        "Backup_every_hour = True #@param {type:\"boolean\"}\n",
        "Silent_Start = True #@param {type:\"boolean\"}\n",
        "\n",
        "%cd \"/content\"\n",
        "\n",
        "#Mount Google Drive as folder\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "import psutil, os, time\n",
        "\n",
        "p = psutil.Process(os.getpid())\n",
        "uptime = time.time() - p.create_time()\n",
        "\n",
        "if (Backup_every_hour):\n",
        "  if not os.path.exists('workspace.zip'):\n",
        "    print(\"Creating workspace archive ...\")\n",
        "    !zip -0 -r -q workspace.zip workspace\n",
        "    print(\"Archive created!\")\n",
        "  else:\n",
        "    print(\"Archive exist!\")\n",
        "\n",
        "if (Backup_every_hour):\n",
        "  print(\"Time to end session: \"+str(round((43200-uptime)/3600))+\" hours\")\n",
        "  backup_time = str(3600)\n",
        "  backup_cmd = \" --execute-program -\"+backup_time+\" \\\"import os; os.system('zip -0 -r -q workspace.zip workspace/model'); os.system('cp /content/workspace.zip /content/drive/My\\ Drive/'); print('Backed up!') \\\"\" \n",
        "elif (round(39600-uptime) > 0):\n",
        "  print(\"Time to backup: \"+str(round((39600-uptime)/3600))+\" hours\")\n",
        "  backup_time = str(round(39600-uptime))\n",
        "  backup_cmd = \" --execute-program \"+backup_time+\" \\\"import os; os.system('zip -0 -r -q workspace.zip workspace'); os.system('cp /content/workspace.zip /content/drive/My\\ Drive/'); print('Backed up!') \\\"\" \n",
        "else:\n",
        "  print(\"Session expires in less than an hour.\")\n",
        "  backup_cmd = \"\"\n",
        "    \n",
        "cmd = \"DeepFaceLab/main.py train --training-data-src-dir workspace/data_src/aligned --training-data-dst-dir workspace/data_dst/aligned --pretraining-data-dir pretrain --model-dir workspace/model --model \"+Model\n",
        "\n",
        "if Model == \"Quick96\":\n",
        "  cmd+= \" --pretrained-model-dir pretrain_Q96\"\n",
        "\n",
        "if Silent_Start:\n",
        "  cmd+= \" --silent-start\"\n",
        "\n",
        "if (backup_cmd != \"\"):\n",
        "  train_cmd = (cmd+backup_cmd)\n",
        "else:\n",
        "  train_cmd = (cmd)\n",
        "\n",
        "!python $train_cmd"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "avAcSL_uvtq_"
      },
      "source": [
        "## Merge frames"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A3Y8K22Sv9Gn",
        "cellView": "form"
      },
      "source": [
        "#@title Merge\n",
        "Model = \"SAEHD\" #@param [\"SAEHD\", \"AMP\", \"Quick96\" ]\n",
        "\n",
        "cmd = \"DeepFaceLab/main.py merge --input-dir workspace/data_dst --output-dir workspace/data_dst/merged --output-mask-dir workspace/data_dst/merged_mask --aligned-dir workspace/data_dst/aligned --model-dir workspace/model --model \"+Model\n",
        "\n",
        "%cd \"/content\"\n",
        "!python $cmd"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JNeGfiZpxlnz",
        "cellView": "form"
      },
      "source": [
        "#@title Get result video \n",
        "Mode = \"result video\" #@param [\"result video\", \"result_mask video\"]\n",
        "Copy_to_Drive = True #@param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "if Mode == \"result video\":\n",
        "  !python DeepFaceLab/main.py videoed video-from-sequence --input-dir workspace/data_dst/merged --output-file workspace/result.mp4 --reference-file workspace/data_dst.mp4 --include-audio\n",
        "  if Copy_to_Drive:\n",
        "    !cp /content/workspace/result.mp4 /content/drive/My\\ Drive/\n",
        "elif Mode == \"result_mask video\":\n",
        "  !python DeepFaceLab/main.py videoed video-from-sequence --input-dir workspace/data_dst/merged_mask --output-file workspace/result_mask.mp4 --reference-file workspace/data_dst.mp4\n",
        "  if Copy_to_Drive:\n",
        "    !cp /content/workspace/result_mask.mp4 /content/drive/My\\ Drive/\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}