wave2lip_fine_tune.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f4416c96",
   "metadata": {},
   "source": [
    "# An sample to finetune wave2lip on SageMaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95febd44",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## Update sagemaker python sdk version\n",
    "!pip install -U sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "158e94ab",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "import boto3\n",
    "from sagemaker import get_execution_role\n",
    "\n",
    "sess = sagemaker.Session()\n",
    "role = get_execution_role()\n",
    "sagemaker_default_bucket = sess.default_bucket()\n",
    "\n",
    "account = sess.boto_session.client(\"sts\").get_caller_identity()[\"Account\"]\n",
    "region = sess.boto_session.region_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82a0c829",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## download training script from github\n",
    "!rm -rf ./wav2lip\n",
    "!git clone https://github.com/Rudrabha/Wav2Lip.git"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bc5a450",
   "metadata": {},
   "source": [
    "## Download pretrained model(expert Discriminator & face detect) and upload to s3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "245e0bed-ee4f-43d5-afbb-c61907353552",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd ./wav2lip/face_detection && wget https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth\n",
    "!cd ./wav2lip/models && wget https://iiitaphyd-my.sharepoint.com/:u:/g/personal/radrabha_m_research_iiit_ac_in/EQRvmiZg-HRAjvI6zqN9eTEBP74KefynCwPWVmF57l-AYA?e=ZRPHKP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38de4454",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "!chmod +x ./s5cmd\n",
    "!./s5cmd sync ./wav2lip/ s3://${sagemaker_default_bucket}/models/wav2lip/ "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9588d498",
   "metadata": {},
   "source": [
    "## Prepare docker image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2057f24",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile Dockerfile\n",
    "## You should change below region code to the region you used, here sample is use us-west-2\n",
    "From 763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training:1.13.1-transformers4.26.0-gpu-py39-cu117-ubuntu20.04 \n",
    "#From pytorch/pytorch:1.5-cuda10.1-cudnn7-runtime\n",
    "\n",
    "ENV LANG=C.UTF-8\n",
    "ENV PYTHONUNBUFFERED=TRUE\n",
    "ENV PYTHONDONTWRITEBYTECODE=TRUE\n",
    "\n",
    "RUN apt-get install ffmpeg\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b8ee553",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## You should change below region code to the region you used, here sample is use us-west-2\n",
    "!aws ecr get-login-password --region us-west-2 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-west-2.amazonaws.com"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "985b0fb6",
   "metadata": {},
   "source": [
    "**Build image and push to ECR.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53800617",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## define repo name, should contain *sagemaker* in the name\n",
    "repo_name = \"sagemaker-wav2lip-demo\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9d98c7c",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%script env repo_name=$repo_name bash\n",
    "\n",
    "#!/usr/bin/env bash\n",
    "\n",
    "# This script shows how to build the Docker image and push it to ECR to be ready for use\n",
    "# by SageMaker.\n",
    "\n",
    "# The argument to this script is the image name. This will be used as the image on the local\n",
    "# machine and combined with the account and region to form the repository name for ECR.\n",
    "# The name of our algorithm\n",
    "algorithm_name=${repo_name}\n",
    "\n",
    "account=$(aws sts get-caller-identity --query Account --output text)\n",
    "\n",
    "# Get the region defined in the current configuration (default to us-west-2 if none defined)\n",
    "region=$(aws configure get region)\n",
    "region=${region:-us-west-2}\n",
    "\n",
    "fullname=\"${account}.dkr.ecr.${region}.amazonaws.com/${algorithm_name}:latest\"\n",
    "\n",
    "# If the repository doesn't exist in ECR, create it.\n",
    "aws ecr describe-repositories --repository-names \"${algorithm_name}\" > /dev/null 2>&1\n",
    "\n",
    "if [ $? -ne 0 ]\n",
    "then\n",
    "    aws ecr create-repository --repository-name \"${algorithm_name}\" > /dev/null\n",
    "fi\n",
    "\n",
    "# Get the login command from ECR and execute it directly\n",
    "aws ecr get-login-password --region ${region}|docker login --username AWS --password-stdin ${fullname}\n",
    "\n",
    "# Build the docker image locally with the image name and then push it to ECR\n",
    "# with the full name.\n",
    "\n",
    "docker build -t ${algorithm_name} .\n",
    "docker tag ${algorithm_name} ${fullname}\n",
    "\n",
    "docker push ${fullname}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06c51767",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Train!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6aef55e2",
   "metadata": {},
   "source": [
    "use s3 dataset path, which is aligned with data_root of wav2lip \n",
    "should be like:\n",
    "data_root (mvlrs_v1)\n",
    "├── main, pretrain (we use only main folder in this work)\n",
    "|\t├── list of folders\n",
    "|\t│   ├── five-digit numbered video IDs ending with (.mp4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55cedcb2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile train.sh\n",
    "#!/bin/bash\n",
    "\n",
    "chmod +x ./s5cmd\n",
    "#./s5cmd sync  s3://${sagemaker_default_bucket}/models/wav2lip/ ./wav2lip/\n",
    "pip install -r ./wav2lip/requirements.txt\n",
    "\n",
    "###data process\n",
    "python ./wav2lip/preprocess.py --data_root /opt/ml/data/inputs/data_root/main --preprocessed_root /tmp/lrs2_preprocessed/\n",
    "###train the expert discriminator\n",
    "python ./wav2lip/color_syncnet_train.py --data_root /tmp/lrs2_preprocessed/ --checkpoint_dir /tmp/trained_syncnet/\n",
    "###train wav2lip \n",
    "python ./wav2lip/hq_wav2lip_train.py --data_root /tmp/lrs2_preprocessed/ --checkpoint_dir /tmp/trained_wav2lip/ --syncnet_checkpoint_path /tmp/trained_syncnet/\n",
    "\n",
    "./s5cmd sync /tmp/trained_wav2lip/ s3://$MODEL_S3_BUCKET/models/wav2lip/output/$(date +%Y-%m-%d-%H-%M-%S)/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "517889b1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## The image uri which is build and pushed above\n",
    "image_uri = \"{}.dkr.ecr.{}.amazonaws.com/{}:latest\".format(account, region, repo_name)\n",
    "image_uri"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad795754",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## set train_data_path to your training dataset path in s3\n",
    "train_data_path = f's3://{sagemaker_default_bucket}/wav2lip/train_data/'\n",
    "\n",
    "inputs = {'data_root': train_data_path}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ab36100",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from sagemaker.estimator import Estimator\n",
    "\n",
    "environment = {\n",
    "              'MODEL_S3_BUCKET': sagemaker_default_bucket # The bucket to store pretrained model and fine-tune model\n",
    "}\n",
    "\n",
    "base_job_name = 'wav2lip-demo'         \n",
    "\n",
    "instance_type = 'ml.p4d.24xlarge'\n",
    "\n",
    "estimator = Estimator(role=role,\n",
    "                      entry_point='train.sh',\n",
    "                      source_dir='./',\n",
    "                      base_job_name=base_job_name,\n",
    "                      instance_count=1,\n",
    "                      instance_type=instance_type,\n",
    "                      image_uri=image_uri,\n",
    "                      environment=environment,\n",
    "                      disable_profiler=True,\n",
    "                      debugger_hook_config=False,\n",
    "                      max_run=24*60*60*2)\n",
    "\n",
    "estimator.fit(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a3e673c",
   "metadata": {},
   "source": [
    "You could find the model path in S3 from above logs."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p39",
   "language": "python",
   "name": "conda_pytorch_p39"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}