{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G7E7eZwJ7F5c"
      },
      "outputs": [],
      "source": [
        "#@title *Prepare* { display-mode: \"form\", run: \"auto\" }\n",
        "\n",
        "#@markdown # **Prepare the environment**\n",
        "#@markdown Execute to install necessary packages, modules, fonts\n",
        "\n",
        "#@markdown Mount Google Drive at /content/drive (your drive folder at /content/drive/MyDrive):\n",
        "mount_gdrive = True #@param {type:\"boolean\"}\n",
        "#@markdown Force mount again. Useful for bug cases:\n",
        "force_remount = False #@param {type:\"boolean\"}\n",
        "\n",
        "# remove the old log\n",
        "!rm -rf error_log.txt > /dev/null\n",
        "\n",
        "import os\n",
        "import subprocess\n",
        "from google.colab import drive\n",
        "\n",
        "# mount Google Drive\n",
        "if mount_gdrive:\n",
        "    print(\"Mounting your Google Drive | Waiting user Allow Access | \", end='')\n",
        "    try:\n",
        "        drive.mount('/content/drive/', force_remount=force_remount)\n",
        "    except Exception as e:\n",
        "        print(f\"[✗]: {e}\")\n",
        "\n",
        "# update code\n",
        "print(\"Ensuring the LeGen code is existing and updated...\", end='')\n",
        "repo_url = \"https://github.com/matheusbach/legen.git\"\n",
        "local_folder = \"/content/src\"  # LeGen source path\n",
        "\n",
        "# Create directory if it does not exist\n",
        "os.makedirs(local_folder, exist_ok=True)\n",
        "\n",
        "# Try git status in the directory\n",
        "git_task = \"git fetch\"\n",
        "process = subprocess.Popen(git_task, cwd=local_folder, shell=True)\n",
        "return_code = process.wait()\n",
        "if return_code == 0:\n",
        "  git_task = \"git fetch && git reset --hard origin/main && git pull\"\n",
        "else:\n",
        "  !rm -rf local_folder > /dev/null\n",
        "  os.makedirs(local_folder, exist_ok=True)\n",
        "  git_task = f\"git clone {repo_url} {local_folder}\"\n",
        "\n",
        "# If it is a git repo, fetch, reset, and pull. Else, clone.\n",
        "with open('/content/error_log.txt', 'a') as f:\n",
        "    process = subprocess.Popen(git_task, cwd=local_folder, shell=True, stderr=f)\n",
        "    return_code = process.wait()\n",
        "    print(\"[✔]\" if return_code == 0 else \"[✗]\")\n",
        "\n",
        "# install pip requirements.txt updating\n",
        "print(\"Installing or updating pip requirements...\", end='')\n",
        "with open('/content/error_log.txt', 'a') as f:\n",
        "    process = subprocess.Popen('pip3 install --upgrade $(grep -v \"git+https://github.com/matheusbach/whisperx.git\" requirements.txt | grep -v \"torch\") && pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 && pip3 install git+https://github.com/matheusbach/whisperx.git --upgrade', cwd=local_folder, shell=True, stderr=f)\n",
        "    return_code = process.wait()\n",
        "    print(\"[✔]\" if return_code == 0 else \"[✗]\")\n",
        "\n",
        "# install ffmpeg and xvfb\n",
        "print(\"Installing FFmpeg and xvfb...\", end='')\n",
        "with open('/content/error_log.txt', 'a') as f:\n",
        "    process = subprocess.Popen('apt update -y ; apt install ffmpeg xvfb -y', shell=True, stderr=f)\n",
        "    return_code = process.wait()\n",
        "    print(\"[✔]\" if return_code == 0 else \"[✗]\")\n",
        "\n",
        "# install pip requirements.txt updating\n",
        "print(\"Installing fonts...\", end='')\n",
        "with open('/content/error_log.txt', 'a') as f:\n",
        "    process = subprocess.Popen('echo ttf-mscorefonts-installer msttcorefonts/accepted-mscorefonts-eula select true | debconf-set-selections && apt install -y ttf-mscorefonts-installer && fc-cache -f -v', shell=True, stderr=f)\n",
        "    return_code = process.wait()\n",
        "    print(\"[✔]\" if return_code == 0 else \"[✗]\")\n",
        "\n",
        "# create a virtual display\n",
        "os.system('Xvfb :1 -screen 0 2560x1440x8  &') # create virtual display with size 1600x1200 and 8 bit color. Color can be changed to 24, 16 or 8\n",
        "os.environ['DISPLAY'] = ':1.0' # tell X clients to use our virtual DISPLAY :1.0.\n",
        "\n",
        "print(\"\\nPreparation tasks done.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Bz77fHta9JxG"
      },
      "outputs": [],
      "source": [
        "#@title *Configure* { display-mode: \"form\", run: \"auto\" }\n",
        "#@markdown # **Define Software Settings**\n",
        "#@markdown ---\n",
        "#@markdown ## General Options\n",
        "#@markdown Set where your files are located (your Drive is the base /content/drive/MyDrive):\n",
        "import os\n",
        "from IPython.display import display\n",
        "import ipywidgets as widgets\n",
        "input_path = \"/content/drive/MyDrive/LeGen/media\" #@param {type:\"string\"}\n",
        "#@markdown Set where output soft subs (.srt and embed) videos will be saved on your Drive:\n",
        "output_softsubs_path = \"/content/drive/MyDrive/LeGen/softsubs\" #@param {type:\"string\"}\n",
        "#@markdown Set where output hard subs (burned in) videos will be saved on your Drive:\n",
        "output_hardsubs_path = \"/content/drive/MyDrive/LeGen/hardsubs\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown ---\n",
        "overwrite = False #@param {type:\"boolean\"}\n",
        "normalize = False #@param {type:\"boolean\"}\n",
        "copy_extra_files = True #@param {type:\"boolean\"}\n",
        "\n",
        "# generate_srt_files = True #@param {type:\"boolean\"}\n",
        "generate_softsub = True #@param {type:\"boolean\"}\n",
        "generate_hardsub = True  # @param {type:\"boolean\"}\n",
        "\n",
        "# buttons for directory create\n",
        "button = widgets.Button(description=\"Create directories!\")\n",
        "output = widgets.Output()\n",
        "\n",
        "\n",
        "def on_button_clicked(b):\n",
        "  # Create input and output directories if it does not exist\n",
        "  os.makedirs(input_path, exist_ok=True)\n",
        "  if generate_softsub:\n",
        "    os.makedirs(output_softsubs_path, exist_ok=True)\n",
        "  if generate_hardsub:\n",
        "    os.makedirs(output_hardsubs_path, exist_ok=True)\n",
        "  with output:\n",
        "    print(\"Done.\")\n",
        "\n",
        "\n",
        "button.on_click(on_button_clicked)\n",
        "display(button, output)\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown ## Transcription Settings:\n",
        "transcription_engine = 'WhisperX' # @param [\"Whisper\", \"WhisperX\"]\n",
        "transcription_device = 'auto' #@param [\"auto\", \"cpu\", \"cuda\"]\n",
        "transcription_model = 'large' #@param [\"tiny\", \"small\", \"medium\", \"large\", \"large-v1\", \"large-v2\", \"large-v3\"]\n",
        "compute_type = 'default' # @param [\"default\", \"int8\", \"int16\", \"float16\", \"float32\"]\n",
        "batch_size = 12 # @param {type: \"number\"}\n",
        "transcription_input_lang = 'auto detect' #@param [\"auto detect\", \"af\", \"am\", \"ar\", \"as\", \"az\", \"ba\", \"be\", \"bg\", \"bn\", \"bo\", \"br\", \"bs\", \"ca\", \"cs\", \"cy\", \"da\", \"de\", \"el\", \"en\", \"es\", \"et\", \"eu\", \"fa\", \"fi\", \"fo\", \"fr\", \"gl\", \"gu\", \"ha\", \"haw\", \"he\", \"hi\", \"hr\", \"ht\", \"hu\", \"hy\", \"id\", \"is\", \"it\", \"ja\", \"jw\", \"ka\", \"kk\", \"km\", \"kn\", \"ko\", \"la\", \"lb\", \"ln\", \"lo\", \"lt\", \"lv\", \"mg\", \"mi\", \"mk\", \"ml\", \"mn\", \"mr\", \"ms\", \"mt\", \"my\", \"ne\", \"nl\", \"nn\", \"no\", \"oc\", \"pa\", \"pl\", \"ps\", \"pt\", \"ro\", \"ru\", \"sa\", \"sd\", \"si\", \"sk\", \"sl\", \"sn\", \"so\", \"sq\", \"sr\", \"su\", \"sv\", \"sw\", \"ta\", \"te\", \"tg\", \"th\", \"tk\", \"tl\", \"tr\", \"tt\", \"uk\", \"ur\", \"uz\", \"vi\", \"yi\", \"yo\", \"zh\"]\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown ## Translation Settings:\n",
        "#@markdown Set the destination langrage code. Set to same as original to skip translation\n",
        "target_language_code = 'pt' #@param [\"af\", \"sq\", \"am\", \"ar\", \"hy\", \"as\", \"ay\", \"az\", \"bm\", \"eu\", \"be\", \"bn\", \"bho\", \"bs\", \"bg\", \"ca\", \"ceb\", \"zh-CN\", \"zh-TW\", \"co\", \"hr\", \"cs\", \"da\", \"dv\", \"doi\", \"nl\", \"en\", \"eo\", \"et\", \"ee\", \"fil\", \"fi\", \"fr\", \"fy\", \"gl\", \"ka\", \"de\", \"el\", \"gn\", \"gu\", \"ht\", \"ha\", \"haw\", \"he\", \"hi\", \"hmn\", \"hu\", \"is\", \"ig\", \"ilo\", \"id\", \"ga\", \"it\", \"ja\", \"jv\", \"kn\", \"kk\", \"km\", \"rw\", \"gom\", \"ko\", \"kri\", \"ku\", \"ckb\", \"ky\", \"lo\", \"la\", \"lv\", \"ln\", \"lt\", \"lg\", \"lb\", \"mk\", \"mai\", \"mg\", \"ms\", \"ml\", \"mt\", \"mi\", \"mr\", \"mni-Mtei\", \"lus\", \"mn\", \"my\", \"ne\", \"no\", \"ny\", \"or\", \"om\", \"ps\", \"fa\", \"pl\", \"pt\", \"pa\", \"qu\", \"ro\", \"ru\", \"sm\", \"sa\", \"gd\", \"nso\", \"sr\", \"st\", \"sn\", \"sd\", \"si\", \"sk\", \"sl\", \"so\", \"es\", \"su\", \"sw\", \"sv\", \"tl\", \"tg\", \"ta\", \"tt\", \"te\", \"th\", \"ti\", \"ts\", \"tr\", \"tk\", \"ak\", \"uk\", \"ur\", \"ug\", \"uz\", \"vi\", \"cy\", \"xh\", \"yi\", \"yo\", \"zu\"]\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown ## Video Settings:\n",
        "codec_video = \"h264\"  #@param [\"h264\", \"hevc\", \"mpeg4\"]\n",
        "video_hardware_api = \"auto\"  #@param [\"auto\", \"none\", \"nvenc\", \"vaapi\", \"amf\", \"v4l2m2m\", \"qsv\", \"videotoolbox\", \"cuvid\"]\n",
        "\n",
        "codec_audio = \"aac\"  #@param [\"aac\", \"libopus\", \"libmp3lame\", \"pcm_s16le\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n4KSJ9BL7JOs"
      },
      "outputs": [],
      "source": [
        "#@title *Run* { display-mode: \"form\" }\n",
        "#@markdown # **Run LeGen.py**\n",
        "\n",
        "print(\"Starting LeGen...\")\n",
        "import torch\n",
        "try:\n",
        "    import tensorflow  # required in Colab to avoid protobuf compatibility issues\n",
        "except ImportError:\n",
        "    pass\n",
        "\n",
        "# Create input and output directories if it does not exist\n",
        "os.makedirs(input_path, exist_ok=True)\n",
        "if generate_softsub:\n",
        "  os.makedirs(output_softsubs_path, exist_ok=True)\n",
        "if generate_hardsub:\n",
        "  os.makedirs(output_hardsubs_path, exist_ok=True)\n",
        "\n",
        "#build query\n",
        "query = f\" -i '{input_path}'\"\n",
        "query += f\" --output_softsubs '{output_softsubs_path}'\"\n",
        "query += f\" --output_hardsubs '{output_hardsubs_path}'\"\n",
        "query += \" --overwrite\" if overwrite else \"\"\n",
        "query += \" --norm\" if normalize else \"\"\n",
        "query += \" --copy_files\" if not copy_extra_files else \"\"\n",
        "query += \" --disable_softsubs\" if not generate_softsub else \"\"\n",
        "query += \" --disable_hardsubs\" if not generate_hardsub else \"\"\n",
        "query += f\" -ts:e {transcription_engine.lower()}\"\n",
        "query += f\" -ts:d {transcription_device.lower()}\"\n",
        "query += f\" -ts:m {transcription_model}\"\n",
        "query += f\" -ts:c {compute_type}\"\n",
        "query += f\" -ts:b {batch_size}\"\n",
        "query += f\" --input_lang {transcription_input_lang}\" if transcription_input_lang != \"auto detect\" else \"\"\n",
        "query += f\" --translate {target_language_code.lower()}\"\n",
        "query += f\" -c:v {codec_video}\" + (\"\" if video_hardware_api == \"none\" else f\"_{video_hardware_api}\" if video_hardware_api != \"auto\" else \"_nvenc\" if torch.cuda.is_available() else \"\")\n",
        "query += f\" -c:a {codec_audio}\"\n",
        "\n",
        "#run python script\n",
        "print(f\"command line: python3 /content/src/legen.py {query}\", end=\"\\n\\n\")\n",
        "!python3 /content/src/legen.py $query"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}