DFL_2.0_latest_1.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "DFL_Colab_Latest.ipynb의 사본",
      "provenance": [],
      "collapsed_sections": [
        "JuVn21kt40Gw",
        "hqwOlJG4MdLC",
        "WTuyUxgdLA13",
        "avAcSL_uvtq_"
      ],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/dalasjoe-1/dfl/blob/master/DFL_2.0_latest_1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IYtWMzOvLQ3s"
      },
      "source": [
        "## ALL\n",
        "\n",
        "This cell runs JS code to automatic reconnect to runtime."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jtClEMAMLVHw",
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "76daa4b7-bdc9-45b7-d304-898195fa7116"
      },
      "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', force_remount=True)\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",
        "\n",
        "\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(\"Done.\")\n",
        "\n",
        "\n",
        "#@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_FFHQ/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",
        "\n",
        "  # fix linux warning\n",
        "  # /usr/lib/python3.6/multiprocessing/semaphore_tracker.py:143: UserWarning: semaphore_tracker: There appear to be 1 leaked semaphores to clean up at shutdown\n",
        "  fin = open(\"/usr/lib/python3.6/multiprocessing/semaphore_tracker.py\", \"rt\")\n",
        "  data = fin.read()\n",
        "  data = data.replace('if cache:', 'if False:')\n",
        "  fin.close()\n",
        "\n",
        "  fin = open(\"/usr/lib/python3.6/multiprocessing/semaphore_tracker.py\", \"wt\")\n",
        "  fin.write(data)\n",
        "  fin.close()\n",
        "else:\n",
        "  %cd /content/DeepFaceLab\n",
        "  !git pull\n",
        "\n",
        "!pip uninstall -y tensorflow\n",
        "!pip install -r /content/DeepFaceLab/requirements-colab.txt\n",
        "!pip install --upgrade scikit-image\n",
        "!apt-get install cuda-10-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!\")\n",
        "\n",
        "#@title Training\n",
        "Model = \"SAEHD\" #@param [\"SAEHD\", \"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 -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 -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 -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\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[01:12:29][#1515158][1526ms][0.3673][0.4468]\n",
            "[01:19:19][#1516401][0361ms][0.3712][0.2890]"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L_xtDLKMWKYS"
      },
      "source": [
        "## Merge\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cNy-j81IWBjr",
        "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', force_remount=True)\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",
        "\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(\"Done.\")\n",
        "\n",
        "#@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_FFHQ/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",
        "\n",
        "  # fix linux warning\n",
        "  # /usr/lib/python3.6/multiprocessing/semaphore_tracker.py:143: UserWarning: semaphore_tracker: There appear to be 1 leaked semaphores to clean up at shutdown\n",
        "  fin = open(\"/usr/lib/python3.6/multiprocessing/semaphore_tracker.py\", \"rt\")\n",
        "  data = fin.read()\n",
        "  data = data.replace('if cache:', 'if False:')\n",
        "  fin.close()\n",
        "\n",
        "  fin = open(\"/usr/lib/python3.6/multiprocessing/semaphore_tracker.py\", \"wt\")\n",
        "  fin.write(data)\n",
        "  fin.close()\n",
        "else:\n",
        "  %cd /content/DeepFaceLab\n",
        "  !git pull\n",
        "\n",
        "!pip uninstall -y tensorflow\n",
        "!pip install -r /content/DeepFaceLab/requirements-colab.txt\n",
        "!pip install --upgrade scikit-image\n",
        "!apt-get install cuda-10-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!\")\n",
        "\n",
        "#@title Merge\n",
        "Model = \"SAEHD\" #@param [\"SAEHD\", \"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\n",
        "\n",
        "#@title Get result video and copy to Drive \n",
        "\n",
        "fileName = \"hjoo_cam8_SAEHD_seamless_dst_20_200_scale0_rct_sharp0.mp4\"\n",
        "\n",
        "!python DeepFaceLab/main.py videoed video-from-sequence --input-dir workspace/data_dst/merged --output-file workspace/$fileName --reference-file workspace/data_dst.mp4 --bitrate 16\n",
        "!cp /content/workspace/$fileName /content/drive/My\\ Drive/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KSPlD3nMRlE7",
        "cellView": "form"
      },
      "source": [
        "#@title Merge\n",
        "Model = \"SAEHD\" #@param [\"SAEHD\", \"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\n",
        "\n",
        "#@title Get result video and copy to Drive \n",
        "\n",
        "fileName = \"hjoo_cam8_SAEHD_seamless_dst_20_200_scale0_rct_sharp0.mp4\"\n",
        "\n",
        "!python DeepFaceLab/main.py videoed video-from-sequence --input-dir workspace/data_dst/merged --output-file workspace/$fileName --reference-file workspace/data_dst.mp4 --bitrate 16\n",
        "!cp /content/workspace/$fileName /content/drive/My\\ Drive/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JDWlEes1au1Y"
      },
      "source": [
        "# Merge Only"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BgQvE5QnacQS"
      },
      "source": [
        ""
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y_V6-aUEayVw",
        "cellView": "form"
      },
      "source": [
        "#@title Merge\n",
        "Model = \"SAEHD\" #@param [\"SAEHD\", \"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\n",
        "\n",
        "#@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",
        "fileName=\"eun_cam6_SAEHD_seamless_FANprd*FANdst_20_200_scale1_sharp5-2.mp4\"\n",
        "\n",
        "if Mode == \"result video\":\n",
        "  !python DeepFaceLab/main.py videoed video-from-sequence --input-dir workspace/data_dst/merged --output-file workspace/$fileName --reference-file workspace/data_dst.mp4 --include-audio --bitrate 16\n",
        "  if Copy_to_Drive:\n",
        "    !cp /content/workspace/$fileName /content/drive/My\\ Drive/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RPUHtHYUa8lX"
      },
      "source": [
        ""
      ],
      "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_FFHQ/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",
        "\n",
        "  # fix linux warning\n",
        "  # /usr/lib/python3.6/multiprocessing/semaphore_tracker.py:143: UserWarning: semaphore_tracker: There appear to be 1 leaked semaphores to clean up at shutdown\n",
        "  fin = open(\"/usr/lib/python3.6/multiprocessing/semaphore_tracker.py\", \"rt\")\n",
        "  data = fin.read()\n",
        "  data = data.replace('if cache:', 'if False:')\n",
        "  fin.close()\n",
        "\n",
        "  fin = open(\"/usr/lib/python3.6/multiprocessing/semaphore_tracker.py\", \"wt\")\n",
        "  fin.write(data)\n",
        "  fin.close()\n",
        "else:\n",
        "  %cd /content/DeepFaceLab\n",
        "  !git pull\n",
        "\n",
        "!pip install -r /content/DeepFaceLab/requirements-colab.txt\n",
        "!pip install --upgrade scikit-image\n",
        "!apt-get install cuda-10-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!\")\n",
        "\n"
      ],
      "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": "both"
      },
      "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', force_remount=True)\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": "both"
      },
      "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', force_remount=True)\n",
        "\n",
        "def zip_and_copy(path, mode):\n",
        "  zip_cmd=\"-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": "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": "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": "both"
      },
      "source": [
        "#@title Training\n",
        "Model = \"SAEHD\" #@param [\"SAEHD\", \"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 -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 -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 -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\", \"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": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bB_MfUqwbWN_"
      },
      "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",
        "* Apply or remove trained XSeg model to the extracted faces\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "IrfqcM_bbRfS"
      },
      "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": {
        "cellView": "form",
        "id": "iZUCNvKWbhBG"
      },
      "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": {
        "cellView": "form",
        "id": "YeA3da6JbjsZ"
      },
      "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": {
        "cellView": "form",
        "id": "U25lvDHubkpV"
      },
      "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": []
    }
  ]
}