DeepFaceLab_Colab_deepartist_ai_AT_gmail_com.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "DeepFaceLab-Colab_deepartist.ai@gmail.com.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/deepartist/DeepFaceLab_Colab/blob/main/DeepFaceLab_Colab_deepartist_ai_AT_gmail_com.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ddmtN7ibpqI6"
      },
      "source": [
        "# DeepFaceLab (Training) Colab Version by [Deepartist](mailto:deepartist.ai@gmail.com)\n",
        "This is based on [DFL-Colab](https://github.com/chervonij/DFL-Colab), [DeepFaceLab-Linux](https://github.com/nagadit/DeepFaceLab_Linux)\n",
        "\n",
        "Features:\n",
        "*   GDrive rsyncing \n",
        "*   Backups via cronjob\n",
        "*   Training configs (MVE branch only)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BH6jGww8sBar"
      },
      "source": [
        "# Installation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GYzGfYygpCsT",
        "cellView": "form"
      },
      "source": [
        "#@title Check GPU and CPU  \n",
        "gpu = !nvidia-smi -L\n",
        "gpu= gpu[0]\n",
        "gpu_mem = !nvidia-smi --query-gpu=memory.total --format=csv\n",
        "if len(gpu_mem) >1:\n",
        "  gpu_mem = gpu_mem[1]\n",
        "  gpu_mem_free = !nvidia-smi --query-gpu=memory.free --format=csv\n",
        "  gpu_mem_free = gpu_mem_free[1]\n",
        "cores = !grep -c ^processor /proc/cpuinfo\n",
        "cores = cores[0]\n",
        "mhz = !grep \"^[c]pu MHz\" /proc/cpuinfo\n",
        "mhz = mhz[0].split(\":\")[-1]\n",
        "mhz\n",
        "!echo \"$gpu\"\n",
        "!echo Total Memory: \"$gpu_mem\" Free: \"$gpu_mem_free\"\n",
        "!echo CPU: \"$cores\" cores @\"$mhz\" MHz\n",
        "!echo RAM: $(($(getconf _PHYS_PAGES) * $(getconf PAGE_SIZE) / (1024 * 1024))) MB"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IRFCT4eYcMXr"
      },
      "source": [
        "**Run to mount your Google drive with auth key**\n",
        "\n",
        "1. Connect to your google drive to be able to import and export files to/from google drive\n",
        "2. Click on Link below\n",
        "3. Login and confirm\n",
        "4. Ctrl+V auth key and hit return key\n",
        "\n",
        "**OR skip and mount your Google drive permantly**\n",
        "\n",
        "1. Create new notebook\n",
        "2. Change runtime to GPU (Runtime -> Change runtime type -> Hardware Accelarator -> GPU)\n",
        "3. Click on files, and mount your google drive permantly \n",
        "4. Switch to this notenook\n",
        "4. Click in the first cell (DeepFaceLab Colab Version by Deepartist)\n",
        "5. Select all cells with shift + arrow down\n",
        "6. Copy all cells from this notebook (ctrl+c)\n",
        "8. Click in first cell in your new notebook and paste (ctrl+v)\n",
        "7. Save your notebook (in drive) -> No more auth token coyping\n",
        "\n",
        "**OR SIMPLY SKIP** if you dont want import and export files to/from google drive"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O2Dlj_TnVLFs",
        "cellView": "form"
      },
      "source": [
        "#@title Connect to your gdrive OR ignrore if permantly mounted\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive', use_metadata_server=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BCa4jhzKtX6S",
        "cellView": "form"
      },
      "source": [
        "#@title Enter puplic or private DeepFaceLab Git\n",
        "\n",
        "\n",
        "gitlink = 'https://github.com/MachineEditor/DeepFaceLab.git' #@param['https://github.com/MachineEditor/DeepFaceLab.git', 'https://github.com/iperov/DeepFaceLab.git', 'https://github.com/faceshiftlabs/DeepFaceLab.git', 'https://github.com/janfschr/DeepFaceLab_Fork.git']\n",
        "gitlink_overwrite = \"\" #@param{type:\"string\"}\n",
        "if gitlink_overwrite is not None and gitlink_overwrite.strip() is not \"\":\n",
        "  gitlink = gitlink_overwrite.strip()\n",
        "\n",
        "#@markdown Enter specific branch , empty is master \n",
        "branch = \"\" #@param {type:\"string\"}\n",
        "\n",
        "if branch is not None and branch.strip() is not \"\":\n",
        "  gitlink+= \" -b \"+branch.strip()\n",
        "\n",
        "#@markdown Enter specific commit (full or short sha)\n",
        "commit = \"\" #@param {type:\"string\"}\n",
        "\n",
        "if commit is not None and commit.strip() is not \"\":\n",
        "  git_commit = commit.strip()\n",
        "#@markdown Check if private, enter user and enter password or enter hidden in prombt \n",
        "is_private = False #@param {type:\"boolean\"}\n",
        "user = 'janfschr' #@param {type:\"string\"}\n",
        "clear_password = 'not_safu' #@param {type:\"string\"}\n",
        "\n",
        "if is_private:\n",
        "  import urllib.parse\n",
        "  if user is None or user.strip() is \"\": raise ValueError('Empty username')\n",
        "  if clear_password is None or clear_password.strip() =='':\n",
        "    from getpass import getpass\n",
        "    password = getpass('Enter your password').strip()\n",
        "    password = urllib.parse.quote(password)\n",
        "  else: password = clear_password.strip()\n",
        "  gitlink = gitlink.replace(\"https://\", \"\")\n",
        "  hidden = \"******\"\n",
        "  print(f\"https://{user.strip()}:{password if clear_password else hidden}@\"+gitlink)\n",
        "  gitlink = f\"https://{user.strip()}:{password}@\"+gitlink\n",
        "else: \n",
        "  print(gitlink)\n",
        "\n",
        "!echo \"$gitlink\" > /content/gitlink\n",
        "!echo \"$git_commit\" > /content/git_commit"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mkSvrFbjuyCd",
        "cellView": "form"
      },
      "source": [
        "#@title Install or update DeepFaceLab from Github\n",
        "#%tensorflow_version 2.x\n",
        "!rm -rf /content/sample_data\n",
        "this_version = \"1.2.0\"\n",
        "\n",
        "from pathlib import Path\n",
        "Mode = \"forced_reinstall\" #@param [\"install\", \"update\", \"forced_reinstall\"]\n",
        "clear_output = False #@param {type:\"boolean\"}\n",
        "##@markdown Experimental: Keep tensorflow version of colab \n",
        "#keep_installed_tf = True #@param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "#@markdown Check if you dont use Quick96 / generic Xseg or want to do a pretraining\n",
        "\n",
        "skip_pretrain = True #@param {type:\"boolean\"}\n",
        "skip_pretrain_q96 = True #@param {type:\"boolean\"}\n",
        "skip_generic_xseg = True #@param {type:\"boolean\"}\n",
        "#@markdown Experimental: Deleted unneeded files and libs\n",
        "free_up_diskspace = False #@param {type:\"boolean\"}\n",
        "##@markdown Some packages (e.g. Tensorflow) needs a restart, check if this should happen automatically, \n",
        "##@markdown if you dont see your files, refresh browser\n",
        "#auto_restart = False #@param{type:\"boolean\"}\n",
        "\n",
        "install_path = \"/content/\"\n",
        "dfl_dir = \"DeepFaceLab/\"\n",
        "\n",
        "!echo \"$install_path\" > /content/install_path\n",
        "\n",
        "!echo \"$dfl_dir\" > /content/dfl_dir\n",
        "gitlink = !cat /content/gitlink\n",
        "gitlink = gitlink[0]\n",
        "\n",
        "if gitlink is None or gitlink.strip() == \"\" or \"cat\" in gitlink: \n",
        "  raise ValueError(\"Run cell above for creating gitlink \")\n",
        "\n",
        "git_commit = !cat /content/git_commit\n",
        "git_commit = git_commit[0]\n",
        "\n",
        "if git_commit is None or git_commit.strip() == \"\" or \"cat\" in git_commit: \n",
        "  git_commit = None\n",
        "\n",
        "if Mode == \"forced_reinstall\": \n",
        "  !rm -rf $dfl_dir\n",
        "\n",
        "if \"install\" in Mode:\n",
        "  !mkdir -p $install_path\n",
        "  %cd $install_path\n",
        "  #!git clone https://github.com/iperov/DeepFaceLab.git\n",
        "  \n",
        "  !git clone $gitlink $dfl_dir\n",
        "  if git_commit is not None:\n",
        "    %cd $dfl_dir\n",
        "    !git checkout $git_commit\n",
        "    %cd /content\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",
        "  \n",
        "  req_file = install_path + \"DeepFaceLab/requirements-colab.txt\"\n",
        "  #!wget -q https://raw.githubusercontent.com/deepartist/DeepFaceLab_Colab/main/requirements-colab.txt\n",
        "  #req_file = \"/content/requirements-colab.txt\"\n",
        "  !pip uninstall -y tensorflow\n",
        "  !pip install -r $req_file\n",
        "  !pip install --upgrade scikit-image\n",
        "  !pip install --upgrade pyyaml\n",
        "  #!sudo apt install cuda-10-0\n",
        "  #configure cron\n",
        "  !sudo apt install cron -y\n",
        "  !sudo service cron start\n",
        "  !sudo service cron stop\n",
        "  !cp /etc/crontab /etc/crontab.ORG\n",
        "\n",
        "  #create default workspace\n",
        "  workspace_path = install_path+\"workspace/\"\n",
        "  !echo \"$workspace_path\" > /content/workspace_path\n",
        "  if not Path(workspace_path).exists():\n",
        "    !mkdir -p $workspace_path\"\"data_src/aligned\n",
        "    !mkdir -p $workspace_path\"\"data_dst/aligned\n",
        "    !mkdir -p $workspace_path\"\"model\n",
        "\n",
        "\n",
        "  #install scripts\n",
        "  %cd /content/\n",
        "  !git clone https://github.com/deepartist/DeepFaceLab_Colab.git\n",
        "  git_notebook_version = !cat /content/DeepFaceLab_Colab/version\n",
        "  git_notebook_version = git_notebook_version[0]\n",
        "  need_update = git_notebook_version != this_version\n",
        "  !mv DeepFaceLab_Colab/scripts /content/scripts\n",
        "  !rm -rf DeepFaceLab_Colab\n",
        "  #TODO check for update in git \n",
        "\n",
        "\n",
        "\n",
        "\n",
        "elif Mode==\"update\":\n",
        "  %cd $dfl_dir\n",
        "  !git pull\n",
        "\n",
        "  install_path = !cat /content/install_path\n",
        "  install_path = install_path[0]\n",
        "  workspace_path = !cat /content/workspace_path\n",
        "  workspace_path = workspace_path[0]\n",
        "\n",
        "pretrain_install_path = install_path+\"pretrain\"\n",
        "pretrain_q_96_install_path = install_path+\"pretrain_Q96\"\n",
        "generic_xseg_install_path = install_path+\"genericXseg\"\n",
        "\n",
        "#check if pretrain folder exists\n",
        "ls = !ls $install_path\n",
        "if not \"pretrain\" in ls.nlstr and not skip_pretrain:\n",
        "  print(\"Downloading Pretrain faceset ... \")\n",
        "  !wget -q --no-check-certificate -r \"https://github.com/deepartist/DeepFaceLab_Colab/releases/download/1.0.0/pretrain_FFHQ.zip\" -O /content/pretrain_faceset.zip\n",
        "  !mkdir $pretrain_install_path\n",
        "  !unzip -q /content/pretrain_faceset.zip -d $pretrain_install_path\n",
        "  !rm /content/pretrain_faceset.zip\n",
        "else:\n",
        "  print(\"Folder pretrain already exists in %s \"%install_path)\n",
        "\n",
        "if not \"pretrain_Q96\" in ls.nlstr and not skip_pretrain_q96:\n",
        "  print(\"Downloading Q96 pretrained model ...\")\n",
        "  !wget -q --no-check-certificate -r 'https://github.com/deepartist/DeepFaceLab_Colab/releases/download/1.0.0/pretrain_Quick96.zip' -O /content/pretrain_Q96.zip\n",
        "  !mkdir $pretrain_q_96_install_path\n",
        "  !unzip -q /content/pretrain_Q96.zip -d $pretrain_q_96_install_path\n",
        "  !rm /content/pretrain_Q96.zip \n",
        "else: \n",
        "  print(\"Folder pretrain_Q96 already exists in %s \"%install_path)\n",
        "\n",
        "\n",
        "if not \"generic_xseg\" in ls.nlstr and not skip_generic_xseg:\n",
        "  #TODO\n",
        "  print(\"Downloading generic pretrained model ...\")\n",
        "  !wget -q --no-check-certificate -r 'https://github.com/deepartist/DeepFaceLab_Colab/releases/download/1.0.0/genericXseg.zip' -O /content/genericXseg.zip\n",
        "  !mkdir $generic_xseg_install_path\n",
        "  !unzip -q /content/genericXseg.zip -d $generic_xseg_install_path\n",
        "  !rm /content/generic_xseg.zip \n",
        "\n",
        "if clear_output:\n",
        "  from google.colab import output\n",
        "  output.clear()\n",
        "\n",
        "print(\"\\nDone!\")\n",
        "print(\"Install path: \")\n",
        "%cd $install_path\n",
        "\n",
        "if need_update:\n",
        "  print(\"\\nOutdated notebook version, get the notebook from https://github.com/deepartist/DeepFaceLab_Colab/\")\n",
        "\n",
        "if free_up_diskspace: #highly experimental!\n",
        "  print(\"\\nCleaning up some space!\")\n",
        "  !sudo rm -rf /usr/local/lib/python2.7\n",
        "  !sudo rm -rf /usr/local/lib/python3.8\n",
        "  !sudo rm -rf /usr/local/lib/python3.6\n",
        "  # !sudo rm -rf /usr/local/cuda-10.1\n",
        "  # !sudo rm -rf /usr/local/cuda-10.0\n",
        "  !sudo rm -rf /tensorflow-1.15.2\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e67Ap04x33Q3"
      },
      "source": [
        "#Manage workspace\n",
        "\n",
        "*   You can upload to colab (zip or folder with faceset.pak) OR\n",
        "*   use your gdrive directly (zip or folder with faceset.pak)\n",
        "*   or delete and recreate your workspace and \n",
        "*   sync with gdrive \n",
        "\n",
        "Dont use gdrive with a lot of single images, it's slow and you will exceed your quota \n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2Yicq70nwg7U"
      },
      "source": [
        "## **Up/Download**\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J_0p8GLH_r2b",
        "cellView": "form"
      },
      "source": [
        "#@title Upload zipped workspace\n",
        "%cd /content/\n",
        "import os \n",
        "from pathlib import Path\n",
        "required_folders = [\"data_dst\", \"data_src\", \"model\"]\n",
        "workspace_path = !cat /content/workspace_path\n",
        "workspace_path = workspace_path[0]\n",
        "#@markdown Delete workspace on colab drive\n",
        "delete_local_workspace = True #@param {type:\"boolean\"}\n",
        "\n",
        "from google.colab import files\n",
        "uploaded = files.upload()\n",
        "filename = uploaded[0]\n",
        "!rm -rf /content/tmp\n",
        "!mkdir -p /content/tmp\n",
        "!7z e /content/$filename -o /content/tmp\n",
        "#check if folder contains \"workspace\"\n",
        "tmp_ls = !ls /content/tmp\n",
        "tmp_ls = tmp_ls[0].split(\" \")\n",
        "\n",
        "if delete_local_workspace:\n",
        "  !rm -rf $workspace_path\n",
        "\n",
        "if len(tmp_ls) == 0 and \"workspace\" in tmp_ls[0]:\n",
        "  !rsync -ah /content/tmp $workspace_path\n",
        "  for f in required_folders:\n",
        "    !rsync -ah \"/content/tmp/workspace/\"$f $workspace_path\n",
        "elif not any(f not in \" \".join(tmp_ls) for f in required_folders):\n",
        "  for f in required_folders:\n",
        "    !rsync -ah \"/content/tmp/\"$f $workspace_path\n",
        "else:\n",
        "  raise ValueError(f\"Extracted zip is no workspace\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "48jAKXto_2bo",
        "cellView": "form"
      },
      "source": [
        "#@title Zip and Download your workspace\n",
        "import os\n",
        "!rm -rf /content/download/*\n",
        "!mkdir -p /content/download\n",
        "download_only_model = True #@param {type:\"boolean\"}\n",
        "if not download_only_model:\n",
        "  !7z a -mx0 -tzip  /content/download/workspace.zip -w $workspace_path\n",
        "else:\n",
        "  model_path = os.path.join(workspace_path, \"model/\")\n",
        "  !7z a -mx0 -tzip  /content/download/workspace.zip -w $model_path\n",
        "#!rsync --ah --info=progress2 -rf $workspace_path /content/tmp/\n",
        "from google.colab import files\n",
        "files.download(workspace_path)\n",
        "\n",
        "print(f\"Or download manually: /content/download/workspace.zip\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iWyPVL8hw3rU"
      },
      "source": [
        "##**Gdrive** "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fPHIbafi_8e4",
        "cellView": "form"
      },
      "source": [
        "#@title Sync workspace **FROM** drive \n",
        "\n",
        "#@markdown Delete workspace on colab drive\n",
        "delete_local_workspace = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown Use zipped or workspace folder with packed faces\n",
        "\n",
        "#@markdown No need to enter /content/drive/MyDrive/ \n",
        "\n",
        "#@markdown Enter folder or path containing the workspace folder OR full path zip file\n",
        "%cd /content/\n",
        "import os \n",
        "from pathlib import Path\n",
        "required_folders = [\"data_dst\", \"data_src\", \"model\"]\n",
        "gdrive_root =  \"/content/drive/MyDrive/\"\n",
        "workspace_on_gdrive = \"deepfakes/workspaces/\" #@param {type:\"string\"}\n",
        "#@markdown Only sync specific model files\n",
        "model_name = \"\" #@param {type:\"string\"}\n",
        "model_name.strip()\n",
        "#@markdown Check if you want to run the model directly via gdrive, \n",
        "\n",
        "#@markdown if yes don't use zip and deactviate save_preview_history\n",
        "model_on_gdrive = False #@param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "\n",
        "if not Path(gdrive_root).exists():\n",
        "  raise ValueError(\"Google drive not mounted\")\n",
        "\n",
        "workspace_path = !cat /content/workspace_path\n",
        "workspace_path = workspace_path[0]\n",
        "if workspace_path is None or workspace_path.strip() == \"\" or \"cat\" in workspace_path: \n",
        "  raise ValueError(\"No workspace set!\")\n",
        "\n",
        "if workspace_on_gdrive.split(\".\")[-1] == \"zip\":\n",
        "  model_on_gdrive = False\n",
        "  if not workspace_on_gdrive.startswith(\"/content/drive/MyDrive/\") or workspace_on_gdrive.startswith(\"/content/drive/My Drive/\"):\n",
        "    workspace_full = os.path.join(gdrive_root, workspace_on_gdrive)\n",
        "  else:  \n",
        "    workspace_full = workspace_on_gdrive \n",
        "  !rsync -ah --info=progress2 $workspace_full\"/\"$f /content\n",
        "  filename = Path(workspace_full).basename\n",
        "  !rm -rf /content/tmp\n",
        "  !mkdir -p /content/tmp\n",
        "  !7z e /content/$filename -o /content/tmp\n",
        "  #check if folder contains \"workspace\"\n",
        "  tmp_ls = !ls /content/tmp\n",
        "  tmp_ls = tmp_ls[0].split(\" \")\n",
        "\n",
        "  \n",
        "\n",
        "  if len(tmp_ls) == 0 and \"workspace\" in tmp_ls[0]:\n",
        "    if delete_local_workspace:\n",
        "      !rm -rf $workspace_path\n",
        "      !mkdir -p $workspace_path\n",
        "    !rsync -ah /content/tmp $workspace_path\n",
        "    for f in required_folders:\n",
        "      !rsync -ah \"/content/tmp/workspace/\"$f $workspace_path\n",
        "  elif not any(f not in \" \".join(tmp_ls) for f in required_folders):\n",
        "    if delete_local_workspace:\n",
        "      !rm -rf $workspace_path\n",
        "      !mkdir -p $workspace_path\n",
        "    for f in required_folders:\n",
        "      !rsync -ah \"/content/tmp/\"$f $workspace_path\n",
        "  else:\n",
        "    raise ValueError(f\"Extracted zip is no workspace\")\n",
        "  pass\n",
        "else:\n",
        "  if workspace_on_gdrive[:-1] == \"/\":\n",
        "    workspace_on_gdrive=workspace_on_gdrive[:-1]\n",
        "  workspace_full = os.path.join(gdrive_root, workspace_on_gdrive)\n",
        "  ls_folder = !ls $workspace_full\n",
        "  ls_folder = ls_folder[0].split(\" \")\n",
        "  if len(ls_folder) == 0: raise ValueError(f\"{workspace_full} empty\")\n",
        "  elif len(ls_folder) > 1 and not any(f not in \" \".join(ls_folder) for f in required_folders):\n",
        "    #workspace_full = Path(workspace_full).parent\n",
        "    pass\n",
        "  elif len(ls_folder) == 1 and \"workspace\" == ls_folder[0]:\n",
        "    workspace_full = os.path.join(workspace_full, \"workspace\")\n",
        "    ls_folder = !ls $workspace_full\n",
        "    ls_folder = ls_folder[0].split(\" \")\n",
        "    if len(ls_folder) > 1 and not any(f not in \" \".join(ls_folder) for f in required_folders):\n",
        "      pass\n",
        "    else:\n",
        "      raise ValueError(f\"{workspace_full} is no workspace\")\n",
        "  else:\n",
        "    raise ValueError(f\"{workspace_full} is no workspace\")\n",
        "\n",
        "\n",
        "  if delete_local_workspace:\n",
        "    !rm -rf $workspace_path\n",
        "    !mkdir -p $workspace_path\n",
        "  for f in required_folders:\n",
        "    if model_on_gdrive and \"model\" in f: pass\n",
        "    print(f\"Syncing {workspace_full}/{f}\")\n",
        "    if \"model\" in f and model_name is not \"\":\n",
        "      !rsync -ah --info=progress2 --include='$model_name*' --exclude='*' $workspace_full\"/\"$f\"/\" $workspace_path\"/model/\"\n",
        "    else:  \n",
        "      !rsync -ah --info=progress2 $workspace_full\"/\"$f\"/\" $workspace_path\"/\"$f/\"/\"\n",
        "\n",
        "!echo $workspace_full\"/model\" > /content/model_path_gdrive\n",
        "!echo $model_on_gdrive > /content/model_on_gdrive\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BIRf1RFpUupd",
        "cellView": "form"
      },
      "source": [
        "#@title Sync workspace **TO** drive \n",
        "#@markdown Use zipped or workspace folder with packed faces\n",
        "\n",
        "#@markdown No need to enter /content/drive/MyDrive/ \n",
        "\n",
        "#@markdown Enter folder or path containing the workspace folder OR full path zip file\n",
        "\n",
        "import os \n",
        "gdrive_root =  \"/content/drive/MyDrive/\"\n",
        "workspace_on_gdrive = \"deepfakes/workspaces/\" #@param {type:\"string\"}\n",
        "workspace_full = os.path.join(gdrive_root, workspace_on_gdrive)\n",
        "from pathlib import Path\n",
        "if not Path(gdrive_root).exists():\n",
        "  raise ValueError(\"Google drive not mounted\")\n",
        "\n",
        "workspace_path = !cat /content/workspace_path\n",
        "workspace_path = workspace_path[0]\n",
        "if workspace_path is None or workspace_path.strip() == \"\" or \"cat\" in workspace_path: \n",
        "  raise ValueError(\"No workspace set!\")\n",
        "\n",
        "                              \n",
        "if zip in workspace_on_gdrive.split(\".\")[-1]:\n",
        "  !7z a -tzip $workspace_on_gdrive -w workspace_path\n",
        "else: \n",
        "  !rsync -ah --info=progress2  $workspace_path $workspace_full\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L2ObADmC7A03"
      },
      "source": [
        "#**Training**"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Modify config\n",
        "#@markdown Change model, set pathes, run this  cell, set config values, then hit build config \n",
        "install_path = !cat /content/install_path\n",
        "install_path = install_path[0]\n",
        "\n",
        "dfl_dir = !cat /content/dfl_dir\n",
        "dfl_dir = dfl_dir[0]\n",
        "\n",
        "gitlink = !cat /content/gitlink\n",
        "gitlink = gitlink[0]\n",
        "\n",
        "\n",
        "import ipywidgets as widgets\n",
        "from ipywidgets import interact_manual\n",
        "import yaml\n",
        "import os\n",
        "from pathlib import Path\n",
        "widgets_dict = {}\n",
        "\n",
        "if \"MachineEditor\" not in gitlink:\n",
        "  raise ValueError('Config only works with MVE branch')\n",
        "\n",
        "Model = \"AMP\" #@param [\"SAEHD\", \"AMP\", \"Quick96\", \"XSeg\"]\n",
        "#@markdown if empty, formated config from model folder is used \n",
        "load_model_config_path = \"\" #@param{\"type\":\"string\"}\n",
        "\n",
        "save_config_path = \"/content/config.yaml\" #@param{\"type\":\"string\"}\n",
        "\n",
        "config_file = install_path + dfl_dir + \"models/Model_\" +Model+\"/config_schema.json\"\n",
        "if load_model_config_path is not None and load_model_config_path.strip() != '' and Path(load_model_config_path).exists():\n",
        "  formatted_file = load_model_config_path\n",
        "else:\n",
        "  print(f\"Configfile at Path '{load_model_config_path}' not found\")\n",
        "  formatted_file = install_path + dfl_dir + \"models/Model_\" +Model+\"/formatted_config.yaml\"\n",
        "\n",
        "with open(formatted_file, 'r') as file:\n",
        "  nested_defs = yaml.safe_load(file)\n",
        "\n",
        "defs = {}\n",
        "\n",
        "def __iterate_read_dict(nested_dict, new_dict=None, keep_dicts=False):\n",
        "  if new_dict == None:\n",
        "    new_dict = {}\n",
        "  for k, v in nested_dict.items():\n",
        "    if isinstance(v, dict):\n",
        "      if keep_dicts:\n",
        "        new_dict[k] = dict([])\n",
        "      new_dict.update(__iterate_read_dict(v, new_dict))\n",
        "    else:            \n",
        "      new_dict[k] = v\n",
        "  return new_dict\n",
        "\n",
        "defs = __iterate_read_dict(nested_defs, defs, keep_dicts=True)\n",
        "\n",
        "defs_c = defs.copy()\n",
        "\n",
        "\n",
        "with open(config_file, 'r') as file:\n",
        "  data = yaml.safe_load(file)\n",
        "\n",
        "props = data[\"definitions\"][\"dfl_config\"][\"properties\"]\n",
        "\n",
        "def build_widgets(key, defs, props, v=None, v_type=None):\n",
        "  if v == None:\n",
        "    v = props[key]\n",
        "  if v_type == None:\n",
        "    v_type =  v.get(\"type\", None)\n",
        "  #print(v_type, key)\n",
        "  if v_type == \"boolean\":\n",
        "\n",
        "    val = defs.get(key, None)\n",
        "    w = widgets.Checkbox(value=val, description=key, style= {'description_width': 'initial'})\n",
        "    \n",
        "  elif v_type == \"integer\":\n",
        "    min = v.get(\"minimum\", 0)\n",
        "    max = v.get(\"maximum\", 2**31-1)\n",
        "\n",
        "    step = v.get(\"multipleOf\", 1)\n",
        "    val = defs.get(key, min+step)\n",
        "    w = widgets.BoundedIntText(value=val, min=min, max=max, step=step, description=key, style= {'description_width': 'initial'})\n",
        "\n",
        "  elif v_type == \"number\":\n",
        "    min = v.get(\"minimum\", 0)\n",
        "    max = v.get(\"maximum\", 2**31-1)\n",
        "    if key == \"lr\":\n",
        "      step = 0.00000001\n",
        "    elif \"gan\" in key:\n",
        "      step = 0.001\n",
        "    else: \n",
        "      step = 0.01 \n",
        "\n",
        "    val = defs.get(key, min+step)\n",
        "    w = widgets.BoundedFloatText(value=val, min=min, max=max, step=step, description=key, style= {'description_width': 'initial'})\n",
        "\n",
        "  elif v_type == \"string\":\n",
        "    opts = v.get(\"enum\", None)\n",
        "    \n",
        "    if opts is not None:\n",
        "      val = defs.get(key, opts[0])\n",
        "      w = widgets.Combobox(options=opts, value=val, placeholder=val, description=key, ensure_option=True, style= {'description_width': 'initial'})\n",
        "    else:\n",
        "      val = defs.get(key, \"\")\n",
        "      placeh = \"\"\n",
        "      if key == \"archi\":\n",
        "        placeh = \"'liae' or 'df' with options 'udt'\"\n",
        "      w = widgets.Text(value=val, placeholder=placeh, description=key, style= {'description_width': 'initial'})\n",
        "  return w\n",
        "\n",
        "\n",
        "for k, v in defs.items():\n",
        "  if k in props:\n",
        "    cur_widget = build_widgets(k, defs, props)\n",
        "  elif isinstance(v, dict):\n",
        "    #print(\"Pair \" + k, v)\n",
        "    cur_widget = widgets.HTML('<b>'+k+':<b>')\n",
        "  else:\n",
        "    if isinstance(v,float):\n",
        "      cur_widget = build_widgets(k, defs, props, v={}, v_type=\"number\")\n",
        "    elif isinstance(v, bool):\n",
        "      cur_widget = build_widgets(k, defs, props, v={}, v_type=\"boolean\")\n",
        "    elif isinstance(v, int):\n",
        "      cur_widget = build_widgets(k, defs, props, v={}, v_type=\"integer\")\n",
        "    elif isinstance(v, str):\n",
        "      cur_widget = build_widgets(k, defs, props, v={}, v_type=\"string\")\n",
        "\n",
        "  widgets_dict[k] = cur_widget\n",
        "\n",
        "#build ui\n",
        "ui = None\n",
        "h_arr = []\n",
        "v_arr = []\n",
        "for i, kv in enumerate(widgets_dict.items()):\n",
        "  k = kv[0]\n",
        "  v = kv[1]\n",
        "  v_arr.append(widgets_dict[k])\n",
        "  if (i+1) % ((len(widgets_dict.keys())//3)+1) == 0: \n",
        "    h_arr.append(widgets.VBox(v_arr))\n",
        "    v_arr = []\n",
        "h_arr.append(widgets.VBox(v_arr))\n",
        "\n",
        "display(widgets.HTML(value=\"<h2>Change values here and then click button 'run interact' </h2>\"))\n",
        "display(widgets.HBox(h_arr))\n",
        "\n",
        "def __update_nested_dict(nested_dict, key, val):\n",
        "    if key in nested_dict:\n",
        "        nested_dict[key] = val\n",
        "        return True\n",
        "    for k, v in nested_dict.items():\n",
        "          if isinstance(v, dict):\n",
        "              if __update_nested_dict(v, key, val):\n",
        "                  return True\n",
        "    return False\n",
        "\n",
        "def build_config(conf_path):\n",
        "  #copy formatted file \n",
        "  !cp $formatted_file $conf_path\n",
        "  with open(conf_path, 'r') as file:\n",
        "    nested = yaml.safe_load(file)\n",
        "    #print(nested)\n",
        "\n",
        "  for k, opt in widgets_dict.items():\n",
        "    if isinstance(opt, widgets.HTML): continue\n",
        "    val = opt.value\n",
        "    key = k\n",
        "\n",
        "    if not __update_nested_dict(nested, key, val):\n",
        "      nested[key] = val\n",
        "\n",
        "  with open(conf_path, 'w') as file:\n",
        "    yaml.dump(nested, file, sort_keys=False)\n",
        "  !echo \"$conf_path\" > /content/conf_path\n",
        "  print(\"Saved!\")\n",
        " \n",
        "display(widgets.HTML(value=\"<h3>Set save path and build config</h3>\"))\n",
        "\n",
        "im = interact_manual.options(manual_name=\"Build Config!\")\n",
        "im(build_config,conf_path=widgets.Text(value=save_config_path, description=\"save path\"))\n",
        "\n"
      ],
      "metadata": {
        "id": "ubZpRnkXXRny",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xNMFdQtp5Lza",
        "cellView": "form"
      },
      "source": [
        "%load_ext tensorboard\n",
        "install_path = !cat /content/install_path\n",
        "install_path = install_path[0]\n",
        "\n",
        "model_path_gdrive = !cat /content/model_path_gdrive\n",
        "model_path_gdrive = model_path_gdrive[0]\n",
        "\n",
        "\n",
        "if model_path_gdrive is None or model_path_gdrive.strip() == \"\" or \"cat\" in model_path_gdrive: \n",
        "  model_path_gdrive = False\n",
        "\n",
        "\n",
        "conf_path = !cat /content/conf_path\n",
        "conf_path = conf_path[0]\n",
        "\n",
        "if conf_path is None or conf_path.strip() == \"\" or \"cat\" in conf_path: \n",
        "  conf_path = None\n",
        "  \n",
        "workspace_path = !cat /content/workspace_path\n",
        "workspace_path = workspace_path[0]\n",
        "if workspace_path is None or workspace_path.strip() == \"\" or \"cat\" in workspace_path: \n",
        "  raise ValueError(\"No workspace set!\")\n",
        "\n",
        "\n",
        "model_on_gdrive = !cat /content/model_on_gdrive\n",
        "model_on_gdrive = model_on_gdrive[0]\n",
        "if model_on_gdrive is None or model_on_gdrive.strip() == \"\" or \"cat\" in model_on_gdrive: \n",
        "  model_on_gdrive = False\n",
        "else:\n",
        "  model_on_gdrive=(\"True\" == model_on_gdrive)\n",
        "\n",
        "\n",
        "\n",
        "#@title Training\n",
        "Model = \"AMP\" #@param [\"SAEHD\", \"SAEHD\", \"AMP\", \"Quick96\", \"XSeg\"]\n",
        "SRC_SRC_training = False #@param {type:\"boolean\"}\n",
        "tensorboard = False #@param {type:\"boolean\"}\n",
        "Silent_Start = False #@param {type:\"boolean\"}\n",
        "\n",
        "\n",
        "#@markdown Config options:\n",
        "autogen_config = False #@param {type:\"boolean\"}\n",
        "overwrite_config = False #@param {type:\"boolean\"}\n",
        "overwrite_config_path = \"/content/config.yaml\" #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "#@markdown Backup options\n",
        "Backup_hourly_on_gdrive = True #@param {type:\"boolean\"}\n",
        "Colab_Pro = True #@param {type:\"boolean\"}\n",
        "#@markdown if colab pro runtime is set to 23.5 hours\n",
        "\n",
        "\n",
        "if Colab_Pro:\n",
        "  max_time = int(23.5 * 60 * 60)\n",
        "else: \n",
        "  max_time = int(11.5 * 60 * 60)\n",
        "\n",
        "\n",
        "%cd $install_path\n",
        "\n",
        "import psutil, os, time\n",
        "from pathlib import Path\n",
        "\n",
        "gdrive_root =  \"/content/drive/MyDrive/\"\n",
        "if not Path(gdrive_root).exists():\n",
        "  raise ValueError(\"Google drive not mounted\")\n",
        "\n",
        "model_path = os.path.join(workspace_path, \"model\")\n",
        "!mkdir -p $model_path\n",
        "#@markdown Overwrite backup location,\n",
        "#@markdown no need to enter /content/drive/MyDrive/ \n",
        "\n",
        "overwrite_backup_location = False #@param {type:\"boolean\"}\n",
        "workspace_on_gdrive = \"deepfakes/workspaces/\" #@param {type:\"string\"}\n",
        "zip_before_backup = False #@param {type:\"boolean\"}\n",
        "if not workspace_on_gdrive.startswith(\"/content/drive/MyDrive/\") or workspace_on_gdrive.startswith(\"/content/drive/My Drive/\"):\n",
        "  workspace_full = os.path.join(gdrive_root, workspace_on_gdrive)\n",
        "else:  \n",
        "  workspace_full = workspace_on_gdrive \n",
        "\n",
        "\n",
        "\n",
        "p = psutil.Process(os.getpid())\n",
        "uptime = time.time() - p.create_time()\n",
        "\n",
        "\n",
        "#create cron rsyncjob \n",
        "!sudo service cron stop\n",
        "!rm -rf /etc/crontab\n",
        "!cp /etc/crontab.ORG /etc/crontab\n",
        "\n",
        "data_dst_opt = \"data_dst\" if not SRC_SRC_training else \"data_src\"\n",
        "\n",
        "if not model_on_gdrive:\n",
        "\n",
        "  if zip_before_backup:\n",
        "    !echo 'rm -rf /content/tmp; mkdir -p /content/tmp' > /content/cronjob.sh\n",
        "    !echo '7z a -mx0 -tzip /content/tmp/model.zip $model_path' >> /content/cronjob.sh\n",
        "    if overwrite_backup_location or not model_path_gdrive:\n",
        "      !echo 'bash $install_path\"scripts/0_rsync.sh\" $model_path/ \"$workspace_full/\" *' >> /content/cronjob.sh\n",
        "    else:\n",
        "      !echo 'bash $install_path\"scripts/0_rsync.sh\" $model_path/ \"$model_path_gdrive/\" $Model' >> /content/cronjob.sh\n",
        "  elif overwrite_backup_location or not model_path_gdrive: \n",
        "    !echo 'bash $install_path\"scripts/0_rsync.sh\" $model_path/ \"$workspace_full/\" $Model' > /content/cronjob.sh\n",
        "  else:\n",
        "    !echo 'bash $install_path\"scripts/0_rsync.sh\" $model_path/ \"$model_path_gdrive/\" $Model' > /content/cronjob.sh\n",
        "\n",
        "  cron_min = \"*\"\n",
        "  cron_hour =\"*\"\n",
        "  if (Backup_hourly_on_gdrive):\n",
        "    print(\"Time to backup: \"+str(round((max_time-uptime)/3600))+\" hours\")\n",
        "    cron_min = !date +\"%M\"\n",
        "    cron_min = (int(cron_min[0])+5) %59\n",
        "  elif (round(max_time-3300-uptime) > 0):\n",
        "    print(\"Time to backup: \"+str(round((max_time-uptime)/3600))+\" hours\")\n",
        "    hours_left = int(round(max_time-uptime)/3600)\n",
        "    cron_hour = !(date +\"%H\")\n",
        "    cron_hour = (int(cron_hour[0])+hours_left)%23\n",
        "  else:\n",
        "    raise Exception(\"Session expires in less than an hour. Do a factory reset!\")\n",
        "    backup_cmd = \"\"\n",
        "\n",
        "  \n",
        "  #save cronjob\n",
        "  !echo '$cron_min $cron_hour * * * root bash /content/cronjob.sh >> /content/cron_backup_log.txt' >> /etc/crontab\n",
        "  !sudo service cron start\n",
        "  crontab = !cat /etc/crontab\n",
        "  print(f\"Cronjob entry: {crontab[-1]}\")\n",
        "\n",
        "  \n",
        "  cmd = f\"DeepFaceLab/main.py train --training-data-src-dir {workspace_path}data_src/aligned --training-data-dst-dir {workspace_path}{data_dst_opt}/aligned --pretraining-data-dir pretrain --model-dir {workspace_path}model --model \"+Model\n",
        "\n",
        "else:\n",
        "  print(\"Model on gdrive, not extra backup needed. Make Sure to deactivate preview history or risk getting drive quota error!\")\n",
        "  cmd = f\"DeepFaceLab/main.py train --training-data-src-dir {workspace_path}data_src/aligned --training-data-dst-dir {workspace_path}{data_dst_opt}/aligned --pretraining-data-dir pretrain --model-dir {model_path_gdrive} --model \"+Model\n",
        "\n",
        "if tensorboard:\n",
        "  %tensorboard --logdir $workspace_path\"model/tb_logs\"\n",
        "  from tensorboard import notebook\n",
        "  notebook.display(height=600) \n",
        "  cmd+= f' --tensorboard-logdir \"{workspace_path}model/tb_logs\"'\n",
        "\n",
        "if autogen_config:\n",
        "  cmd+= f' --auto-gen-config'\n",
        "\n",
        "if overwrite_config and overwrite_config_path is not None and Path(overwrite_config_path).exists():\n",
        "    cmd+= f' --config-training-file {overwrite_config_path}'\n",
        "\n",
        "\n",
        "if Model == \"Quick96\":\n",
        "  cmd+= \" --pretrained-model-dir pretrain_Q96\"\n",
        "\n",
        "if Silent_Start:\n",
        "  cmd+= \" --silent-start\"\n",
        "\n",
        "print(f\"Running comamnd {cmd}\")\n",
        "\n",
        "train_cmd = cmd\n",
        "\n",
        "!python $train_cmd\n",
        "!sudo bash /content/cronjob.sh\n",
        "!sudo service cron stop"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E-054PYI7DpX"
      },
      "source": [
        "#Scripts\n",
        "\n",
        "Run these scripts interactivly!\n",
        "GUI related stuff may do not work (like xseg or interactive merger and training with preview)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_hYk1xptr80l",
        "cellView": "form"
      },
      "source": [
        "#@title Run cell and than select script, hit \"Run script\" \n",
        "#@markdown Do not use for training or gui related stuff !\n",
        "\n",
        "#@markdown use arrow keys \n",
        "import ipywidgets as widgets\n",
        "from ipywidgets import interact_manual\n",
        "from os.path import join\n",
        "from glob import glob\n",
        "import os \n",
        "\n",
        "install_path = !cat /content/install_path\n",
        "install_path = install_path[0]\n",
        "script_dir = install_path+\"/scripts/*.sh\"\n",
        "\n",
        "#audio_in = \"eminem_drugs_cut.mp3\"\n",
        "\n",
        "scripts_arr = glob(script_dir)   \n",
        "scripts_arr = sorted(list(map(os.path.basename, scripts_arr)))\n",
        "if len(scripts_arr) == 0:\n",
        "  raise Exception(\"No scripts found, did you clone the repo?\")\n",
        "  pass\n",
        "script_widget = widgets.Select(options=scripts_arr, description=\"Select script\", layout={'height':\"350px\", 'width': 'max-content'})\n",
        "\n",
        "display(script_widget)\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "azlh67nhspVi",
        "cellView": "form"
      },
      "source": [
        "!#@title run script\n",
        "\n",
        "#@markdown Click Runtime->Inerrupt Execution (somtime multiple times) to stop\n",
        "import os \n",
        "%cd /content\n",
        "install_path = !cat /content/install_path\n",
        "install_path = install_path[0]\n",
        "script_dir = os.path.join(install_path, \"scripts\")\n",
        "%cd $script_dir\n",
        "script = script_widget.value\n",
        "print(f\"Running script {script}\")\n",
        "!sudo bash $script\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u_ev-AgJu9NT"
      },
      "source": [
        "# Tools\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iUAIErHkfFda"
      },
      "source": [
        "##Shell"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "XD4YUlTavCzj"
      },
      "source": [
        "#@title build shell\n",
        "from IPython.display import JSON\n",
        "from google.colab import output\n",
        "from subprocess import getoutput\n",
        "import os\n",
        "\n",
        "def shell(command):\n",
        "  if command.startswith('cd'):\n",
        "    path = command.strip().split(maxsplit=1)[1]\n",
        "    os.chdir(path)\n",
        "    return JSON([''])\n",
        "  return JSON([getoutput(command)])\n",
        "output.register_callback('shell', shell)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "OmFHbH1vvGMu"
      },
      "source": [
        "#@title Run simple Shell\n",
        "%%html\n",
        "<div id=term_demo></div>\n",
        "<script src=\"https://code.jquery.com/jquery-latest.js\"></script>\n",
        "<script src=\"https://cdn.jsdelivr.net/npm/jquery.terminal/js/jquery.terminal.min.js\"></script>\n",
        "<link href=\"https://cdn.jsdelivr.net/npm/jquery.terminal/css/jquery.terminal.min.css\" rel=\"stylesheet\"/>\n",
        "<script>\n",
        "  $('#term_demo').terminal(async function(command) {\n",
        "      if (command !== '') {\n",
        "          try {\n",
        "              let res = await google.colab.kernel.invokeFunction('shell', [command])\n",
        "              let out = res.data['application/json'][0]\n",
        "              this.echo(new String(out))\n",
        "          } catch(e) {\n",
        "              this.error(new String(e));\n",
        "          }\n",
        "      } else {\n",
        "          this.echo('');\n",
        "      }\n",
        "  }, {\n",
        "      greetings: 'Welcome to Colab Shell',\n",
        "      name: 'colab_demo',\n",
        "      height: 250,\n",
        "      prompt: 'colab > '\n",
        "  });"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}