auto_translate.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "gpuType": "A100",
      "mount_file_id": "1ztQQGtwOW9BOFm-znjyqJGIb75Fi4BcQ",
      "authorship_tag": "ABX9TyM815xEAxDeJTYCAvgXL/91",
      "include_colab_link": true
    },
    "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/elishaROBINSON/toys/blob/master/auto_translate.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%pip install pytube pydub numpy wavio tqdm wave\n",
        "%pip install -q git+https://github.com/facebookresearch/seamless_communication\n",
        "\n",
        "# Import necessary libraries\n",
        "import argparse\n",
        "import logging\n",
        "import torch\n",
        "import torchaudio\n",
        "import wave\n",
        "from seamless_communication.models.inference import Translator\n",
        "from fairseq2.memory import MemoryBlock\n",
        "from pydub import AudioSegment\n",
        "import io\n",
        "from tqdm import tqdm\n",
        "from fairseq2.memory import MemoryBlock\n",
        "import torchaudio\n",
        "import torchaudio.transforms as T\n",
        "from torchaudio.io import StreamReader\n",
        "from pytube import YouTube\n",
        "import os\n",
        "\n",
        "# Configure logging\n",
        "logging.basicConfig(\n",
        "    level=logging.INFO,\n",
        "    format=\"%(asctime)s %(levelname)s -- %(name)s: %(message)s\",\n",
        ")\n",
        "\n",
        "# Define the YouTube URL to download\n",
        "url = \"\"\n",
        "\n",
        "# Create a logger\n",
        "logger = logging.getLogger(__name__)\n",
        "\n",
        "# Function to download a video from YouTube\n",
        "def Download(link):\n",
        "    youtubeObject = YouTube(link)\n",
        "    youtubeObject = youtubeObject.streams.get_highest_resolution()\n",
        "    try:\n",
        "        youtubeObject.download(filename=\"input.mp4\")\n",
        "    except:\n",
        "        print(\"An error has occurred\")\n",
        "    print(\"Download is completed successfully\")\n",
        "\n",
        "# Call the download function to get the video\n",
        "Download(url)\n",
        "\n",
        "# Use FFmpeg to convert the video to WAV format\n",
        "os.system('ffmpeg -i input.mp4 -vn -acodec pcm_s16le -ar 44100 -ac 2 test_file.wav')\n",
        "\n",
        "# Load the resulting WAV file\n",
        "waveform, sample_rate = torchaudio.load(\"test_file.wav\")\n",
        "\n",
        "# Function to translate text to speech\n",
        "def translate(wav_data):\n",
        "    if torch.cuda.is_available():\n",
        "        device = torch.device(\"cuda:0\")\n",
        "        dtype = torch.float16\n",
        "        logger.info(f\"Running inference on the GPU in {dtype}.\")\n",
        "    else:\n",
        "        device = torch.device(\"cpu\")\n",
        "        dtype = torch.float32\n",
        "        logger.info(f\"Running inference on the CPU in {dtype}.\")\n",
        "\n",
        "    # Initialize the Translator\n",
        "    translator = Translator('seamlessM4T_large', 'vocoder_36langs', device, dtype)\n",
        "\n",
        "    # Perform translation\n",
        "    translated_text, wav, sr = translator.predict(\n",
        "        wav_data,\n",
        "        \"s2st\",\n",
        "        \"eng\",\n",
        "        src_lang=None,\n",
        "        spkr=2,\n",
        "        sample_rate=sample_rate,\n",
        "        ngram_filtering=True,\n",
        "    )\n",
        "    return wav, sr\n",
        "\n",
        "# Define the duration of each chunk in seconds\n",
        "chunk_duration = 25\n",
        "\n",
        "# Convert chunk duration to samples\n",
        "chunk_samples = int(chunk_duration * sample_rate)\n",
        "\n",
        "# Initialize a tensor to store the translated audio\n",
        "store = torch.empty(0, 1, dtype=torch.float32, device='cuda:0')\n",
        "\n",
        "# Iterate through the audio waveform in chunks\n",
        "for chunk in tqdm(range(0, waveform.size(1), chunk_samples)):\n",
        "    val = waveform[:, chunk:chunk+chunk_samples].to('cuda:0')\n",
        "    val = val.t()\n",
        "    wav, sr = translate(val)\n",
        "    store = torch.cat((store, wav[0].t()), dim=0)\n",
        "\n",
        "# Save the translated audio as a WAV file\n",
        "torchaudio.save(\n",
        "    f\"translated_out.wav\",\n",
        "    store.t().to(torch.float32).cpu(),\n",
        "    sample_rate=int(sr/3),\n",
        ")"
      ],
      "metadata": {
        "id": "APz4uGs8_B4I"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}