{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# EMAGE: Co-Speech 3D Gesture Generation\n", "\n", "This notebook demonstrates three models for generating body gestures from speech:\n", "- **CaMN**: Upper body gesture generation\n", "- **DisCo**: Upper body gesture generation with diffusion\n", "- **EMAGE**: Full body + face gesture generation\n", "\n", "[Project Page](https://pantomatrix.github.io/EMAGE/) | [GitHub](https://github.com/PantoMatrix/PantoMatrix) | [Paper](https://arxiv.org/abs/2401.00374) | [HF Space](https://huggingface.co/spaces/H-Liu1997/EMAGE)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 1. Setup Environment\n", "\n", "Install Python 3.10, create virtual environment, and install dependencies." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Install Python 3.10 and create virtual environment\n", "!sudo add-apt-repository -y ppa:deadsnakes/ppa > /dev/null 2>&1\n", "!sudo apt-get update -qq\n", "!sudo apt-get install -y python3.10 python3.10-venv python3.10-dev > /dev/null 2>&1\n", "\n", "ENV_PATH = \"/content/py310_env\"\n", "!python3.10 -m venv {ENV_PATH}\n", "\n", "PYTHON = f\"{ENV_PATH}/bin/python\"\n", "PIP = f\"{ENV_PATH}/bin/pip\"\n", "\n", "# Install dependencies\n", "!{PIP} install -q --upgrade pip\n", "!{PIP} install -q torch==2.1.2 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121\n", "!{PIP} install -q numpy==1.23.0 librosa soundfile transformers huggingface_hub\n", "!{PIP} install -q smplx trimesh scipy easydict omegaconf\n", "\n", "# Verify\n", "!{PYTHON} -c \"import torch; print(f'Python 3.10 + PyTorch {torch.__version__} + CUDA: {torch.cuda.is_available()}')\"" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Clone code repositories\n", "!apt-get install -y git-lfs > /dev/null 2>&1\n", "!git lfs install\n", "\n", "!git clone https://github.com/PantoMatrix/PantoMatrix.git /content/PantoMatrix\n", "!git clone https://huggingface.co/H-Liu1997/emage_evaltools /content/PantoMatrix/emage_evaltools\n", "%cd /content/PantoMatrix/emage_evaltools\n", "!git lfs pull\n", "%cd /content\n", "\n", "print(\"Code ready!\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%%writefile /content/run_inference.py\n", "import sys, os\n", "sys.path.insert(0, '/content/PantoMatrix')\n", "os.chdir('/content/PantoMatrix')\n", "\n", "import torch, numpy as np, librosa, argparse\n", "import torch.nn.functional as F\n", "from torchvision.io import write_video\n", "from models.camn_audio import CamnAudioModel\n", "from models.disco_audio import DiscoAudioModel\n", "from models.emage_audio import EmageAudioModel, EmageVQVAEConv, EmageVAEConv, EmageVQModel\n", "from emage_utils.motion_io import beat_format_save\n", "from emage_utils.npz2pose import render2d\n", "\n", "def main():\n", " parser = argparse.ArgumentParser()\n", " parser.add_argument('--audio', type=str, required=True)\n", " parser.add_argument('--model', type=str, default='camn', choices=['camn', 'disco', 'emage'])\n", " parser.add_argument('--output_dir', type=str, default='/content/outputs')\n", " args = parser.parse_args()\n", " \n", " os.makedirs(args.output_dir, exist_ok=True)\n", " device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n", " print(f'Using device: {device}')\n", " \n", " if args.model == 'camn':\n", " model = CamnAudioModel.from_pretrained('H-Liu1997/camn_audio').to(device).eval()\n", " sr, fps, seed = model.cfg.audio_sr, model.cfg.pose_fps, model.cfg.seed_frames\n", " audio, _ = librosa.load(args.audio, sr=sr)\n", " audio_t = torch.from_numpy(audio).float().unsqueeze(0).to(device)\n", " with torch.no_grad():\n", " motion = model(audio_t, torch.zeros(1,1).long().to(device), seed_frames=seed)['motion_axis_angle']\n", " npz_path = os.path.join(args.output_dir, 'camn_output.npz')\n", " beat_format_save(npz_path, motion.cpu().numpy().reshape(motion.shape[1], -1), upsample=30//fps)\n", " \n", " elif args.model == 'disco':\n", " model = DiscoAudioModel.from_pretrained('H-Liu1997/disco_audio').to(device).eval()\n", " sr, fps, seed = model.cfg.audio_sr, model.cfg.pose_fps, model.cfg.seed_frames\n", " audio, _ = librosa.load(args.audio, sr=sr)\n", " audio_t = torch.from_numpy(audio).float().unsqueeze(0).to(device)\n", " with torch.no_grad():\n", " motion = model(audio_t, torch.zeros(1,1).long().to(device), seed_frames=seed, seed_motion=None)['motion_axis_angle']\n", " npz_path = os.path.join(args.output_dir, 'disco_output.npz')\n", " beat_format_save(npz_path, motion.cpu().numpy().reshape(motion.shape[1], -1), upsample=30//fps)\n", " \n", " else: # emage\n", " vq_models = {k: EmageVQVAEConv.from_pretrained('H-Liu1997/emage_audio', subfolder=f'emage_vq/{k}').to(device).eval() \n", " for k in ['face', 'upper', 'lower', 'hands']}\n", " global_ae = EmageVAEConv.from_pretrained('H-Liu1997/emage_audio', subfolder='emage_vq/global').to(device).eval()\n", " vq = EmageVQModel(face_model=vq_models['face'], upper_model=vq_models['upper'],\n", " lower_model=vq_models['lower'], hands_model=vq_models['hands'], global_model=global_ae).to(device).eval()\n", " model = EmageAudioModel.from_pretrained('H-Liu1997/emage_audio').to(device).eval()\n", " sr, fps = model.cfg.audio_sr, model.cfg.pose_fps\n", " audio, _ = librosa.load(args.audio, sr=sr)\n", " audio_t = torch.from_numpy(audio).float().unsqueeze(0).to(device)\n", " with torch.no_grad():\n", " lat = model.inference(audio_t, torch.zeros(1,1).long().to(device), vq, masked_motion=None, mask=None)\n", " get = lambda k, c: lat[f'rec_{k}'] if getattr(model.cfg, f'l{k[0]}') > 0 and getattr(model.cfg, f'c{k[0]}') == 0 else None\n", " idx = lambda k: torch.max(F.log_softmax(lat[f'cls_{k}'], dim=2), dim=2)[1] if getattr(model.cfg, f'c{k[0]}') > 0 else None\n", " pred = vq.decode(face_latent=get('face','f'), upper_latent=get('upper','u'), lower_latent=get('lower','l'), hands_latent=get('hands','h'),\n", " face_index=idx('face'), upper_index=idx('upper'), lower_index=idx('lower'), hands_index=idx('hands'),\n", " get_global_motion=True, ref_trans=torch.zeros(1,3).to(device))\n", " motion = pred['motion_axis_angle']\n", " npz_path = os.path.join(args.output_dir, 'emage_output.npz')\n", " beat_format_save(npz_path, motion.cpu().numpy().reshape(motion.shape[1], -1), upsample=30//fps,\n", " expressions=pred['expression'].cpu().numpy().reshape(motion.shape[1], -1),\n", " trans=pred['trans'].cpu().numpy().reshape(motion.shape[1], -1))\n", " \n", " # Render 2D visualization\n", " motion_dict = np.load(npz_path, allow_pickle=True)\n", " v2d = render2d(motion_dict, (720, 480), face_only=False, remove_global=True)\n", " video_path = npz_path.replace('.npz', '_2d.mp4')\n", " write_video(video_path, v2d.permute(0, 2, 3, 1), fps=30)\n", " print(f'Saved: {npz_path}')\n", " print(f'Video: {video_path}')\n", "\n", "if __name__ == '__main__': main()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 2. Run Inference\n", "\n", "Choose your audio and model, then run inference." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Audio file (use example or upload your own)\n", "audio_path = \"/content/PantoMatrix/examples/audio/2_scott_0_103_103_10s.wav\"\n", "\n", "# Uncomment to upload your own audio:\n", "# from google.colab import files\n", "# uploaded = files.upload()\n", "# audio_path = \"/content/\" + list(uploaded.keys())[0]\n", "\n", "from IPython.display import Audio\n", "display(Audio(audio_path))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Run CaMN (Upper Body)\n", "PYTHON = \"/content/py310_env/bin/python\"\n", "!{PYTHON} /content/run_inference.py --audio {audio_path} --model camn\n", "\n", "from IPython.display import Video\n", "display(Video(\"/content/outputs/camn_output_2d.mp4\", embed=True, width=600))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Run DisCo (Upper Body with Diffusion)\n", "PYTHON = \"/content/py310_env/bin/python\"\n", "!{PYTHON} /content/run_inference.py --audio {audio_path} --model disco\n", "\n", "from IPython.display import Video\n", "display(Video(\"/content/outputs/disco_output_2d.mp4\", embed=True, width=600))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Run EMAGE (Full Body + Face)\n", "PYTHON = \"/content/py310_env/bin/python\"\n", "!{PYTHON} /content/run_inference.py --audio {audio_path} --model emage\n", "\n", "from IPython.display import Video\n", "display(Video(\"/content/outputs/emage_output_2d.mp4\", embed=True, width=600))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 3. Download Results\n", "\n", "Download motion files (`.npz`) for use with Blender." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "print(\"Generated files:\")\n", "for f in os.listdir(\"/content/outputs\"):\n", " print(f\" /content/outputs/{f}\")\n", "\n", "# Uncomment to download:\n", "# from google.colab import files\n", "# files.download(\"/content/outputs/camn_output.npz\")\n", "# files.download(\"/content/outputs/disco_output.npz\")\n", "# files.download(\"/content/outputs/emage_output.npz\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Notes\n", "\n", "- **Environment**: Python 3.10.x + PyTorch 2.1.2 + CUDA 12.1\n", "- **Motion Format**: `.npz` files contain SMPL-X format motion data\n", "- **Visualization**: Use the [Blender Add-on](https://huggingface.co/datasets/H-Liu1997/BEAT2_Tools/blob/main/smplx_blender_addon_20230921.zip) for high-quality rendering\n", "- **Interactive Demo**: [HuggingFace Space](https://huggingface.co/spaces/H-Liu1997/EMAGE)" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python", "version": "3.10.0" } }, "nbformat": 4, "nbformat_minor": 4 }