JoyVASA_jupyter.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/JoyVASA-jupyter/blob/main/JoyVASA_jupyter.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VjYy0F2gZIPR"
   },
   "outputs": [],
   "source": [
    "%cd /content\n",
    "!git clone https://github.com/jdh-algo/JoyVASA\n",
    "!pip install tyro gradio onnx onnxruntime onnxruntime-gpu pykalman colorama\n",
    "\n",
    "!apt install aria2 -qqy\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/JoyVASA/motion_generator/motion_generator_hubert_chinese.pt -d /content/JoyVASA/pretrained_weights/JoyVASA/motion_generator -o motion_generator_hubert_chinese.pt\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/JoyVASA/motion_template/motion_template.pkl -d /content/JoyVASA/pretrained_weights/JoyVASA/motion_template -o motion_template.pkl\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/TencentGameMate%3Achinese-hubert-base/chinese-hubert-base-fairseq-ckpt.pt -d /content/JoyVASA/pretrained_weights/TencentGameMate:chinese-hubert-base -o chinese-hubert-base-fairseq-ckpt.pt\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/raw/main/TencentGameMate%3Achinese-hubert-base/config.json -d /content/JoyVASA/pretrained_weights/TencentGameMate:chinese-hubert-base -o config.json\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/raw/main/TencentGameMate%3Achinese-hubert-base/preprocessor_config.json -d /content/JoyVASA/pretrained_weights/TencentGameMate:chinese-hubert-base -o preprocessor_config.json\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/TencentGameMate%3Achinese-hubert-base/pytorch_model.bin -d /content/JoyVASA/pretrained_weights/TencentGameMate:chinese-hubert-base -o pytorch_model.bin\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/insightface/models/buffalo_l/2d106det.onnx -d /content/JoyVASA/pretrained_weights/insightface/models/buffalo_l -o 2d106det.onnx\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/insightface/models/buffalo_l/det_10g.onnx -d /content/JoyVASA/pretrained_weights/insightface/models/buffalo_l -o det_10g.onnx\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/liveportrait/landmark.onnx -d /content/JoyVASA/pretrained_weights/liveportrait -o landmark.onnx\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/liveportrait/retargeting_models/stitching_retargeting_module.pth -d /content/JoyVASA/pretrained_weights/liveportrait/retargeting_models -o stitching_retargeting_module.pth\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/liveportrait/base_models/appearance_feature_extractor.pth -d /content/JoyVASA/pretrained_weights/liveportrait/base_models -o appearance_feature_extractor.pth\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/liveportrait/base_models/motion_extractor.pth -d /content/JoyVASA/pretrained_weights/liveportrait/base_models -o motion_extractor.pth\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/liveportrait/base_models/spade_generator.pth -d /content/JoyVASA/pretrained_weights/liveportrait/base_models -o spade_generator.pth\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/liveportrait/base_models/warping_module.pth -d /content/JoyVASA/pretrained_weights/liveportrait/base_models -o warping_module.pth\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/liveportrait_animals/xpose.pth -d /content/JoyVASA/pretrained_weights/liveportrait_animals -o xpose.pth\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/liveportrait_animals/retargeting_models/stitching_retargeting_module.pth -d /content/JoyVASA/pretrained_weights/liveportrait_animals/retargeting_models -o stitching_retargeting_module.pth\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/liveportrait_animals/base_models/appearance_feature_extractor.pth -d /content/JoyVASA/pretrained_weights/liveportrait_animals/base_models -o appearance_feature_extractor.pth\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/liveportrait_animals/base_models/motion_extractor.pth -d /content/JoyVASA/pretrained_weights/liveportrait_animals/base_models -o motion_extractor.pth\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/liveportrait_animals/base_models/spade_generator.pth -d /content/JoyVASA/pretrained_weights/liveportrait_animals/base_models -o spade_generator.pth\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/liveportrait_animals/base_models/warping_module.pth -d /content/JoyVASA/pretrained_weights/liveportrait_animals/base_models -o warping_module.pth\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/raw/main/wav2vec2-base-960h/config.json -d /content/JoyVASA/pretrained_weights/wav2vec2-base-960h -o config.json\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/raw/main/wav2vec2-base-960h/feature_extractor_config.json -d /content/JoyVASA/pretrained_weights/wav2vec2-base-960h -o feature_extractor_config.json\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/wav2vec2-base-960h/model.safetensors -d /content/JoyVASA/pretrained_weights/wav2vec2-base-960h -o model.safetensors\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/raw/main/wav2vec2-base-960h/preprocessor_config.json -d /content/JoyVASA/pretrained_weights/wav2vec2-base-960h -o preprocessor_config.json\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/wav2vec2-base-960h/pytorch_model.bin -d /content/JoyVASA/pretrained_weights/wav2vec2-base-960h -o pytorch_model.bin\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/raw/main/wav2vec2-base-960h/special_tokens_map.json -d /content/JoyVASA/pretrained_weights/wav2vec2-base-960h -o special_tokens_map.json\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/resolve/main/wav2vec2-base-960h/tf_model.h5 -d /content/JoyVASA/pretrained_weights/wav2vec2-base-960h -o tf_model.h5\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/raw/main/wav2vec2-base-960h/tokenizer_config.json -d /content/JoyVASA/pretrained_weights/wav2vec2-base-960h -o tokenizer_config.json\n",
    "!aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/JoyVASA/raw/main/wav2vec2-base-960h/vocab.json -d /content/JoyVASA/pretrained_weights/wav2vec2-base-960h -o vocab.json\n",
    "\n",
    "%cd /content/JoyVASA/src/utils/dependencies/XPose/models/UniPose/ops\n",
    "!python setup.py build install"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Runtime/Restart session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /content/JoyVASA\n",
    "\n",
    "import os\n",
    "import gradio as gr\n",
    "import os.path as osp\n",
    "from src.gradio_pipeline import GradioPipeline, GradioPipelineAnimal\n",
    "from src.config.crop_config import CropConfig\n",
    "from src.config.argument_config import ArgumentConfig\n",
    "from src.config.inference_config import InferenceConfig\n",
    "\n",
    "args = ArgumentConfig(\n",
    "    gradio_temp_dir=\"gradio_temp\",\n",
    "    server_port=7860,\n",
    "    share=True,\n",
    "    server_name=\"127.0.0.1\",\n",
    ")\n",
    "\n",
    "inference_cfg = InferenceConfig()\n",
    "crop_cfg = CropConfig()\n",
    "\n",
    "gradio_pipeline_human = GradioPipeline(\n",
    "    inference_cfg=inference_cfg,\n",
    "    crop_cfg=crop_cfg,\n",
    "    args=args\n",
    ")\n",
    "gradio_pipeline_animal = GradioPipelineAnimal(\n",
    "    inference_cfg=inference_cfg,\n",
    "    crop_cfg=crop_cfg,\n",
    "    args=args\n",
    ")\n",
    "\n",
    "def gpu_wrapped_execute_a2v(*args, **kwargs):\n",
    "    if args[5] == \"animal\":\n",
    "        return gradio_pipeline_animal.execute_a2v(*args, **kwargs)\n",
    "    else:\n",
    "        return gradio_pipeline_human.execute_a2v(*args, **kwargs)\n",
    "\n",
    "with gr.Blocks(theme=gr.themes.Soft(font=[gr.themes.GoogleFont(\"Plus Jakarta Sans\")])) as demo:\n",
    "    with gr.Row():\n",
    "        with gr.Accordion(open=True, label=\"🖼️ Reference Image\"):\n",
    "            input_image = gr.Image(type=\"filepath\", width=512, label=\"Reference Image\")\n",
    "        with gr.Accordion(open=True, label=\"🎵 Input Audio\"):\n",
    "            input_audio = gr.Audio(type=\"filepath\", label=\"Input Audio\")\n",
    "        with gr.Accordion(open=True, label=\"🎬 Output Video\",):\n",
    "            output_video = gr.Video(autoplay=False, interactive=False, width=512)     \n",
    "    with gr.Column():\n",
    "        with gr.Accordion(open=True, label=\"Key Animation Options\"):\n",
    "            with gr.Row():\n",
    "                animation_mode = gr.Radio(['human', 'animal'], value=\"human\", label=\"Animation Mode\") \n",
    "                flag_do_crop_input = gr.Checkbox(value=True, label=\"do crop (image)\")\n",
    "                cfg_scale = gr.Number(value=4.0, label=\"cfg_scale\", minimum=0.0, maximum=10.0, step=0.5)\n",
    "        with gr.Accordion(open=False, label=\"Optional Animation Options\"):\n",
    "            with gr.Row():\n",
    "                driving_option_input = gr.Radio(['expression-friendly', 'pose-friendly'], value=\"expression-friendly\", label=\"driving option\")\n",
    "                driving_multiplier = gr.Number(value=1.0, label=\"driving multiplier\", minimum=0.0, maximum=2.0, step=0.02)\n",
    "            with gr.Row():\n",
    "                flag_normalize_lip = gr.Checkbox(value=True, label=\"normalize lip\")\n",
    "                flag_relative_motion = gr.Checkbox(value=True, label=\"relative motion\")\n",
    "                flag_remap_input = gr.Checkbox(value=True, label=\"paste-back\")\n",
    "                flag_stitching_input = gr.Checkbox(value=True, label=\"stitching\")\n",
    "        with gr.Accordion(open=False, label=\"Optional Options for Reference Image\"):\n",
    "            with gr.Row():\n",
    "                scale = gr.Number(value=2.3, label=\"image crop scale\", minimum=1.8, maximum=4.0, step=0.05)\n",
    "                vx_ratio = gr.Number(value=0.0, label=\"image crop x\", minimum=-0.5, maximum=0.5, step=0.01)\n",
    "                vy_ratio = gr.Number(value=-0.125, label=\"image crop y\", minimum=-0.5, maximum=0.5, step=0.01)\n",
    "\n",
    "    with gr.Row():\n",
    "        process_button_generate = gr.Button(\"🚀 Generate\", variant=\"primary\")\n",
    "    generation_func = gpu_wrapped_execute_a2v\n",
    "    process_button_generate.click(\n",
    "        fn=generation_func,\n",
    "        inputs=[\n",
    "            input_image,\n",
    "            input_audio,\n",
    "            flag_normalize_lip,\n",
    "            flag_relative_motion,\n",
    "            driving_multiplier,\n",
    "            animation_mode,\n",
    "            driving_option_input,\n",
    "            flag_do_crop_input,\n",
    "            scale,\n",
    "            vx_ratio,\n",
    "            vy_ratio,\n",
    "            flag_stitching_input,\n",
    "            flag_remap_input,\n",
    "            cfg_scale,\n",
    "        ],\n",
    "        outputs=[\n",
    "            output_video,\n",
    "        ],\n",
    "        show_progress=True\n",
    "    )\n",
    "\n",
    "demo.launch(\n",
    "    server_port=args.server_port,\n",
    "    share=args.share,\n",
    "    server_name=args.server_name,\n",
    "    inline=False\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}