{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SadTalker: Real-Time Audio-Driven Talking Head Animation\n",
    "### Complete Setup & Benchmark Notebook for Vast.ai / JupyterLab / Google Colab\n",
    "\n",
    "This notebook contains the complete, tested execution pipeline for [OpenTalker/SadTalker](https://github.com/OpenTalker/SadTalker) on an NVIDIA GPU.\n",
    "- Includes fixes for `torchvision.transforms.functional_tensor`\n",
    "- Includes Gradio 3.41.2 lock to prevent `Row.style` errors\n",
    "- Uses GFPGAN face enhancer for high-resolution 512px rendering"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Verify NVIDIA GPU & CUDA Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Install System Prerequisites (Git, FFmpeg, Wget)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!apt update && apt install -y git ffmpeg wget"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Clone Repository & Navigate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/OpenTalker/SadTalker.git\n",
    "%cd SadTalker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Install Python Dependencies & PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torch torchvision torchaudio\n",
    "!pip install -r requirements.txt\n",
    "!pip install \"gradio==3.41.2\" TTS==0.13.3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Apply Critical Patch for TorchVision in BasicSR\n",
    "Fixes: `ModuleNotFoundError: No module named 'torchvision.transforms.functional_tensor'`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys, glob\n",
    "site_packages = [p for p in sys.path if 'site-packages' in p]\n",
    "for sp in site_packages:\n",
    "    target = os.path.join(sp, 'basicsr', 'data', 'degradations.py')\n",
    "    if os.path.exists(target):\n",
    "        with open(target, 'r') as f:\n",
    "            content = f.read()\n",
    "        if 'torchvision.transforms.functional_tensor' in content:\n",
    "            content = content.replace('from torchvision.transforms.functional_tensor import rgb_to_grayscale', 'from torchvision.transforms.functional import rgb_to_grayscale')\n",
    "            with open(target, 'w') as f:\n",
    "                f.write(content)\n",
    "            print(f'Successfully patched {target}')\n",
    "        else:\n",
    "            print(f'Already patched or clean: {target}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Download Pretrained Model Checkpoints\n",
    "Downloads checkpoints into `checkpoints/` and `gfpgan/`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!bash scripts/download_models.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Run CLI Inference with Timing Benchmark\n",
    "Generates enhanced talking-head video."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "!python inference.py \\\n",
    "  --driven_audio ./examples/driven_audio/bus_chinese.wav \\\n",
    "  --source_image ./examples/source_image/full_body_1.png \\\n",
    "  --result_dir ./results \\\n",
    "  --enhancer gfpgan"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Launch Local Gradio WebUI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!python app_sadtalker.py --share"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}