dreamtalk_gradio_colab.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/camenduru/dreamtalk-colab/blob/main/dreamtalk_gradio_colab.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VjYy0F2gZIPR"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "!git clone -b dev https://github.com/camenduru/dreamtalk\n",
        "%cd /content/dreamtalk\n",
        "\n",
        "!wget https://huggingface.co/camenduru/dreamtalk/resolve/main/damo/dreamtalk/checkpoints/denoising_network.pth -O /content/dreamtalk/checkpoints/denoising_network.pth\n",
        "!wget https://huggingface.co/camenduru/dreamtalk/resolve/main/damo/dreamtalk/checkpoints/renderer.pt -O /content/dreamtalk/checkpoints/renderer.pt\n",
        "\n",
        "!pip install -q yacs av gradio\n",
        "\n",
        "# https://huggingface.co/spaces/fffiloni/dreamtalk/blob/main/app.py modified\n",
        "\n",
        "import gradio as gr\n",
        "import subprocess\n",
        "from moviepy.editor import VideoFileClip\n",
        "import datetime\n",
        "\n",
        "def convert_to_mp4_with_aac(input_path, output_path):\n",
        "    video = VideoFileClip(input_path)\n",
        "    video.write_videofile(output_path, codec=\"libx264\", audio_codec=\"aac\")\n",
        "    return output_path\n",
        "\n",
        "def check_file_exists(file_path, audio_list):\n",
        "    return file_path in audio_list\n",
        "\n",
        "def load_audio(audio_listed):\n",
        "    if audio_listed is None:\n",
        "        return None\n",
        "    else:\n",
        "        return f\"data/audio/{audio_listed}\"\n",
        "\n",
        "def execute_command(command: str) -> None:\n",
        "    subprocess.run(command, check=True)\n",
        "\n",
        "def infer(audio_input, image_path, emotional_style):\n",
        "    timestamp = datetime.datetime.now().strftime(\"%Y%m%d%H%M%S\")\n",
        "    output_name = f\"lipsynced_result_{timestamp}\"\n",
        "    command = [\n",
        "        f\"python\",\n",
        "        f\"inference_for_demo_video.py\",\n",
        "        f\"--wav_path={audio_input}\",\n",
        "        f\"--style_clip_path=data/style_clip/3DMM/{emotional_style}\",\n",
        "        f\"--pose_path=data/pose/RichardShelby_front_neutral_level1_001.mat\",\n",
        "        f\"--image_path={image_path}\",\n",
        "        f\"--cfg_scale=1.0\",\n",
        "        f\"--max_gen_len=30\",\n",
        "        f\"--output_name={output_name}\"\n",
        "    ]\n",
        "\n",
        "    execute_command(command)\n",
        "    input_file = f\"output_video/{output_name}.mp4\"\n",
        "    output_file = f\"{output_name}.mp4\"\n",
        "    result = convert_to_mp4_with_aac(input_file, output_file)\n",
        "    return result\n",
        "\n",
        "css=\"\"\"\n",
        "#project-links{\n",
        "    margin: 0 0 12px !important;\n",
        "    column-gap: 8px;\n",
        "    display: flex;\n",
        "    justify-content: center;\n",
        "    flex-wrap: nowrap;\n",
        "    flex-direction: row;\n",
        "    align-items: center;\n",
        "}\n",
        "\"\"\"\n",
        "with gr.Blocks(css=css) as demo:\n",
        "    with gr.Column(elem_id=\"col-container\"):\n",
        "        gr.HTML(\"\"\"\n",
        "        <h2 style=\"text-align: center;\">DreamTalk</h2>\n",
        "        <p style=\"text-align: center;\">When Expressive Talking Head Generation Meets Diffusion Probabilistic Models</p>        \n",
        "        \"\"\")\n",
        "        with gr.Row():\n",
        "            with gr.Column():\n",
        "                image_path = gr.Image(label=\"Image\", type=\"filepath\", sources=[\"upload\"])\n",
        "                audio_input = gr.Audio(label=\"Audio input\", type=\"filepath\", sources=[\"upload\"], value=\"data/audio/acknowledgement_english.m4a\")\n",
        "                with gr.Row():\n",
        "                    audio_list = gr.Dropdown(\n",
        "                        label=\"Choose an audio (optional)\",\n",
        "                        choices=[\n",
        "                            \"German1.wav\", \"German2.wav\", \"German3.wav\", \"German4.wav\",\n",
        "                            \"acknowledgement_chinese.m4a\", \"acknowledgement_english.m4a\",\n",
        "                            \"chinese1_haierlizhi.wav\", \"chinese2_guanyu.wav\",\n",
        "                            \"french1.wav\", \"french2.wav\", \"french3.wav\",\n",
        "                            \"italian1.wav\", \"italian2.wav\", \"italian3.wav\",\n",
        "                            \"japan1.wav\", \"japan2.wav\", \"japan3.wav\",\n",
        "                            \"korean1.wav\", \"korean2.wav\", \"korean3.wav\",\n",
        "                            \"noisy_audio_cafeter_snr_0.wav\", \"noisy_audio_meeting_snr_0.wav\", \"noisy_audio_meeting_snr_10.wav\", \"noisy_audio_meeting_snr_20.wav\", \"noisy_audio_narrative.wav\", \"noisy_audio_office_snr_0.wav\", \"out_of_domain_narrative.wav\",\n",
        "                            \"spanish1.wav\", \"spanish2.wav\", \"spanish3.wav\"\n",
        "                            ],\n",
        "                        value = \"acknowledgement_english.m4a\"\n",
        "                    )\n",
        "                    audio_list.change(\n",
        "                        fn = load_audio,\n",
        "                        inputs = [audio_list],\n",
        "                        outputs = [audio_input]\n",
        "                    )\n",
        "                    emotional_style = gr.Dropdown(\n",
        "                        label = \"emotional style\",\n",
        "                        choices = [\n",
        "                            \"M030_front_angry_level3_001.mat\",\n",
        "                            \"M030_front_contempt_level3_001.mat\",\n",
        "                            \"M030_front_disgusted_level3_001.mat\",\n",
        "                            \"M030_front_fear_level3_001.mat\",\n",
        "                            \"M030_front_happy_level3_001.mat\",\n",
        "                            \"M030_front_neutral_level1_001.mat\",\n",
        "                            \"M030_front_sad_level3_001.mat\",\n",
        "                            \"M030_front_surprised_level3_001.mat\",\n",
        "                            \"W009_front_angry_level3_001.mat\",\n",
        "                            \"W009_front_contempt_level3_001.mat\",\n",
        "                            \"W009_front_disgusted_level3_001.mat\",\n",
        "                            \"W009_front_fear_level3_001.mat\",\n",
        "                            \"W009_front_happy_level3_001.mat\",\n",
        "                            \"W009_front_neutral_level1_001.mat\",\n",
        "                            \"W009_front_sad_level3_001.mat\",\n",
        "                            \"W009_front_surprised_level3_001.mat\",\n",
        "                            \"W011_front_angry_level3_001.mat\",\n",
        "                            \"W011_front_contempt_level3_001.mat\",\n",
        "                            \"W011_front_disgusted_level3_001.mat\",\n",
        "                            \"W011_front_fear_level3_001.mat\",\n",
        "                            \"W011_front_happy_level3_001.mat\",\n",
        "                            \"W011_front_neutral_level1_001.mat\",\n",
        "                            \"W011_front_sad_level3_001.mat\",\n",
        "                            \"W011_front_surprised_level3_001.mat\"\n",
        "                        ],\n",
        "                        value = \"M030_front_neutral_level1_001.mat\"\n",
        "                    )\n",
        "                gr.Examples(\n",
        "                    examples = [\n",
        "                        \"data/src_img/uncropped/face3.png\",\n",
        "                        \"data/src_img/uncropped/male_face.png\",\n",
        "                        \"data/src_img/uncropped/uncut_src_img.jpg\",\n",
        "                        \"data/src_img/cropped/chpa5.png\",\n",
        "                        \"data/src_img/cropped/cut_img.png\",\n",
        "                        \"data/src_img/cropped/f30.png\",\n",
        "                        \"data/src_img/cropped/menglu2.png\",\n",
        "                        \"data/src_img/cropped/nscu2.png\",\n",
        "                        \"data/src_img/cropped/zp1.png\",\n",
        "                        \"data/src_img/cropped/zt12.png\"\n",
        "                    ],\n",
        "                    inputs=[image_path],\n",
        "                    examples_per_page=5\n",
        "                )\n",
        "                with gr.Row():\n",
        "                    gr.ClearButton([audio_input, image_path, audio_list])\n",
        "                    run_btn = gr.Button(\"Run\", elem_id=\"run-btn\")\n",
        "            with gr.Column():\n",
        "                output_video = gr.Video(format=\"mp4\")\n",
        "                gr.HTML(\"\"\"\n",
        "                <p id=\"project-links\" align=\"center\">\n",
        "                  <a href='https://dreamtalk-project.github.io/'><img src='https://img.shields.io/badge/Project-Page-Green'></a> <a href='https://arxiv.org/abs/2312.09767'><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a> <a href='https://youtu.be/VF4vlE6ZqWQ'><img src='https://badges.aleen42.com/src/youtube.svg'></a>\n",
        "                </p>\n",
        "                <img src=\"https://github.com/ali-vilab/dreamtalk/raw/main/media/teaser.gif\" style=\"margin: 0 auto;border-radius: 10px;\" />    \n",
        "                \"\"\")\n",
        "    \n",
        "    run_btn.click(\n",
        "        fn = infer,\n",
        "        inputs = [audio_input, image_path, emotional_style],\n",
        "        outputs = [output_video]\n",
        "    )\n",
        "\n",
        "demo.queue().launch(inline=False, share=True, debug=True)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}