init project

Dockerfile

@@ -0,0 +1,42 @@
+FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    wget \
+    ffmpeg \
+    libsndfile1 \
+    python3-pip \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Miniconda
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh \
+    && bash miniconda.sh -b -p /opt/conda \
+    && rm miniconda.sh
+
+# Add conda to path
+ENV PATH=/opt/conda/bin:$PATH
+
+# Create and activate conda environment
+RUN conda create -n mimictalk python=3.9 -y
+SHELL ["conda", "run", "-n", "mimictalk", "/bin/bash", "-c"]
+
+# Install PyTorch and other dependencies
+RUN pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu121 \
+    && pip install cython openmim==0.3.9 \
+    && mim install mmcv==2.1.0 \
+    && pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable"
+
+# Install requirements
+COPY requirements.txt /tmp/
+RUN pip install -r /tmp/requirements.txt
+
+# Download pre-trained models
+RUN pip install huggingface_hub \
+    && python -c "from huggingface_hub import snapshot_download; snapshot_download('mrbear1024/demo-model', local_dir='/app/models')"
+
+# Set working directory
+WORKDIR /app
+
+# Set default command to activate conda env
+CMD ["conda", "run", "--no-capture-output", "-n", "mimictalk", "python", "inference/app_mimictalk.py"]

README.md

@@ -5,6 +5,196 @@ colorFrom: blue
 colorTo: pink
 sdk: docker
 pinned: false
+app_file: inference/app_mimictalk.py
+models:
+  - mrbear1024/demo-model
+
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+
+# MimicTalk: Mimicking a personalized and expressive 3D talking face in few minutes | NeurIPS 2024
+[![arXiv](https://img.shields.io/badge/arXiv-Paper-%3CCOLOR%3E.svg)](https://arxiv.org/abs/2401.08503)| [![GitHub Stars](https://img.shields.io/github/stars/yerfor/MimicTalk
+)](https://github.com/yerfor/MimicTalk)  | [English Readme](./README.md)
+
+这个仓库是MimicTalk的官方PyTorch实现, 用于实现特定说话人的高表现力的虚拟人视频合成。该仓库代码基于我们先前的工作[Real3D-Portrait](https://github.com/yerfor/Real3DPortrait) (ICLR 2024)，即基于NeRF的one-shot说话人合成，这让Mimictalk的训练加速且效果增强。您可以访问我们的[项目页面](https://mimictalk.github.io/)以观看Demo视频, 阅读我们的[论文](https://arxiv.org/abs/2410.06734)以了解技术细节。
+
+<p align="center">
+    <br>
+    <img src="assets/mimictalk.png" width="100%"/>
+    <br>
+</p>
+
+# 快速上手！
+## 安装环境
+请参照[环境配置文档](docs/prepare_env/install_guide-zh.md)，配置Conda环境`mimictalk`
+## 下载预训练与第三方模型
+### 3DMM BFM模型
+下载3DMM BFM模型：[Google Drive](https://drive.google.com/drive/folders/1o4t5YIw7w4cMUN4bgU9nPf6IyWVG1bEk?usp=sharing) 或 [BaiduYun Disk](https://pan.baidu.com/s/1aqv1z_qZ23Vp2VP4uxxblQ?pwd=m9q5 ) 提取码: m9q5
+
+
+下载完成后，放置全部的文件到`deep_3drecon/BFM`里，文件结构如下：
+```
+deep_3drecon/BFM/
+├── 01_MorphableModel.mat
+├── BFM_exp_idx.mat
+├── BFM_front_idx.mat
+├── BFM_model_front.mat
+├── Exp_Pca.bin
+├── facemodel_info.mat
+├── index_mp468_from_mesh35709.npy
+├── mediapipe_in_bfm53201.npy
+└── std_exp.txt
+```
+
+### 预训练模型
+下载预训练的MimicTalk相关Checkpoints：[Google Drive](https://drive.google.com/drive/folders/1Kc6ueDO9HFDN3BhtJCEKNCZtyKHSktaA?usp=sharing) or [BaiduYun Disk](https://pan.baidu.com/s/1nQKyGV5JB6rJtda7qsThUg?pwd=mimi) 提取码: mimi
+  
+下载完成后，放置全部的文件到`checkpoints`与`checkpoints_mimictalk`里并解压，文件结构如下：
+```
+checkpoints/
+├── mimictalk_orig
+│   └── os_secc2plane_torso
+│       ├── config.yaml
+│       └── model_ckpt_steps_100000.ckpt
+|-- 240112_icl_audio2secc_vox2_cmlr
+│     ├── config.yaml
+│     └── model_ckpt_steps_1856000.ckpt
+└── pretrained_ckpts
+    └── mit_b0.pth
+
+checkpoints_mimictalk/
+└── German_20s
+    ├── config.yaml
+    └── model_ckpt_steps_10000.ckpt
+```
+
+## MimicTalk训练与推理的最简命令
+```
+python inference/train_mimictalk_on_a_video.py # train the model, this may take 10 minutes for 2,000 steps
+python inference/mimictalk_infer.py # infer the model
+```
+
+
+# 训练与推理细节
+我们目前提供了**命令行（CLI）**与**Gradio WebUI**推理方式。音频驱动推理的人像信息来自于`torso_ckpt`，因此需要至少再提供`driving audio`用于推理。另外，可以提供`style video`让模型能够预测与该视频风格一致的说话人动作。
+
+首先，切换至项目根目录并启用Conda环境：
+```bash
+cd <Real3DPortraitRoot>
+conda activate mimictalk
+export PYTHONPATH=./
+export HF_ENDPOINT=https://hf-mirror.com
+```
+
+## Gradio WebUI推理
+启动Gradio WebUI，按照提示上传素材，点击`Training`按钮进行训练；训练完成后点击`Generate`按钮即可推理：
+```bash
+python inference/app_mimictalk.py
+```
+
+## 使用Transformers Pipeline（Hugging Face Space）
+在Hugging Face Space中，我们提供了使用Transformers Pipeline的简化版本，可以直接从Hugging Face Hub下载模型文件：
+
+```bash
+python app.py --model_id mrbear1024/demo-model
+```
+
+您也可以通过环境变量设置模型来源：
+```bash
+export MODEL_ID=mrbear1024/demo-model
+python app.py
+```
+
+### 模型路径说明
+在Hugging Face Space中，模型文件默认会被下载到`/app/models`目录。您可以通过以下方式自定义模型路径：
+
+```bash
+# 使用命令行参数
+python app.py --model_dir /path/to/models
+
+# 使用环境变量
+export MODEL_DIR=/path/to/models
+python app.py
+```
+
+如果您在本地运行，模型文件会被下载到临时目录中。如果您想指定自定义路径，可以使用上述方法。
+
+## 命令行特定说话人训练
+
+需要至少提供`source video`，训练指令：
+```bash
+python inference/train_mimictalk_on_a_video.py \
+--video_id <PATH_TO_SOURCE_VIDEO> \
+--max_updates <UPDATES_NUMBER> \
+--work_dir <PATH_TO_SAVING_CKPT>
+```
+
+一些可选参数注释：
+
+- `--torso_ckpt` 预训练的Real3D-Portrait模型
+- `--max_updates` 训练更新次数
+- `--batch_size` 训练的batch size： `1` 需要约8GB显存; `2`需要约15GB��存
+- `--lr_triplane` triplane的学习率：对于视频输入,  应为0.1; 对于图片输入，应为0.001
+- `--work_dir` 未指定时，将默认存储在`checkpoints_mimictalk/`中
+
+指令示例：
+```bash
+python inference/train_mimictalk_on_a_video.py \
+--video_id data/raw/videos/German_20s.mp4 \
+--max_updates 2000 \
+--work_dir checkpoints_mimictalk/German_20s
+```
+
+## 命令行推理
+
+需要至少提供`driving audio`，可选提供`driving style`，推理指令：
+```bash
+python inference/mimictalk_infer.py \
+--drv_aud <PATH_TO_AUDIO> \
+--drv_style <PATH_TO_STYLE_VIDEO, OPTIONAL> \
+--drv_pose <PATH_TO_POSE_VIDEO, OPTIONAL> \
+--bg_img <PATH_TO_BACKGROUND_IMAGE, OPTIONAL> \
+--out_name <PATH_TO_OUTPUT_VIDEO, OPTIONAL>
+```
+
+一些可选参数注释：
+- `--drv_pose` 指定时提供了运动pose信息，不指定则为静态运动
+- `--bg_img` 指定时提供了背景信息，不指定则为source image提取的背景
+- `--mouth_amp` 嘴部张幅参数，值越大张幅越大
+- `--map_to_init_pose` 值为`True`时，首帧的pose将被映射到source pose，后续帧也作相同变换
+- `--temperature` 代表audio2motion的采样温度，值越大结果越多样，但同时精确度越低
+- `--out_name` 不指定时，结果将保存在`infer_out/tmp/`中
+- `--out_mode` 值为`final`时，只输出说话人视频；值为`concat_debug`时，同时输出一些可视化的中间结果
+
+推理命令例子：
+```bash
+python inference/mimictalk_infer.py \
+--drv_aud data/raw/examples/Obama_5s.wav \
+--drv_pose data/raw/examples/German_20s.mp4 \
+--drv_style data/raw/examples/German_20s.mp4 \
+--bg_img data/raw/examples/bg.png \
+--out_name output.mp4 \
+--out_mode final
+```
+
+# 声明
+任何组织或个人未经本人同意，不得使用本文提及的任何技术生成他人说话的视频，包括但不限于政府领导人、政界人士、社会名流等。如不遵守本条款，则可能违反版权法。
+
+# 引用我们
+如果这个仓库对你有帮助，请考虑引用我们的工作：
+```
+@inproceedings{ye2024mimicktalk,
+    author    = {Ye, Zhenhui and Zhong, Tianyun and Ren, Yi and Yang, Jiaqi and Li, Weichuang and Huang, Jiangwei and Jiang, Ziyue and He, Jinzheng and Huang, Rongjie and Liu, Jinglin and Zhang, Chen and Yin, Xiang and Ma, Zejun and Zhao, Zhou},
+    title     = {MimicTalk: Mimicking a personalized and expressive 3D talking face in few minutes},
+    journal   = {NeurIPS},
+    year      = {2024},
+}
+@inproceedings{ye2024real3d,
+  title = {Real3D-Portrait: One-shot Realistic 3D Talking Portrait Synthesis},
+  author = {Ye, Zhenhui and Zhong, Tianyun and Ren, Yi and Yang, Jiaqi and Li, Weichuang and Huang, Jiawei and Jiang, Ziyue and He, Jinzheng and Huang, Rongjie and Liu, Jinglin and others},
+  journal  = {ICLR},
+  year={2024}
+}
+```

inference/app_mimictalk.py

@@ -0,0 +1,365 @@
+import os, sys
+sys.path.append('./')
+import argparse
+import traceback
+import gradio as gr
+# from inference.real3d_infer import GeneFace2Infer
+from inference.mimictalk_infer import AdaptGeneFace2Infer
+from inference.train_mimictalk_on_a_video import LoRATrainer
+from utils.commons.hparams import hparams
+
+class Trainer():
+    def __init__(self):
+        pass
+    
+    def train_once_args(self, *args, **kargs):
+        assert len(kargs) == 0
+        args = [
+            '', # head_ckpt
+            'checkpoints/mimictalk_orig/os_secc2plane_torso/', # torso_ckpt
+            args[0], 
+            '',
+            10000,
+            1,
+            False,
+            0.001,
+            0.005,
+            0.2,
+            'secc2plane_sr',
+            2,
+        ]
+        keys = [
+            'head_ckpt',
+            'torso_ckpt',
+            'video_id',
+            'work_dir',
+            'max_updates',
+            'batch_size',
+            'test',
+            'lr',
+            'lr_triplane',
+            'lambda_lpips',
+            'lora_mode',
+            'lora_r',
+        ]
+        inp = {}
+        info = ""
+        
+        try: # try to catch errors and jump to return 
+            for key_index in range(len(keys)):
+                key = keys[key_index]
+                inp[key] = args[key_index]
+                if '_name' in key:
+                    inp[key] = inp[key] if inp[key] is not None else ''
+            
+            if inp['video_id'] == '':
+                info = "Input Error: Source video is REQUIRED!"
+                raise ValueError  
+
+            inp['out_name'] = ''
+            inp['seed'] = 42
+            
+            if inp['work_dir'] is None or inp['work_dir'] == '':
+                video_id = os.path.basename(inp['video_id'])[:-4] if inp['video_id'].endswith((".mp4", ".png", ".jpg", ".jpeg")) else inp['video_id']
+                inp['work_dir'] = f'checkpoints_mimictalk/{video_id}'
+            try:
+                trainer = LoRATrainer(inp)
+                trainer.training_loop(inp)
+            except Exception as e:
+                content = f"{e}"
+                info = f"Training ERROR: {content}"
+                traceback.print_exc()
+                raise ValueError
+        except Exception as e:
+            if info == "": # unexpected errors
+                content = f"{e}"
+                info = f"WebUI ERROR: {content}"
+                traceback.print_exc()
+
+        
+        # output part
+        if len(info) > 0 : # there is errors    
+            print(info)
+            info_gr = gr.update(visible=True, value=info)
+        else: # no errors
+            info_gr = gr.update(visible=False, value=info)
+
+        torso_model_dir = gr.FileExplorer(glob="checkpoints_mimictalk/**/*.ckpt", value=inp['work_dir'], file_count='single', label='mimictalk model ckpt path or directory')
+
+            
+        return info_gr, torso_model_dir
+
+class Inferer(AdaptGeneFace2Infer):
+    def infer_once_args(self, *args, **kargs):
+        assert len(kargs) == 0
+        keys = [
+            # 'src_image_name',
+            'drv_audio_name',
+            'drv_pose_name',
+            'drv_talking_style_name',
+            'bg_image_name',
+            'blink_mode',
+            'temperature',
+            'cfg_scale',
+            'out_mode',
+            'map_to_init_pose',
+            'low_memory_usage',
+            'hold_eye_opened',
+            'a2m_ckpt',
+            # 'head_ckpt',
+            'torso_ckpt',
+            'min_face_area_percent',
+        ]
+        inp = {}
+        out_name = None
+        info = ""
+        
+        try: # try to catch errors and jump to return 
+            for key_index in range(len(keys)):
+                key = keys[key_index]
+                inp[key] = args[key_index]
+                if '_name' in key:
+                    inp[key] = inp[key] if inp[key] is not None else ''
+            
+            inp['head_ckpt'] = ''
+            
+            # if inp['src_image_name'] == '':
+            #     info = "Input Error: Source image is REQUIRED!"
+            #     raise ValueError
+            if inp['drv_audio_name'] == '' and inp['drv_pose_name'] == '':
+                info = "Input Error: At least one of driving audio or video is REQUIRED!"
+                raise ValueError
+
+
+            if inp['drv_audio_name'] == '' and inp['drv_pose_name'] != '':
+                inp['drv_audio_name'] = inp['drv_pose_name']
+                print("No audio input, we use driving pose video for video driving")
+                
+            if inp['drv_pose_name'] == '':
+                inp['drv_pose_name'] = 'static'    
+            
+            reload_flag = False
+            if inp['a2m_ckpt'] != self.audio2secc_dir:
+                print("Changes of a2m_ckpt detected, reloading model")
+                reload_flag = True
+            if inp['head_ckpt'] != self.head_model_dir:
+                print("Changes of head_ckpt detected, reloading model")
+                reload_flag = True
+            if inp['torso_ckpt'] != self.torso_model_dir:
+                print("Changes of torso_ckpt detected, reloading model")
+                reload_flag = True
+
+            inp['out_name'] = ''
+            inp['seed'] = 42
+            inp['denoising_steps'] = 20
+            print(f"infer inputs : {inp}")
+                
+            try:
+                if reload_flag:
+                    self.__init__(inp['a2m_ckpt'], inp['head_ckpt'], inp['torso_ckpt'], inp=inp, device=self.device)
+            except Exception as e:
+                content = f"{e}"
+                info = f"Reload ERROR: {content}"
+                traceback.print_exc()
+                raise ValueError
+            try:
+                out_name = self.infer_once(inp)
+            except Exception as e:
+                content = f"{e}"
+                info = f"Inference ERROR: {content}"
+                traceback.print_exc()
+                raise ValueError
+        except Exception as e:
+            if info == "": # unexpected errors
+                content = f"{e}"
+                info = f"WebUI ERROR: {content}"
+                traceback.print_exc()
+
+        # output part
+        if len(info) > 0 : # there is errors    
+            print(info)
+            info_gr = gr.update(visible=True, value=info)
+        else: # no errors
+            info_gr = gr.update(visible=False, value=info)
+        if out_name is not None and len(out_name) > 0 and os.path.exists(out_name): # good output
+            print(f"Succefully generated in {out_name}")
+            video_gr = gr.update(visible=True, value=out_name)
+        else:
+            print(out_name)
+            print(os.path.exists(out_name))
+            print(f"Failed to generate")
+            video_gr = gr.update(visible=True, value=out_name)
+            
+        return video_gr, info_gr
+
+def toggle_audio_file(choice):
+    if choice == False:
+        return gr.update(visible=True), gr.update(visible=False)
+    else:
+        return gr.update(visible=False), gr.update(visible=True)
+    
+def ref_video_fn(path_of_ref_video):
+    if path_of_ref_video is not None:
+        return gr.update(value=True)
+    else:
+        return gr.update(value=False)
+
+def mimictalk_demo(
+    audio2secc_dir,
+    head_model_dir,
+    torso_model_dir, 
+    device          = 'cuda',
+    warpfn          = None,
+    ):
+
+    sep_line = "-" * 40
+
+    infer_obj = Inferer(
+        audio2secc_dir=audio2secc_dir, 
+        head_model_dir=head_model_dir,
+        torso_model_dir=torso_model_dir,
+        device=device,
+    )
+    train_obj = Trainer()
+
+    print(sep_line)
+    print("Model loading is finished.")
+    print(sep_line)
+    with gr.Blocks(analytics_enabled=False) as real3dportrait_interface:
+        # gr.Markdown("\
+        #     <div align='center'> <h2> MimicTalk: Mimicking a personalized and expressive 3D talking face in minutes (NIPS 2024) </span> </h2> \
+        #     <a style='font-size:18px;color: #a0a0a0' href=''>Arxiv</a> &nbsp;&nbsp;&nbsp;&nbsp;&nbsp; \
+        #     <a style='font-size:18px;color: #a0a0a0' href='https://mimictalk.github.io/'>Homepage</a>  &nbsp;&nbsp;&nbsp;&nbsp;&nbsp; \
+        #     <a style='font-size:18px;color: #a0a0a0' href='https://github.com/yerfor/MimicTalk/'> Github </div>")
+        
+        gr.Markdown("\
+            <div align='center'> <h2> MimicTalk: Mimicking a personalized and expressive 3D talking face in minutes (NIPS 2024) </span> </h2> \
+            <a style='font-size:18px;color: #a0a0a0' href='https://mimictalk.github.io/'>Homepage</a>  &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;")
+        
+        sources = None
+        with gr.Row():
+            with gr.Column(variant='panel'):
+                with gr.Tabs(elem_id="source_image"):
+                    with gr.TabItem('Upload Training Video'):
+                        with gr.Row():
+                            src_video_name = gr.Video(label="Source video (required for training)", sources=sources, value="data/raw/videos/German_20s.mp4")
+                            # src_video_name = gr.Image(label="Source video (required for training)", sources=sources, type="filepath", value="data/raw/videos/German_20s.mp4")
+                with gr.Tabs(elem_id="driven_audio"):
+                    with gr.TabItem('Upload Driving Audio'):
+                        with gr.Column(variant='panel'):
+                            drv_audio_name = gr.Audio(label="Input audio (required for inference)", sources=sources, type="filepath", value="data/raw/examples/80_vs_60_10s.wav")
+                with gr.Tabs(elem_id="driven_style"):
+                    with gr.TabItem('Upload Style Prompt'):
+                        with gr.Column(variant='panel'):
+                            drv_style_name = gr.Video(label="Driven Style (optional for inference)", sources=sources, value="data/raw/videos/German_20s.mp4")
+                with gr.Tabs(elem_id="driven_pose"):
+                    with gr.TabItem('Upload Driving Pose'):
+                        with gr.Column(variant='panel'):
+                            drv_pose_name = gr.Video(label="Driven Pose (optional for inference)", sources=sources, value="data/raw/videos/German_20s.mp4")
+                with gr.Tabs(elem_id="bg_image"):
+                    with gr.TabItem('Upload Background Image'):
+                        with gr.Row():
+                            bg_image_name = gr.Image(label="Background image (optional for inference)", sources=sources, type="filepath", value=None)
+
+                             
+            with gr.Column(variant='panel'): 
+                with gr.Tabs(elem_id="checkbox"):
+                    with gr.TabItem('General Settings'):
+                        with gr.Column(variant='panel'):
+
+                            blink_mode = gr.Radio(['none', 'period'], value='period', label='blink mode', info="whether to blink periodly") #        
+                            min_face_area_percent = gr.Slider(minimum=0.15, maximum=0.5, step=0.01, label="min_face_area_percent",  value=0.2, info='The minimum face area percent in the output frame, to prevent bad cases caused by a too small face.',)
+                            temperature = gr.Slider(minimum=0.0, maximum=1.0, step=0.025, label="temperature",  value=0.2, info='audio to secc temperature',)
+                            cfg_scale = gr.Slider(minimum=1.0, maximum=3.0, step=0.025, label="talking style cfg_scale",  value=1.5, info='higher -> encourage the generated motion more coherent to talking style',)
+                            out_mode = gr.Radio(['final', 'concat_debug'], value='concat_debug', label='output layout', info="final: only final output ; concat_debug: final output concated with internel features") 
+                            low_memory_usage = gr.Checkbox(label="Low Memory Usage Mode: save memory at the expense of lower inference speed. Useful when running a low audio (minutes-long).", value=False)
+                            map_to_init_pose = gr.Checkbox(label="Whether to map pose of first frame to initial pose", value=True)
+                            hold_eye_opened  = gr.Checkbox(label="Whether to maintain eyes always open")
+                          
+                            train_submit = gr.Button('Train', elem_id="train", variant='primary')      
+                            infer_submit = gr.Button('Generate', elem_id="generate", variant='primary')
+                        
+                    with gr.Tabs(elem_id="genearted_video"):
+                        info_box = gr.Textbox(label="Error", interactive=False, visible=False)
+                        gen_video = gr.Video(label="Generated video", format="mp4", visible=True)
+            with gr.Column(variant='panel'): 
+                with gr.Tabs(elem_id="checkbox"):
+                    with gr.TabItem('Checkpoints'):
+                        with gr.Column(variant='panel'):
+                            ckpt_info_box = gr.Textbox(value="Please select \"ckpt\" under the checkpoint folder ", interactive=False, visible=True, show_label=False)
+                            audio2secc_dir = gr.FileExplorer(glob="checkpoints/**/*.ckpt", value=audio2secc_dir, file_count='single', label='audio2secc model ckpt path or directory')
+                            # head_model_dir = gr.FileExplorer(glob="checkpoints/**/*.ckpt", value=head_model_dir, file_count='single', label='head model ckpt path or directory (will be ignored if torso model is set)')
+                            torso_model_dir = gr.FileExplorer(glob="checkpoints_mimictalk/**/*.ckpt", value=torso_model_dir, file_count='single', label='mimictalk model ckpt path or directory')
+                            # audio2secc_dir = gr.Textbox(audio2secc_dir, max_lines=1, label='audio2secc model ckpt path or directory (will be ignored if torso model is set)')
+                            # head_model_dir = gr.Textbox(head_model_dir, max_lines=1, label='head model ckpt path or directory (will be ignored if torso model is set)')
+                            # torso_model_dir = gr.Textbox(torso_model_dir, max_lines=1, label='torso model ckpt path or directory')
+
+
+        fn = infer_obj.infer_once_args
+        if warpfn:
+            fn = warpfn(fn)
+        infer_submit.click(
+                    fn=fn, 
+                    inputs=[
+                        drv_audio_name,
+                        drv_pose_name,
+                        drv_style_name,
+                        bg_image_name,
+                        blink_mode,
+                        temperature,
+                        cfg_scale,
+                        out_mode,
+                        map_to_init_pose,
+                        low_memory_usage,
+                        hold_eye_opened,
+                        audio2secc_dir,
+                        # head_model_dir,
+                        torso_model_dir,
+                        min_face_area_percent,
+                    ], 
+                    outputs=[
+                        gen_video,
+                        info_box,
+                    ],
+                    )
+        
+        fn_train = train_obj.train_once_args
+
+        train_submit.click(
+            fn=fn_train, 
+            inputs=[
+                src_video_name, 
+
+            ], 
+            outputs=[
+                # gen_video,
+                info_box,
+                torso_model_dir,
+            ],
+        )
+
+    print(sep_line)
+    print("Gradio page is constructed.")
+    print(sep_line)
+
+    return real3dportrait_interface
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--a2m_ckpt", default='checkpoints/240112_icl_audio2secc_vox2_cmlr') # checkpoints/0727_audio2secc/audio2secc_withlm2d100_randomframe
+    parser.add_argument("--head_ckpt", default='') # checkpoints/0729_th1kh/secc_img2plane checkpoints/0720_img2planes/secc_img2plane_two_stage
+    parser.add_argument("--torso_ckpt", default='checkpoints_mimictalk/German_20s/model_ckpt_steps_10000.ckpt') 
+    parser.add_argument("--port", type=int, default=None) 
+    parser.add_argument("--server", type=str, default='127.0.0.1')
+    parser.add_argument("--share", action='store_true', dest='share', help='share srever to Internet')
+
+    args = parser.parse_args()
+    demo = mimictalk_demo(
+        audio2secc_dir=args.a2m_ckpt,
+        head_model_dir=args.head_ckpt,
+        torso_model_dir=args.torso_ckpt,
+        device='cuda:0',
+        warpfn=None,
+    )
+    demo.queue()
+    demo.launch(share=args.share, server_name=args.server, server_port=args.port)

inference/app_real3dportrait.py

@@ -0,0 +1,247 @@
+import os, sys
+sys.path.append('./')
+import argparse
+import gradio as gr
+from inference.real3d_infer import GeneFace2Infer
+from utils.commons.hparams import hparams
+
+class Inferer(GeneFace2Infer):
+    def infer_once_args(self, *args, **kargs):
+        assert len(kargs) == 0
+        keys = [
+            'src_image_name',
+            'drv_audio_name',
+            'drv_pose_name',
+            'bg_image_name',
+            'blink_mode',
+            'temperature',
+            'mouth_amp',
+            'out_mode',
+            'map_to_init_pose',
+            'low_memory_usage',
+            'hold_eye_opened',
+            'a2m_ckpt',
+            'head_ckpt',
+            'torso_ckpt',
+            'min_face_area_percent',
+        ]
+        inp = {}
+        out_name = None
+        info = ""
+        
+        try: # try to catch errors and jump to return 
+            for key_index in range(len(keys)):
+                key = keys[key_index]
+                inp[key] = args[key_index]
+                if '_name' in key:
+                    inp[key] = inp[key] if inp[key] is not None else ''
+            
+            if inp['src_image_name'] == '':
+                info = "Input Error: Source image is REQUIRED!"
+                raise ValueError
+            if inp['drv_audio_name'] == '' and inp['drv_pose_name'] == '':
+                info = "Input Error: At least one of driving audio or video is REQUIRED!"
+                raise ValueError
+
+
+            if inp['drv_audio_name'] == '' and inp['drv_pose_name'] != '':
+                inp['drv_audio_name'] = inp['drv_pose_name']
+                print("No audio input, we use driving pose video for video driving")
+                
+            if inp['drv_pose_name'] == '':
+                inp['drv_pose_name'] = 'static'    
+            
+            reload_flag = False
+            if inp['a2m_ckpt'] != self.audio2secc_dir:
+                print("Changes of a2m_ckpt detected, reloading model")
+                reload_flag = True
+            if inp['head_ckpt'] != self.head_model_dir:
+                print("Changes of head_ckpt detected, reloading model")
+                reload_flag = True
+            if inp['torso_ckpt'] != self.torso_model_dir:
+                print("Changes of torso_ckpt detected, reloading model")
+                reload_flag = True
+
+            inp['out_name'] = ''
+            inp['seed'] = 42
+            
+            print(f"infer inputs : {inp}")
+                
+            try:
+                if reload_flag:
+                    self.__init__(inp['a2m_ckpt'], inp['head_ckpt'], inp['torso_ckpt'], inp=inp, device=self.device)
+            except Exception as e:
+                content = f"{e}"
+                info = f"Reload ERROR: {content}"
+                raise ValueError
+            try:
+                out_name = self.infer_once(inp)
+            except Exception as e:
+                content = f"{e}"
+                info = f"Inference ERROR: {content}"
+                raise ValueError
+        except Exception as e:
+            if info == "": # unexpected errors
+                content = f"{e}"
+                info = f"WebUI ERROR: {content}"
+        
+        # output part
+        if len(info) > 0 : # there is errors    
+            print(info)
+            info_gr = gr.update(visible=True, value=info)
+        else: # no errors
+            info_gr = gr.update(visible=False, value=info)
+        if out_name is not None and len(out_name) > 0 and os.path.exists(out_name): # good output
+            print(f"Succefully generated in {out_name}")
+            video_gr = gr.update(visible=True, value=out_name)
+        else:
+            print(f"Failed to generate")
+            video_gr = gr.update(visible=True, value=out_name)
+            
+        return video_gr, info_gr
+
+def toggle_audio_file(choice):
+    if choice == False:
+        return gr.update(visible=True), gr.update(visible=False)
+    else:
+        return gr.update(visible=False), gr.update(visible=True)
+    
+def ref_video_fn(path_of_ref_video):
+    if path_of_ref_video is not None:
+        return gr.update(value=True)
+    else:
+        return gr.update(value=False)
+
+def real3dportrait_demo(
+    audio2secc_dir,
+    head_model_dir,
+    torso_model_dir, 
+    device          = 'cuda',
+    warpfn          = None,
+    ):
+
+    sep_line = "-" * 40
+
+    infer_obj = Inferer(
+        audio2secc_dir=audio2secc_dir, 
+        head_model_dir=head_model_dir,
+        torso_model_dir=torso_model_dir,
+        device=device,
+    )
+
+    print(sep_line)
+    print("Model loading is finished.")
+    print(sep_line)
+    with gr.Blocks(analytics_enabled=False) as real3dportrait_interface:
+        gr.Markdown("\
+            <div align='center'> <h2> Real3D-Portrait: One-shot Realistic 3D Talking Portrait Synthesis (ICLR 2024 Spotlight) </span> </h2> \
+            <a style='font-size:18px;color: #a0a0a0' href='https://arxiv.org/pdf/2401.08503.pdf'>Arxiv</a> &nbsp;&nbsp;&nbsp;&nbsp;&nbsp; \
+            <a style='font-size:18px;color: #a0a0a0' href='https://real3dportrait.github.io/'>Homepage</a>  &nbsp;&nbsp;&nbsp;&nbsp;&nbsp; \
+            <a style='font-size:18px;color: #a0a0a0' href='https://github.com/yerfor/Real3DPortrait/'> Github </div>")
+        
+        sources = None
+        with gr.Row():
+            with gr.Column(variant='panel'):
+                with gr.Tabs(elem_id="source_image"):
+                    with gr.TabItem('Upload image'):
+                        with gr.Row():
+                            src_image_name = gr.Image(label="Source image (required)", sources=sources, type="filepath", value="data/raw/examples/Macron.png")
+                with gr.Tabs(elem_id="driven_audio"):
+                    with gr.TabItem('Upload audio'):
+                        with gr.Column(variant='panel'):
+                            drv_audio_name = gr.Audio(label="Input audio (required for audio-driven)", sources=sources, type="filepath", value="data/raw/examples/Obama_5s.wav")
+                with gr.Tabs(elem_id="driven_pose"):
+                    with gr.TabItem('Upload video'):
+                        with gr.Column(variant='panel'):
+                            drv_pose_name = gr.Video(label="Driven Pose (required for video-driven, optional for audio-driven)", sources=sources, value="data/raw/examples/May_5s.mp4")
+                with gr.Tabs(elem_id="bg_image"):
+                    with gr.TabItem('Upload image'):
+                        with gr.Row():
+                            bg_image_name = gr.Image(label="Background image (optional)", sources=sources, type="filepath", value="data/raw/examples/bg.png")
+
+                             
+            with gr.Column(variant='panel'): 
+                with gr.Tabs(elem_id="checkbox"):
+                    with gr.TabItem('General Settings'):
+                        with gr.Column(variant='panel'):
+
+                            blink_mode = gr.Radio(['none', 'period'], value='period', label='blink mode', info="whether to blink periodly") #        
+                            min_face_area_percent = gr.Slider(minimum=0.15, maximum=0.5, step=0.01, label="min_face_area_percent",  value=0.2, info='The minimum face area percent in the output frame, to prevent bad cases caused by a too small face.',)
+                            temperature = gr.Slider(minimum=0.0, maximum=1.0, step=0.025, label="temperature",  value=0.2, info='audio to secc temperature',)
+                            mouth_amp = gr.Slider(minimum=0.0, maximum=1.0, step=0.025, label="mouth amplitude",  value=0.45, info='higher -> mouth will open wider, default to be 0.4',)
+                            out_mode = gr.Radio(['final', 'concat_debug'], value='concat_debug', label='output layout', info="final: only final output ; concat_debug: final output concated with internel features") 
+                            low_memory_usage = gr.Checkbox(label="Low Memory Usage Mode: save memory at the expense of lower inference speed. Useful when running a low audio (minutes-long).", value=False)
+                            map_to_init_pose = gr.Checkbox(label="Whether to map pose of first frame to initial pose", value=True)
+                            hold_eye_opened  = gr.Checkbox(label="Whether to maintain eyes always open")
+                                
+                            submit = gr.Button('Generate', elem_id="generate", variant='primary')
+                        
+                    with gr.Tabs(elem_id="genearted_video"):
+                        info_box = gr.Textbox(label="Error", interactive=False, visible=False)
+                        gen_video = gr.Video(label="Generated video", format="mp4", visible=True)
+            with gr.Column(variant='panel'): 
+                with gr.Tabs(elem_id="checkbox"):
+                    with gr.TabItem('Checkpoints'):
+                        with gr.Column(variant='panel'):
+                            ckpt_info_box = gr.Textbox(value="Please select \"ckpt\" under the checkpoint folder ", interactive=False, visible=True, show_label=False)
+                            audio2secc_dir = gr.FileExplorer(glob="checkpoints/**/*.ckpt", value=audio2secc_dir, file_count='single', label='audio2secc model ckpt path or directory')
+                            head_model_dir = gr.FileExplorer(glob="checkpoints/**/*.ckpt", value=head_model_dir, file_count='single', label='head model ckpt path or directory (will be ignored if torso model is set)')
+                            torso_model_dir = gr.FileExplorer(glob="checkpoints/**/*.ckpt", value=torso_model_dir, file_count='single', label='torso model ckpt path or directory')
+                            # audio2secc_dir = gr.Textbox(audio2secc_dir, max_lines=1, label='audio2secc model ckpt path or directory (will be ignored if torso model is set)')
+                            # head_model_dir = gr.Textbox(head_model_dir, max_lines=1, label='head model ckpt path or directory (will be ignored if torso model is set)')
+                            # torso_model_dir = gr.Textbox(torso_model_dir, max_lines=1, label='torso model ckpt path or directory')
+
+
+        fn = infer_obj.infer_once_args
+        if warpfn:
+            fn = warpfn(fn)
+        submit.click(
+                    fn=fn, 
+                    inputs=[
+                        src_image_name, 
+                        drv_audio_name,
+                        drv_pose_name,
+                        bg_image_name,
+                        blink_mode,
+                        temperature,
+                        mouth_amp,
+                        out_mode,
+                        map_to_init_pose,
+                        low_memory_usage,
+                        hold_eye_opened,
+                        audio2secc_dir,
+                        head_model_dir,
+                        torso_model_dir,
+                        min_face_area_percent,
+                    ], 
+                    outputs=[
+                        gen_video,
+                        info_box,
+                    ],
+                    )
+
+    print(sep_line)
+    print("Gradio page is constructed.")
+    print(sep_line)
+
+    return real3dportrait_interface
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--a2m_ckpt", type=str, default='checkpoints/240210_real3dportrait_orig/audio2secc_vae/model_ckpt_steps_400000.ckpt')
+    parser.add_argument("--head_ckpt", type=str, default='')
+    parser.add_argument("--torso_ckpt", type=str, default='checkpoints/240210_real3dportrait_orig/secc2plane_torso_orig/model_ckpt_steps_100000.ckpt') 
+    parser.add_argument("--port", type=int, default=None) 
+    parser.add_argument("--server", type=str, default='127.0.0.1')
+    parser.add_argument("--share", action='store_true', dest='share', help='share srever to Internet')
+
+    args = parser.parse_args()
+    demo = real3dportrait_demo(
+        audio2secc_dir=args.a2m_ckpt,
+        head_model_dir=args.head_ckpt,
+        torso_model_dir=args.torso_ckpt,
+        device='cuda:0',
+        warpfn=None,
+    )
+    demo.queue()
+    demo.launch(share=args.share, server_name=args.server, server_port=args.port)

inference/edit_secc.py

@@ -0,0 +1,147 @@
+import cv2
+import torch
+from utils.commons.image_utils import dilate, erode
+from sklearn.neighbors import NearestNeighbors
+import copy
+import numpy as np
+from utils.commons.meters import Timer
+
+def hold_eye_opened_for_secc(img):
+    img = img.permute(1,2,0).cpu().numpy()
+    img = ((img +1)/2*255).astype(np.uint)
+    face_mask = (img[...,0] != 0) & (img[...,1] != 0) & (img[...,2] != 0)
+    face_xys = np.stack(np.nonzero(face_mask)).transpose(1, 0) # [N_nonbg,2] coordinate of non-face pixels
+    h,w = face_mask.shape
+    # get face and eye mask
+    left_eye_prior_reigon = np.zeros([h,w], dtype=bool)
+    right_eye_prior_reigon = np.zeros([h,w], dtype=bool)
+    left_eye_prior_reigon[h//4:h//2, w//4:w//2] = True
+    right_eye_prior_reigon[h//4:h//2, w//2:w//4*3] = True
+    eye_prior_reigon = left_eye_prior_reigon | right_eye_prior_reigon
+    coarse_eye_mask = (~ face_mask) & eye_prior_reigon
+    coarse_eye_xys = np.stack(np.nonzero(coarse_eye_mask)).transpose(1, 0) # [N_nonbg,2] coordinate of non-face pixels
+
+    opened_eye_mask = cv2.imread('inference/os_avatar/opened_eye_mask.png')
+    opened_eye_mask = torch.nn.functional.interpolate(torch.tensor(opened_eye_mask).permute(2,0,1).unsqueeze(0), size=(img.shape[0], img.shape[1]), mode='nearest')[0].permute(1,2,0).sum(-1).bool().cpu() # [512,512,3]
+    coarse_opened_eye_xys = np.stack(np.nonzero(opened_eye_mask)) # [N_nonbg,2] coordinate of non-face pixels
+    
+    nbrs = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(coarse_eye_xys)
+    dists, _ = nbrs.kneighbors(coarse_opened_eye_xys) # [512*512, 1] distance to nearest non-bg pixel
+    # print(dists.max())
+    non_opened_eye_pixs = dists > max(dists.max()*0.75, 4) # 大于这个距离的opened eye部分会被合上
+    non_opened_eye_pixs = non_opened_eye_pixs.reshape([-1])
+    opened_eye_xys_to_erode = coarse_opened_eye_xys[non_opened_eye_pixs]
+    opened_eye_mask[opened_eye_xys_to_erode[...,0], opened_eye_xys_to_erode[...,1]] = False # shrink 将mask在face-eye边界收缩3pixel，为了平滑
+
+    img[opened_eye_mask] = 0
+    return torch.tensor(img.astype(np.float32) / 127.5 - 1).permute(2,0,1)
+    
+
+# def hold_eye_opened_for_secc(img):
+#     img = copy.copy(img)
+#     eye_mask = cv2.imread('inference/os_avatar/opened_eye_mask.png')
+#     eye_mask = torch.nn.functional.interpolate(torch.tensor(eye_mask).permute(2,0,1).unsqueeze(0), size=(img.shape[-2], img.shape[-1]), mode='nearest')[0].bool().to(img.device) # [3,512,512]
+#     img[eye_mask] = -1
+#     return img
+    
+def blink_eye_for_secc(img, close_eye_percent=0.5):
+    """
+    secc_img: [3,h,w], tensor, -1~1
+    """
+    img = img.permute(1,2,0).cpu().numpy()
+    img = ((img +1)/2*255).astype(np.uint)
+    assert close_eye_percent <= 1.0 and close_eye_percent >= 0.
+    if close_eye_percent == 0: return torch.tensor(img.astype(np.float32) / 127.5 - 1).permute(2,0,1)
+    img = copy.deepcopy(img)
+    face_mask = (img[...,0] != 0) & (img[...,1] != 0) & (img[...,2] != 0)
+    h,w = face_mask.shape
+
+    # get face and eye mask
+    left_eye_prior_reigon = np.zeros([h,w], dtype=bool)
+    right_eye_prior_reigon = np.zeros([h,w], dtype=bool)
+    left_eye_prior_reigon[h//4:h//2, w//4:w//2] = True
+    right_eye_prior_reigon[h//4:h//2, w//2:w//4*3] = True
+    eye_prior_reigon = left_eye_prior_reigon | right_eye_prior_reigon
+    coarse_eye_mask = (~ face_mask) & eye_prior_reigon
+    coarse_left_eye_mask = (~ face_mask) & left_eye_prior_reigon
+    coarse_right_eye_mask = (~ face_mask) & right_eye_prior_reigon
+    coarse_eye_xys = np.stack(np.nonzero(coarse_eye_mask)).transpose(1, 0) # [N_nonbg,2] coordinate of non-face pixels
+    min_h = coarse_eye_xys[:, 0].min()
+    max_h = coarse_eye_xys[:, 0].max()
+    coarse_left_eye_xys = np.stack(np.nonzero(coarse_left_eye_mask)).transpose(1, 0) # [N_nonbg,2] coordinate of non-face pixels
+    left_min_w = coarse_left_eye_xys[:, 1].min()
+    left_max_w = coarse_left_eye_xys[:, 1].max()
+    coarse_right_eye_xys = np.stack(np.nonzero(coarse_right_eye_mask)).transpose(1, 0) # [N_nonbg,2] coordinate of non-face pixels
+    right_min_w = coarse_right_eye_xys[:, 1].min()
+    right_max_w = coarse_right_eye_xys[:, 1].max()
+
+    # 尽力较少需要考虑的face_xyz,以降低KNN的损耗
+    left_eye_prior_reigon = np.zeros([h,w], dtype=bool)
+    more_room = 4 # 过小会导致一些问题
+    left_eye_prior_reigon[min_h-more_room:max_h+more_room, left_min_w-more_room:left_max_w+more_room] = True
+    right_eye_prior_reigon = np.zeros([h,w], dtype=bool)
+    right_eye_prior_reigon[min_h-more_room:max_h+more_room, right_min_w-more_room:right_max_w+more_room] = True
+    eye_prior_reigon = left_eye_prior_reigon | right_eye_prior_reigon
+
+    around_eye_face_mask = face_mask & eye_prior_reigon
+    face_mask = around_eye_face_mask
+    face_xys = np.stack(np.nonzero(around_eye_face_mask)).transpose(1, 0) # [N_nonbg,2] coordinate of non-face pixels
+
+    nbrs = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(coarse_eye_xys)
+    dists, _ = nbrs.kneighbors(face_xys) # [512*512, 1] distance to nearest non-bg pixel
+    face_pixs = dists > 5 # 只有距离最近的eye pixel大于5的才被认为是face，过小会导致一些问题
+    face_pixs = face_pixs.reshape([-1])
+    face_xys_to_erode = face_xys[~face_pixs]
+    face_mask[face_xys_to_erode[...,0], face_xys_to_erode[...,1]] = False # shrink 将mask在face-eye边界收缩3pixel，为了平滑
+    eye_mask = (~ face_mask) & eye_prior_reigon
+
+    h_grid = np.mgrid[0:h, 0:w][0]
+    eye_num_pixel_along_w_axis = eye_mask.sum(axis=0)
+    eye_mask_along_w_axis = eye_num_pixel_along_w_axis != 0
+
+    tmp_h_grid = h_grid.copy()
+    tmp_h_grid[~eye_mask] = 0
+    eye_mean_h_coord_along_w_axis = tmp_h_grid.sum(axis=0) / np.clip(eye_num_pixel_along_w_axis, a_min=1, a_max=h)
+    tmp_h_grid = h_grid.copy()
+    tmp_h_grid[~eye_mask] = 99999
+    eye_min_h_coord_along_w_axis = tmp_h_grid.min(axis=0)
+    tmp_h_grid = h_grid.copy()
+    tmp_h_grid[~eye_mask] = -99999
+    eye_max_h_coord_along_w_axis = tmp_h_grid.max(axis=0)
+
+    eye_low_h_coord_along_w_axis = close_eye_percent * eye_mean_h_coord_along_w_axis + (1-close_eye_percent) * eye_min_h_coord_along_w_axis # upper eye 
+    eye_high_h_coord_along_w_axis = close_eye_percent * eye_mean_h_coord_along_w_axis + (1-close_eye_percent) * eye_max_h_coord_along_w_axis # lower eye 
+
+    tmp_h_grid = h_grid.copy()
+    tmp_h_grid[~eye_mask] = 99999
+    upper_eye_blink_mask = tmp_h_grid <= eye_low_h_coord_along_w_axis
+    tmp_h_grid = h_grid.copy()
+    tmp_h_grid[~eye_mask] = -99999
+    lower_eye_blink_mask = tmp_h_grid >= eye_high_h_coord_along_w_axis
+    eye_blink_mask = upper_eye_blink_mask | lower_eye_blink_mask
+
+    face_xys = np.stack(np.nonzero(around_eye_face_mask)).transpose(1, 0) # [N_nonbg,2] coordinate of non-face pixels
+    eye_blink_xys = np.stack(np.nonzero(eye_blink_mask)).transpose(1, 0) # [N_nonbg,hw] coordinate of non-face pixels
+    nbrs = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(face_xys)
+    distances, indices = nbrs.kneighbors(eye_blink_xys)
+    bg_fg_xys = face_xys[indices[:, 0]]
+    img[eye_blink_xys[:, 0], eye_blink_xys[:, 1], :] = img[bg_fg_xys[:, 0], bg_fg_xys[:, 1], :]
+    return torch.tensor(img.astype(np.float32) / 127.5 - 1).permute(2,0,1)
+
+
+if __name__ == '__main__':
+    import imageio
+    import tqdm
+    img = cv2.imread("assets/cano_secc.png")
+    img = img / 127.5 - 1
+    img = torch.FloatTensor(img).permute(2, 0, 1)
+    fps = 25
+    writer = imageio.get_writer('demo_blink.mp4', fps=fps)
+
+    for i in tqdm.trange(33):
+        blink_percent = 0.03 * i
+        with Timer("Blink", True):
+            out_img = blink_eye_for_secc(img, blink_percent)
+        out_img = ((out_img.permute(1,2,0)+1)*127.5).int().numpy()
+        writer.append_data(out_img)
+    writer.close()

inference/infer_utils.py

@@ -0,0 +1,154 @@
+import os
+import torch
+import torch.nn.functional as F
+import librosa
+import numpy as np
+import importlib
+import tqdm
+import copy
+import cv2
+from scipy.spatial.transform import Rotation
+
+
+def load_img_to_512_hwc_array(img_name):
+    img = cv2.imread(img_name)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img = cv2.resize(img, (512, 512))
+    return img
+
+def load_img_to_normalized_512_bchw_tensor(img_name):
+    img = load_img_to_512_hwc_array(img_name)
+    img = ((torch.tensor(img) - 127.5)/127.5).float().unsqueeze(0).permute(0, 3, 1,2) # [b,c,h,w]
+    return img
+
+def mirror_index(index, len_seq):
+    """
+    get mirror index when indexing a sequence and the index is larger than len_pose
+    args:
+        index: int
+        len_pose: int
+    return:
+        mirror_index: int
+    """
+    turn = index // len_seq
+    res = index % len_seq
+    if turn % 2 == 0:
+        return res # forward indexing
+    else:
+        return len_seq - res - 1 # reverse indexing
+    
+def smooth_camera_sequence(camera, kernel_size=7):
+    """
+    smooth the camera trajectory (i.e., rotation & translation)...
+    args:
+        camera: [N, 25] or [N, 16]. np.ndarray
+        kernel_size: int
+    return: 
+        smoothed_camera: [N, 25] or [N, 16]. np.ndarray
+    """
+    # poses: [N, 25], numpy array
+    N = camera.shape[0]
+    K = kernel_size // 2
+    poses = camera[:, :16].reshape([-1, 4, 4]).copy()
+    trans = poses[:, :3, 3].copy() # [N, 3]
+    rots = poses[:, :3, :3].copy() # [N, 3, 3]
+
+    for i in range(N):
+        start = max(0, i - K)
+        end = min(N, i + K + 1)
+        poses[i, :3, 3] = trans[start:end].mean(0)
+        try:
+            poses[i, :3, :3] = Rotation.from_matrix(rots[start:end]).mean().as_matrix()
+        except:
+            if i == 0:
+                poses[i, :3, :3] = rots[i]
+            else:
+                poses[i, :3, :3] = poses[i-1, :3, :3]
+    poses = poses.reshape([-1, 16])
+    camera[:, :16] = poses
+    return camera
+
+def smooth_features_xd(in_tensor, kernel_size=7):
+    """
+    smooth the feature maps
+    args:
+        in_tensor: [T, c,h,w] or [T, c1,c2,h,w]
+        kernel_size: int
+    return:
+        out_tensor: [T, c,h,w] or [T, c1,c2,h,w]
+    """
+    t = in_tensor.shape[0]
+    ndim = in_tensor.ndim
+    pad = (kernel_size- 1)//2
+    in_tensor = torch.cat([torch.flip(in_tensor[0:pad], dims=[0]), in_tensor, torch.flip(in_tensor[t-pad:t], dims=[0])], dim=0)
+    if ndim == 2: # tc
+        _,c = in_tensor.shape
+        in_tensor = in_tensor.permute(1,0).reshape([-1,1,t+2*pad]) # [c, 1, t]
+    elif ndim == 4: # tchw
+        _,c,h,w = in_tensor.shape
+        in_tensor = in_tensor.permute(1,2,3,0).reshape([-1,1,t+2*pad]) # [c, 1, t]
+    elif ndim == 5: # tcchw, like deformation
+        _,c1,c2, h,w = in_tensor.shape
+        in_tensor = in_tensor.permute(1,2,3,4,0).reshape([-1,1,t+2*pad]) # [c, 1, t]
+    else: raise NotImplementedError()
+    avg_kernel = 1 / kernel_size * torch.Tensor([1.]*kernel_size).reshape([1,1,kernel_size]).float().to(in_tensor.device) # [1, 1, kw]
+    out_tensor = F.conv1d(in_tensor, avg_kernel)
+    if ndim == 2: # tc
+        return out_tensor.reshape([c,t]).permute(1,0)
+    elif ndim == 4: # tchw
+        return out_tensor.reshape([c,h,w,t]).permute(3,0,1,2)
+    elif ndim == 5: # tcchw, like deformation
+        return out_tensor.reshape([c1,c2,h,w,t]).permute(4,0,1,2,3)
+
+
+def extract_audio_motion_from_ref_video(video_name):
+    def save_wav16k(audio_name):
+        supported_types = ('.wav', '.mp3', '.mp4', '.avi')
+        assert audio_name.endswith(supported_types), f"Now we only support {','.join(supported_types)} as audio source!"
+        wav16k_name = audio_name[:-4] + '_16k.wav'
+        extract_wav_cmd = f"ffmpeg -i {audio_name} -f wav -ar 16000 -v quiet -y {wav16k_name} -y"
+        os.system(extract_wav_cmd)
+        print(f"Extracted wav file (16khz) from {audio_name} to {wav16k_name}.")
+        return wav16k_name
+    
+    def get_f0( wav16k_name):
+        from data_gen.utils.process_audio.extract_mel_f0 import extract_mel_from_fname, extract_f0_from_wav_and_mel
+        wav, mel = extract_mel_from_fname(wav16k_name)
+        f0, f0_coarse = extract_f0_from_wav_and_mel(wav, mel)
+        f0 = f0.reshape([-1,1])
+        f0 = torch.tensor(f0)
+        return f0
+    
+    def get_hubert(wav16k_name):
+        from data_gen.utils.process_audio.extract_hubert import get_hubert_from_16k_wav
+        hubert = get_hubert_from_16k_wav(wav16k_name).detach().numpy()
+        len_mel = hubert.shape[0]
+        x_multiply = 8
+        if len_mel % x_multiply == 0:
+            num_to_pad = 0
+        else:
+            num_to_pad = x_multiply - len_mel % x_multiply
+        hubert = np.pad(hubert, pad_width=((0,num_to_pad), (0,0)))
+        hubert = torch.tensor(hubert)
+        return hubert
+
+    def get_exp(video_name):
+        from data_gen.utils.process_video.fit_3dmm_landmark import fit_3dmm_for_a_video
+        drv_motion_coeff_dict = fit_3dmm_for_a_video(video_name, save=False)
+        exp = torch.tensor(drv_motion_coeff_dict['exp'])
+        return exp
+    
+    wav16k_name = save_wav16k(video_name)
+    f0 = get_f0(wav16k_name)
+    hubert = get_hubert(wav16k_name)
+    os.system(f"rm {wav16k_name}")
+    exp = get_exp(video_name)
+    target_length = min(len(exp), len(hubert)//2, len(f0)//2)
+    exp = exp[:target_length]
+    f0 = f0[:target_length*2]
+    hubert = hubert[:target_length*2]
+    return exp.unsqueeze(0), hubert.unsqueeze(0), f0.unsqueeze(0)
+
+
+if __name__ == '__main__':
+    extract_audio_motion_from_ref_video('data/raw/videos/crop_0213.mp4')

inference/mimictalk_infer.py

@@ -0,0 +1,357 @@
+"""
+用于推理 inference/train_mimictalk_on_a_video.py 得到的person-specific模型
+"""
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+# import librosa
+import random
+import time
+import numpy as np
+import importlib
+import tqdm
+import copy
+import cv2
+
+# common utils
+from utils.commons.hparams import hparams, set_hparams
+from utils.commons.tensor_utils import move_to_cuda, convert_to_tensor
+from utils.commons.ckpt_utils import load_ckpt, get_last_checkpoint
+# 3DMM-related utils
+from deep_3drecon.deep_3drecon_models.bfm import ParametricFaceModel
+from data_util.face3d_helper import Face3DHelper
+from data_gen.utils.process_image.fit_3dmm_landmark import fit_3dmm_for_a_image
+from data_gen.utils.process_video.fit_3dmm_landmark import fit_3dmm_for_a_video
+from deep_3drecon.secc_renderer import SECC_Renderer
+from data_gen.eg3d.convert_to_eg3d_convention import get_eg3d_convention_camera_pose_intrinsic
+# Face Parsing 
+from data_gen.utils.mp_feature_extractors.mp_segmenter import MediapipeSegmenter
+from data_gen.utils.process_video.extract_segment_imgs import inpaint_torso_job, extract_background
+# other inference utils
+from inference.infer_utils import mirror_index, load_img_to_512_hwc_array, load_img_to_normalized_512_bchw_tensor
+from inference.infer_utils import smooth_camera_sequence, smooth_features_xd
+from inference.edit_secc import blink_eye_for_secc, hold_eye_opened_for_secc
+from inference.real3d_infer import GeneFace2Infer
+
+
+class AdaptGeneFace2Infer(GeneFace2Infer):
+    def __init__(self, audio2secc_dir, head_model_dir, torso_model_dir, device=None, **kwargs):
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = device
+        self.audio2secc_dir = audio2secc_dir
+        self.head_model_dir = head_model_dir
+        self.torso_model_dir = torso_model_dir
+        self.audio2secc_model = self.load_audio2secc(audio2secc_dir)
+        self.secc2video_model = self.load_secc2video(head_model_dir, torso_model_dir)
+        self.audio2secc_model.to(device).eval()
+        self.secc2video_model.to(device).eval()
+        self.seg_model = MediapipeSegmenter()
+        self.secc_renderer = SECC_Renderer(512)
+        self.face3d_helper = Face3DHelper(use_gpu=True, keypoint_mode='lm68')
+        self.mp_face3d_helper = Face3DHelper(use_gpu=True, keypoint_mode='mediapipe')
+        # self.camera_selector = KNearestCameraSelector()
+
+    def load_secc2video(self, head_model_dir, torso_model_dir):
+        if torso_model_dir != '':
+            config_dir = torso_model_dir if os.path.isdir(torso_model_dir) else os.path.dirname(torso_model_dir)
+            set_hparams(f"{config_dir}/config.yaml", print_hparams=False)
+            hparams['htbsr_head_threshold'] = 1.0
+            self.secc2video_hparams = copy.deepcopy(hparams)
+            ckpt = get_last_checkpoint(torso_model_dir)[0]
+            lora_args = ckpt.get("lora_args", None)
+            from modules.real3d.secc_img2plane_torso import OSAvatarSECC_Img2plane_Torso
+            model = OSAvatarSECC_Img2plane_Torso(self.secc2video_hparams, lora_args=lora_args)
+            load_ckpt(model, f"{torso_model_dir}", model_name='model', strict=True)
+            self.learnable_triplane = nn.Parameter(torch.zeros([1, 3, model.triplane_hid_dim*model.triplane_depth, 256, 256]).float().cuda(), requires_grad=True)
+            load_ckpt(self.learnable_triplane, f"{torso_model_dir}", model_name='learnable_triplane', strict=True)
+            model._last_cano_planes = self.learnable_triplane
+            if head_model_dir != '':
+                print("| Warning: Assigned --torso_ckpt which also contains head, but --head_ckpt is also assigned, skipping the --head_ckpt.")
+        else:
+            from modules.real3d.secc_img2plane_torso import OSAvatarSECC_Img2plane
+            set_hparams(f"{head_model_dir}/config.yaml", print_hparams=False)
+            ckpt = get_last_checkpoint(head_model_dir)[0]
+            lora_args = ckpt.get("lora_args", None)
+            self.secc2video_hparams = copy.deepcopy(hparams)
+            model = OSAvatarSECC_Img2plane(self.secc2video_hparams, lora_args=lora_args)
+            load_ckpt(model, f"{head_model_dir}", model_name='model', strict=True)
+            self.learnable_triplane = nn.Parameter(torch.zeros([1, 3, model.triplane_hid_dim*model.triplane_depth, 256, 256]).float().cuda(), requires_grad=True)
+            model._last_cano_planes = self.learnable_triplane
+            load_ckpt(model._last_cano_planes, f"{head_model_dir}", model_name='learnable_triplane', strict=True)
+        self.person_ds = ckpt['person_ds']
+        return model
+
+    def prepare_batch_from_inp(self, inp):
+        """
+        :param inp: {'audio_source_name': (str)}
+        :return: a dict that contains the condition feature of NeRF
+        """
+        sample = {}
+        # Process Driving Motion
+        if inp['drv_audio_name'][-4:] in ['.wav', '.mp3']:
+            self.save_wav16k(inp['drv_audio_name'])
+            if self.audio2secc_hparams['audio_type'] == 'hubert':
+                hubert = self.get_hubert(self.wav16k_name)
+            elif self.audio2secc_hparams['audio_type'] == 'mfcc':
+                hubert = self.get_mfcc(self.wav16k_name) / 100
+
+            f0 = self.get_f0(self.wav16k_name)
+            if f0.shape[0] > len(hubert):
+                f0 = f0[:len(hubert)]
+            else:
+                num_to_pad = len(hubert) - len(f0)
+                f0 = np.pad(f0, pad_width=((0,num_to_pad), (0,0)))
+            t_x = hubert.shape[0]
+            x_mask = torch.ones([1, t_x]).float() # mask for audio frames
+            y_mask = torch.ones([1, t_x//2]).float() # mask for motion/image frames
+            sample.update({
+                'hubert': torch.from_numpy(hubert).float().unsqueeze(0).cuda(),
+                'f0': torch.from_numpy(f0).float().reshape([1,-1]).cuda(),
+                'x_mask': x_mask.cuda(),
+                'y_mask': y_mask.cuda(),
+                })
+            sample['blink'] = torch.zeros([1, t_x, 1]).long().cuda()
+            sample['audio'] = sample['hubert']
+            sample['eye_amp'] = torch.ones([1, 1]).cuda() * 1.0
+        elif inp['drv_audio_name'][-4:] in ['.mp4']:
+            drv_motion_coeff_dict = fit_3dmm_for_a_video(inp['drv_audio_name'], save=False)
+            drv_motion_coeff_dict = convert_to_tensor(drv_motion_coeff_dict)
+            t_x = drv_motion_coeff_dict['exp'].shape[0] * 2
+            self.drv_motion_coeff_dict = drv_motion_coeff_dict
+        elif inp['drv_audio_name'][-4:] in ['.npy']:
+            drv_motion_coeff_dict = np.load(inp['drv_audio_name'], allow_pickle=True).tolist()
+            drv_motion_coeff_dict = convert_to_tensor(drv_motion_coeff_dict)
+            t_x = drv_motion_coeff_dict['exp'].shape[0] * 2
+            self.drv_motion_coeff_dict = drv_motion_coeff_dict
+
+        # Face Parsing
+        sample['ref_gt_img'] = self.person_ds['gt_img'].cuda()
+        img = self.person_ds['gt_img'].reshape([3, 512, 512]).permute(1, 2, 0)
+        img = (img + 1) * 127.5
+        img = np.ascontiguousarray(img.int().numpy()).astype(np.uint8)
+        segmap = self.seg_model._cal_seg_map(img)
+        sample['segmap'] = torch.tensor(segmap).float().unsqueeze(0).cuda()
+        head_img = self.seg_model._seg_out_img_with_segmap(img, segmap, mode='head')[0]
+        sample['ref_head_img'] = ((torch.tensor(head_img) - 127.5)/127.5).float().unsqueeze(0).permute(0, 3, 1,2).cuda() # [b,c,h,w]
+        inpaint_torso_img, _, _, _ = inpaint_torso_job(img, segmap)
+        sample['ref_torso_img'] = ((torch.tensor(inpaint_torso_img) - 127.5)/127.5).float().unsqueeze(0).permute(0, 3, 1,2).cuda() # [b,c,h,w]
+        
+        if inp['bg_image_name'] == '':
+            bg_img = extract_background([img], [segmap], 'knn')
+        else:
+            bg_img = cv2.imread(inp['bg_image_name'])
+            bg_img = cv2.cvtColor(bg_img, cv2.COLOR_BGR2RGB)
+            bg_img = cv2.resize(bg_img, (512,512))
+        sample['bg_img'] = ((torch.tensor(bg_img) - 127.5)/127.5).float().unsqueeze(0).permute(0, 3, 1,2).cuda() # [b,c,h,w]
+
+        # 3DMM, get identity code and camera pose
+        image_name = f"data/raw/val_imgs/{self.person_ds['video_id']}_img.png"
+        os.makedirs(os.path.dirname(image_name), exist_ok=True)
+        cv2.imwrite(image_name, img[:,:,::-1])
+        coeff_dict = fit_3dmm_for_a_image(image_name, save=False)
+        coeff_dict['id'] = self.person_ds['id'].reshape([1,80]).numpy()
+
+        assert coeff_dict is not None
+        src_id = torch.tensor(coeff_dict['id']).reshape([1,80]).cuda()
+        src_exp = torch.tensor(coeff_dict['exp']).reshape([1,64]).cuda()
+        src_euler = torch.tensor(coeff_dict['euler']).reshape([1,3]).cuda()
+        src_trans = torch.tensor(coeff_dict['trans']).reshape([1,3]).cuda()
+        sample['id'] = src_id.repeat([t_x//2,1])
+
+        # get the src_kp for torso model
+        sample['src_kp'] = self.person_ds['src_kp'].cuda().reshape([1, 68, 3]).repeat([t_x//2,1,1])[..., :2] # [B, 68, 2]
+
+        # get camera pose file
+        random.seed(time.time())
+        if inp['drv_pose_name'] in ['nearest', 'topk']:
+            camera_ret = get_eg3d_convention_camera_pose_intrinsic({'euler': torch.tensor(coeff_dict['euler']).reshape([1,3]), 'trans': torch.tensor(coeff_dict['trans']).reshape([1,3])})
+            c2w, intrinsics = camera_ret['c2w'], camera_ret['intrinsics']
+            camera = np.concatenate([c2w.reshape([1,16]), intrinsics.reshape([1,9])], axis=-1)
+            coeff_names, distance_matrix = self.camera_selector.find_k_nearest(camera, k=100)
+            coeff_names = coeff_names[0] # squeeze
+            if inp['drv_pose_name'] == 'nearest':
+                inp['drv_pose_name'] = coeff_names[0]
+            else:
+                inp['drv_pose_name'] = random.choice(coeff_names)
+            # inp['drv_pose_name'] = coeff_names[0]
+        elif inp['drv_pose_name'] == 'random':
+            inp['drv_pose_name'] = self.camera_selector.random_select()
+        else:
+            inp['drv_pose_name'] = inp['drv_pose_name']
+
+        print(f"| To extract pose from {inp['drv_pose_name']}")
+
+        # extract camera pose 
+        if inp['drv_pose_name'] == 'static':
+            sample['euler'] = torch.tensor(coeff_dict['euler']).reshape([1,3]).cuda().repeat([t_x//2,1]) # default static pose
+            sample['trans'] = torch.tensor(coeff_dict['trans']).reshape([1,3]).cuda().repeat([t_x//2,1])
+        else: # from file
+            if inp['drv_pose_name'].endswith('.mp4'):
+                # extract coeff from video
+                drv_pose_coeff_dict = fit_3dmm_for_a_video(inp['drv_pose_name'], save=False)
+            else:
+                # load from npy
+                drv_pose_coeff_dict = np.load(inp['drv_pose_name'], allow_pickle=True).tolist()
+            print(f"| Extracted pose from {inp['drv_pose_name']}")
+            eulers = convert_to_tensor(drv_pose_coeff_dict['euler']).reshape([-1,3]).cuda()
+            trans = convert_to_tensor(drv_pose_coeff_dict['trans']).reshape([-1,3]).cuda()
+            len_pose = len(eulers)
+            index_lst = [mirror_index(i, len_pose) for i in range(t_x//2)]
+            sample['euler'] = eulers[index_lst]
+            sample['trans'] = trans[index_lst]
+
+        # fix the z axis
+        sample['trans'][:, -1] = sample['trans'][0:1, -1].repeat([sample['trans'].shape[0]])
+
+        # mapping to the init pose
+        if inp.get("map_to_init_pose", 'False') == 'True':
+            diff_euler = torch.tensor(coeff_dict['euler']).reshape([1,3]).cuda() - sample['euler'][0:1]
+            sample['euler'] = sample['euler'] + diff_euler
+            diff_trans = torch.tensor(coeff_dict['trans']).reshape([1,3]).cuda() - sample['trans'][0:1]
+            sample['trans'] = sample['trans'] + diff_trans
+
+        # prepare camera
+        camera_ret = get_eg3d_convention_camera_pose_intrinsic({'euler':sample['euler'].cpu(), 'trans':sample['trans'].cpu()})
+        c2w, intrinsics = camera_ret['c2w'], camera_ret['intrinsics']
+        # smooth camera
+        camera_smo_ksize = 7
+        camera = np.concatenate([c2w.reshape([-1,16]), intrinsics.reshape([-1,9])], axis=-1)
+        camera = smooth_camera_sequence(camera, kernel_size=camera_smo_ksize) # [T, 25]
+        camera = torch.tensor(camera).cuda().float()
+        sample['camera'] = camera
+
+        return sample
+
+    @torch.no_grad()
+    def forward_secc2video(self, batch, inp=None):
+        num_frames = len(batch['drv_secc'])
+        camera = batch['camera']
+        src_kps = batch['src_kp']
+        drv_kps = batch['drv_kp']
+        cano_secc_color = batch['cano_secc']
+        src_secc_color = batch['src_secc']
+        drv_secc_colors = batch['drv_secc']
+        ref_img_gt = batch['ref_gt_img']
+        ref_img_head = batch['ref_head_img']
+        ref_torso_img = batch['ref_torso_img']
+        bg_img = batch['bg_img']
+        segmap = batch['segmap']
+        
+        # smooth torso drv_kp
+        torso_smo_ksize = 7
+        drv_kps = smooth_features_xd(drv_kps.reshape([-1, 68*2]), kernel_size=torso_smo_ksize).reshape([-1, 68, 2])
+
+        # forward renderer
+        img_raw_lst = []
+        img_lst = []
+        depth_img_lst = []
+        with torch.no_grad():
+            for i in tqdm.trange(num_frames, desc="MimicTalk is rendering frames"):
+                kp_src = torch.cat([src_kps[i:i+1].reshape([1, 68, 2]), torch.zeros([1, 68,1]).to(src_kps.device)],dim=-1)
+                kp_drv = torch.cat([drv_kps[i:i+1].reshape([1, 68, 2]), torch.zeros([1, 68,1]).to(drv_kps.device)],dim=-1)
+                cond={'cond_cano': cano_secc_color,'cond_src': src_secc_color, 'cond_tgt': drv_secc_colors[i:i+1].cuda(),
+                        'ref_torso_img': ref_torso_img, 'bg_img': bg_img, 'segmap': segmap,
+                        'kp_s': kp_src, 'kp_d': kp_drv}
+
+                ########################################################################################################
+                ### 相比real3d_infer只修改了这行👇，即cano_triplane来自cache里的learnable_triplane,而不是img预测的plane ####
+                ########################################################################################################
+                gen_output = self.secc2video_model.forward(img=None, camera=camera[i:i+1], cond=cond, ret={}, cache_backbone=False, use_cached_backbone=True)
+                
+                img_lst.append(gen_output['image'])
+                img_raw_lst.append(gen_output['image_raw'])
+                depth_img_lst.append(gen_output['image_depth'])
+
+        # save demo video
+        depth_imgs = torch.cat(depth_img_lst)
+        imgs = torch.cat(img_lst)
+        imgs_raw = torch.cat(img_raw_lst)
+        secc_img = torch.cat([torch.nn.functional.interpolate(drv_secc_colors[i:i+1], (512,512)) for i in range(num_frames)])
+        
+        if inp['out_mode'] == 'concat_debug':
+            secc_img = secc_img.cpu()
+            secc_img = ((secc_img + 1) * 127.5).permute(0, 2, 3, 1).int().numpy()
+
+            depth_img = F.interpolate(depth_imgs, (512,512)).cpu()
+            depth_img = depth_img.repeat([1,3,1,1])
+            depth_img = (depth_img - depth_img.min()) / (depth_img.max() - depth_img.min())
+            depth_img = depth_img * 2 - 1
+            depth_img = depth_img.clamp(-1,1)
+
+            secc_img = secc_img / 127.5 - 1
+            secc_img = torch.from_numpy(secc_img).permute(0, 3, 1, 2)
+            imgs = torch.cat([ref_img_gt.repeat([imgs.shape[0],1,1,1]).cpu(), secc_img, F.interpolate(imgs_raw, (512,512)).cpu(), depth_img, imgs.cpu()], dim=-1)
+        elif inp['out_mode'] == 'final':
+            imgs = imgs.cpu()
+        elif inp['out_mode'] == 'debug':
+            raise NotImplementedError("to do: save separate videos")
+        imgs = imgs.clamp(-1,1)
+
+        import imageio
+        import uuid
+        debug_name = f'{uuid.uuid1()}.mp4'
+        out_imgs = ((imgs.permute(0, 2, 3, 1) + 1)/2 * 255).int().cpu().numpy().astype(np.uint8)
+        writer = imageio.get_writer(debug_name, fps=25, format='FFMPEG', codec='h264')
+        for i in tqdm.trange(len(out_imgs), desc="Imageio is saving video"):
+            writer.append_data(out_imgs[i])
+        writer.close()
+        
+        out_fname = 'infer_out/tmp/' + os.path.basename(inp['drv_pose_name'])[:-4] + '.mp4' if inp['out_name'] == '' else inp['out_name']
+        try:
+            os.makedirs(os.path.dirname(out_fname), exist_ok=True)
+        except: pass
+        if inp['drv_audio_name'][-4:] in ['.wav', '.mp3']:
+            # os.system(f"ffmpeg -i {debug_name} -i {inp['drv_audio_name']} -y -v quiet -shortest {out_fname}")
+            cmd = f"/usr/bin/ffmpeg -i {debug_name} -i {self.wav16k_name} -y -r 25 -ar 16000 -c:v copy -c:a libmp3lame -pix_fmt yuv420p -b:v 2000k  -strict experimental -shortest {out_fname}"
+            os.system(cmd)
+            os.system(f"rm {debug_name}")
+        else:
+            ret = os.system(f"ffmpeg -i {debug_name} -i {inp['drv_audio_name']} -map 0:v -map 1:a -y -v quiet -shortest {out_fname}")
+            if ret != 0: # 没有成功从drv_audio_name里面提取到音频, 则直接输出无音频轨道的纯视频
+                os.system(f"mv {debug_name} {out_fname}")
+        print(f"Saved at {out_fname}")
+        return out_fname
+
+if __name__ == '__main__':
+    import argparse, glob, tqdm
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--a2m_ckpt", default='checkpoints/240112_icl_audio2secc_vox2_cmlr') # checkpoints/0727_audio2secc/audio2secc_withlm2d100_randomframe
+    parser.add_argument("--head_ckpt", default='') # checkpoints/0729_th1kh/secc_img2plane checkpoints/0720_img2planes/secc_img2plane_two_stage
+    parser.add_argument("--torso_ckpt", default='checkpoints_mimictalk/German_20s') 
+    parser.add_argument("--bg_img", default='') # data/raw/val_imgs/bg3.png
+    parser.add_argument("--drv_aud", default='data/raw/examples/80_vs_60_10s.wav')
+    parser.add_argument("--drv_pose", default='data/raw/examples/German_20s.mp4') # nearest | topk | random | static | vid_name
+    parser.add_argument("--drv_style", default='data/raw/examples/angry.mp4') # nearest | topk | random | static | vid_name
+    parser.add_argument("--blink_mode", default='period') # none | period
+    parser.add_argument("--temperature", default=0.3, type=float) # nearest | random
+    parser.add_argument("--denoising_steps", default=20, type=int) # nearest | random
+    parser.add_argument("--cfg_scale", default=1.5, type=float) # nearest | random
+    parser.add_argument("--out_name", default='') # nearest | random
+    parser.add_argument("--out_mode", default='concat_debug') # concat_debug | debug | final 
+    parser.add_argument("--hold_eye_opened", default='False') # concat_debug | debug | final 
+    parser.add_argument("--map_to_init_pose", default='True') # concat_debug | debug | final 
+    parser.add_argument("--seed", default=None, type=int) # random seed, default None to use time.time()
+ 
+    args = parser.parse_args()
+
+    inp = {
+            'a2m_ckpt': args.a2m_ckpt,
+            'head_ckpt': args.head_ckpt,
+            'torso_ckpt': args.torso_ckpt,
+            'bg_image_name': args.bg_img,
+            'drv_audio_name': args.drv_aud,
+            'drv_pose_name': args.drv_pose,
+            'drv_talking_style_name': args.drv_style,
+            'blink_mode': args.blink_mode,
+            'temperature': args.temperature,
+            'denoising_steps': args.denoising_steps,
+            'cfg_scale': args.cfg_scale,
+            'out_name': args.out_name,
+            'out_mode': args.out_mode,
+            'map_to_init_pose': args.map_to_init_pose,
+            'hold_eye_opened': args.hold_eye_opened,
+            'seed': args.seed,
+            }
+    AdaptGeneFace2Infer.example_run(inp)

inference/real3d_infer.py

@@ -0,0 +1,667 @@
+import os
+import torch
+import torch.nn.functional as F
+import torchshow as ts
+import librosa
+import random
+import time
+import numpy as np
+import importlib
+import tqdm
+import copy
+import cv2
+import math
+
+# common utils
+from utils.commons.hparams import hparams, set_hparams
+from utils.commons.tensor_utils import move_to_cuda, convert_to_tensor
+from utils.commons.ckpt_utils import load_ckpt, get_last_checkpoint
+# 3DMM-related utils
+from deep_3drecon.deep_3drecon_models.bfm import ParametricFaceModel
+from data_util.face3d_helper import Face3DHelper
+from data_gen.utils.process_image.fit_3dmm_landmark import fit_3dmm_for_a_image
+from data_gen.utils.process_video.fit_3dmm_landmark import fit_3dmm_for_a_video
+from data_gen.utils.process_image.extract_lm2d import extract_lms_mediapipe_job
+from data_gen.utils.process_image.fit_3dmm_landmark import index_lm68_from_lm468
+from deep_3drecon.secc_renderer import SECC_Renderer
+from data_gen.eg3d.convert_to_eg3d_convention import get_eg3d_convention_camera_pose_intrinsic
+# Face Parsing 
+from data_gen.utils.mp_feature_extractors.mp_segmenter import MediapipeSegmenter
+from data_gen.utils.process_video.extract_segment_imgs import inpaint_torso_job, extract_background
+# other inference utils
+from inference.infer_utils import mirror_index, load_img_to_512_hwc_array, load_img_to_normalized_512_bchw_tensor
+from inference.infer_utils import smooth_camera_sequence, smooth_features_xd
+from inference.edit_secc import blink_eye_for_secc, hold_eye_opened_for_secc
+
+
+def read_first_frame_from_a_video(vid_name):
+    frames = []
+    cap = cv2.VideoCapture(vid_name)
+    ret, frame_bgr = cap.read()
+    frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)
+    return frame_rgb
+
+def analyze_weights_img(gen_output):
+    img_raw = gen_output['image_raw']
+    mask_005_to_03 = torch.bitwise_and(gen_output['weights_img']>0.05, gen_output['weights_img']<0.3).repeat([1,3,1,1])
+    mask_005_to_05 = torch.bitwise_and(gen_output['weights_img']>0.05, gen_output['weights_img']<0.5).repeat([1,3,1,1])
+    mask_005_to_07 = torch.bitwise_and(gen_output['weights_img']>0.05, gen_output['weights_img']<0.7).repeat([1,3,1,1])
+    mask_005_to_09 = torch.bitwise_and(gen_output['weights_img']>0.05, gen_output['weights_img']<0.9).repeat([1,3,1,1])
+    mask_005_to_10 = torch.bitwise_and(gen_output['weights_img']>0.05, gen_output['weights_img']<1.0).repeat([1,3,1,1])
+
+    img_raw_005_to_03 = img_raw.clone()
+    img_raw_005_to_03[~mask_005_to_03] = -1
+    img_raw_005_to_05 = img_raw.clone()
+    img_raw_005_to_05[~mask_005_to_05] = -1
+    img_raw_005_to_07 = img_raw.clone()
+    img_raw_005_to_07[~mask_005_to_07] = -1
+    img_raw_005_to_09 = img_raw.clone()
+    img_raw_005_to_09[~mask_005_to_09] = -1
+    img_raw_005_to_10 = img_raw.clone()
+    img_raw_005_to_10[~mask_005_to_10] = -1
+    ts.save([img_raw_005_to_03[0], img_raw_005_to_05[0], img_raw_005_to_07[0], img_raw_005_to_09[0], img_raw_005_to_10[0]])
+
+
+def cal_face_area_percent(img_name):
+    img = cv2.resize(cv2.imread(img_name)[:,:,::-1], (512,512))
+    lm478 = extract_lms_mediapipe_job(img) / 512
+    min_x = lm478[:,0].min()
+    max_x = lm478[:,0].max()
+    min_y = lm478[:,1].min()
+    max_y = lm478[:,1].max()
+    area = (max_x - min_x) * (max_y - min_y)
+    return area
+
+def crop_img_on_face_area_percent(img_name, out_name='temp/cropped_src_img.png', min_face_area_percent=0.2):
+    try:
+        os.makedirs(os.path.dirname(out_name), exist_ok=True)
+    except: pass
+    face_area_percent = cal_face_area_percent(img_name)
+    if face_area_percent >= min_face_area_percent:
+        print(f"face area percent {face_area_percent} larger than threshold {min_face_area_percent}, directly use the input image...")
+        cmd = f"cp {img_name} {out_name}"
+        os.system(cmd)
+        return out_name
+    else:
+        print(f"face area percent {face_area_percent} smaller than threshold {min_face_area_percent}, crop the input image...")
+        img = cv2.resize(cv2.imread(img_name)[:,:,::-1], (512,512))
+        lm478 = extract_lms_mediapipe_job(img).astype(int)
+        min_x = lm478[:,0].min()
+        max_x = lm478[:,0].max()
+        min_y = lm478[:,1].min()
+        max_y = lm478[:,1].max()
+        face_area = (max_x - min_x) * (max_y - min_y)
+        target_total_area = face_area / min_face_area_percent
+        target_hw = int(target_total_area**0.5)
+        center_x, center_y = (min_x+max_x)/2, (min_y+max_y)/2
+        shrink_pixels = 2 * max(-(center_x - target_hw/2), center_x + target_hw/2 - 512, -(center_y - target_hw/2), center_y + target_hw/2-512)
+        shrink_pixels = max(0, shrink_pixels)
+        hw = math.floor(target_hw - shrink_pixels)
+        new_min_x = int(center_x - hw/2)
+        new_max_x = int(center_x + hw/2)
+        new_min_y = int(center_y - hw/2)
+        new_max_y = int(center_y + hw/2)
+
+        img = img[new_min_y:new_max_y, new_min_x:new_max_x]
+        img = cv2.resize(img, (512, 512))
+        cv2.imwrite(out_name, img[:,:,::-1])
+        return out_name
+
+
+class GeneFace2Infer:
+    def __init__(self, audio2secc_dir, head_model_dir, torso_model_dir, device=None, inp=None):
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.audio2secc_model = self.load_audio2secc(audio2secc_dir)
+        self.secc2video_model = self.load_secc2video(head_model_dir, torso_model_dir, inp)
+        self.audio2secc_model.to(device).eval()
+        self.secc2video_model.to(device).eval()
+        self.seg_model = MediapipeSegmenter()
+        self.secc_renderer = SECC_Renderer(512)
+        self.face3d_helper = Face3DHelper(use_gpu=True, keypoint_mode='lm68')
+        self.mp_face3d_helper = Face3DHelper(use_gpu=True, keypoint_mode='mediapipe')
+        self.camera_selector = KNearestCameraSelector()
+
+    def load_audio2secc(self, audio2secc_dir):
+        config_name = f"{audio2secc_dir}/config.yaml" if not audio2secc_dir.endswith(".ckpt") else f"{os.path.dirname(audio2secc_dir)}/config.yaml"
+        set_hparams(f"{config_name}", print_hparams=False)
+        self.audio2secc_dir = audio2secc_dir
+        self.audio2secc_hparams = copy.deepcopy(hparams)
+        from modules.audio2motion.vae import VAEModel, PitchContourVAEModel
+        from modules.audio2motion.cfm.icl_audio2motion_model import InContextAudio2MotionModel
+        if self.audio2secc_hparams['audio_type'] == 'hubert':
+            audio_in_dim = 1024
+        elif self.audio2secc_hparams['audio_type'] == 'mfcc':
+            audio_in_dim = 13
+
+        if 'icl' in hparams['task_cls']:
+            self.use_icl_audio2motion = True
+            model = InContextAudio2MotionModel(hparams['icl_model_type'], hparams=self.audio2secc_hparams)
+        else:
+            self.use_icl_audio2motion = False
+            if hparams.get("use_pitch", False) is True:
+                model = PitchContourVAEModel(hparams, in_out_dim=64, audio_in_dim=audio_in_dim)
+            else:
+                model = VAEModel(in_out_dim=64, audio_in_dim=audio_in_dim)
+        load_ckpt(model, f"{audio2secc_dir}", model_name='model', strict=True)
+        return model
+
+    def load_secc2video(self, head_model_dir, torso_model_dir, inp):
+        if inp is None:
+            inp = {}
+        self.head_model_dir = head_model_dir
+        self.torso_model_dir = torso_model_dir
+        if torso_model_dir != '':
+            if torso_model_dir.endswith(".ckpt"):
+                set_hparams(f"{os.path.dirname(torso_model_dir)}/config.yaml", print_hparams=False)
+            else:
+                set_hparams(f"{torso_model_dir}/config.yaml", print_hparams=False)
+            if inp.get('head_torso_threshold', None) is not None:
+                hparams['htbsr_head_threshold'] = inp['head_torso_threshold']
+            self.secc2video_hparams = copy.deepcopy(hparams)
+            from modules.real3d.secc_img2plane_torso import OSAvatarSECC_Img2plane_Torso
+            model = OSAvatarSECC_Img2plane_Torso()
+            load_ckpt(model, f"{torso_model_dir}", model_name='model', strict=False)
+            if head_model_dir != '':
+                print("| Warning: Assigned --torso_ckpt which also contains head, but --head_ckpt is also assigned, skipping the --head_ckpt.")
+        else:
+            from modules.real3d.secc_img2plane_torso import OSAvatarSECC_Img2plane
+            if head_model_dir.endswith(".ckpt"):
+                set_hparams(f"{os.path.dirname(head_model_dir)}/config.yaml", print_hparams=False)
+            else:
+                set_hparams(f"{head_model_dir}/config.yaml", print_hparams=False)
+            if inp.get('head_torso_threshold', None) is not None:
+                hparams['htbsr_head_threshold'] = inp['head_torso_threshold']
+            self.secc2video_hparams = copy.deepcopy(hparams)
+            model = OSAvatarSECC_Img2plane()
+            load_ckpt(model, f"{head_model_dir}", model_name='model', strict=False)
+        return model
+
+    def infer_once(self, inp):
+        self.inp = inp
+        samples = self.prepare_batch_from_inp(inp)
+        seed = inp['seed'] if inp['seed'] is not None else int(time.time())
+        random.seed(seed)
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+        out_name = self.forward_system(samples, inp)
+        return out_name
+    
+    def prepare_batch_from_inp(self, inp):
+        """
+        :param inp: {'audio_source_name': (str)}
+        :return: a dict that contains the condition feature of NeRF
+        """
+        cropped_name = 'temp/cropped_src_img_512.png'
+        crop_img_on_face_area_percent(inp['src_image_name'], cropped_name, min_face_area_percent=inp['min_face_area_percent'])
+        inp['src_image_name'] = cropped_name
+
+        sample = {}
+        # Process Driving Motion
+        if inp['drv_audio_name'][-4:] in ['.wav', '.mp3']:
+            self.save_wav16k(inp['drv_audio_name'])
+            if self.audio2secc_hparams['audio_type'] == 'hubert':
+                hubert = self.get_hubert(self.wav16k_name)
+            elif self.audio2secc_hparams['audio_type'] == 'mfcc':
+                hubert = self.get_mfcc(self.wav16k_name) / 100
+
+            f0 = self.get_f0(self.wav16k_name)
+            if f0.shape[0] > len(hubert):
+                f0 = f0[:len(hubert)]
+            else:
+                num_to_pad = len(hubert) - len(f0)
+                f0 = np.pad(f0, pad_width=((0,num_to_pad), (0,0)))
+            t_x = hubert.shape[0]
+            x_mask = torch.ones([1, t_x]).float() # mask for audio frames
+            y_mask = torch.ones([1, t_x//2]).float() # mask for motion/image frames
+            sample.update({
+                'hubert': torch.from_numpy(hubert).float().unsqueeze(0).cuda(),
+                'f0': torch.from_numpy(f0).float().reshape([1,-1]).cuda(),
+                'x_mask': x_mask.cuda(),
+                'y_mask': y_mask.cuda(),
+                })
+            sample['blink'] = torch.zeros([1, t_x, 1]).long().cuda()
+            sample['audio'] = sample['hubert']
+            sample['eye_amp'] = torch.ones([1, 1]).cuda() * 1.0
+            sample['mouth_amp'] = torch.ones([1, 1]).cuda() * inp['mouth_amp']
+        elif inp['drv_audio_name'][-4:] in ['.mp4']:
+            drv_motion_coeff_dict = fit_3dmm_for_a_video(inp['drv_audio_name'], save=False)
+            drv_motion_coeff_dict = convert_to_tensor(drv_motion_coeff_dict)
+            t_x = drv_motion_coeff_dict['exp'].shape[0] * 2
+            self.drv_motion_coeff_dict = drv_motion_coeff_dict
+        elif inp['drv_audio_name'][-4:] in ['.npy']:
+            drv_motion_coeff_dict = np.load(inp['drv_audio_name'], allow_pickle=True).tolist()
+            drv_motion_coeff_dict = convert_to_tensor(drv_motion_coeff_dict)
+            t_x = drv_motion_coeff_dict['exp'].shape[0] * 2
+            self.drv_motion_coeff_dict = drv_motion_coeff_dict
+        else:
+            raise ValueError()
+
+        # Face Parsing
+        image_name = inp['src_image_name']
+        if image_name.endswith(".mp4"):
+            img = read_first_frame_from_a_video(image_name)
+            image_name = inp['src_image_name'] = image_name[:-4] + '.png'
+            cv2.imwrite(image_name, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+        sample['ref_gt_img'] = load_img_to_normalized_512_bchw_tensor(image_name).cuda()
+        img = load_img_to_512_hwc_array(image_name)
+        segmap = self.seg_model._cal_seg_map(img)
+        sample['segmap'] = torch.tensor(segmap).float().unsqueeze(0).cuda()
+        head_img = self.seg_model._seg_out_img_with_segmap(img, segmap, mode='head')[0]
+        sample['ref_head_img'] = ((torch.tensor(head_img) - 127.5)/127.5).float().unsqueeze(0).permute(0, 3, 1,2).cuda() # [b,c,h,w]
+        inpaint_torso_img, _, _, _ = inpaint_torso_job(img, segmap)
+        sample['ref_torso_img'] = ((torch.tensor(inpaint_torso_img) - 127.5)/127.5).float().unsqueeze(0).permute(0, 3, 1,2).cuda() # [b,c,h,w]
+        
+        if inp['bg_image_name'] == '':
+            bg_img = extract_background([img], [segmap], 'lama')
+        else:
+            bg_img = cv2.imread(inp['bg_image_name'])
+            bg_img = cv2.cvtColor(bg_img, cv2.COLOR_BGR2RGB)
+            bg_img = cv2.resize(bg_img, (512,512))
+        sample['bg_img'] = ((torch.tensor(bg_img) - 127.5)/127.5).float().unsqueeze(0).permute(0, 3, 1,2).cuda() # [b,c,h,w]
+
+        # 3DMM, get identity code and camera pose
+        coeff_dict = fit_3dmm_for_a_image(image_name, save=False)
+        assert coeff_dict is not None
+        src_id = torch.tensor(coeff_dict['id']).reshape([1,80]).cuda()
+        src_exp = torch.tensor(coeff_dict['exp']).reshape([1,64]).cuda()
+        src_euler = torch.tensor(coeff_dict['euler']).reshape([1,3]).cuda()
+        src_trans = torch.tensor(coeff_dict['trans']).reshape([1,3]).cuda()
+        sample['id'] = src_id.repeat([t_x//2,1])
+
+        # get the src_kp for torso model
+        src_kp = self.face3d_helper.reconstruct_lm2d(src_id, src_exp, src_euler, src_trans) # [1, 68, 2]
+        src_kp = (src_kp-0.5) / 0.5 # rescale to -1~1
+        sample['src_kp'] = torch.clamp(src_kp, -1, 1).repeat([t_x//2,1,1])
+
+        # get camera pose file
+        # random.seed(time.time())
+        if inp['drv_pose_name'] in ['nearest', 'topk']:
+            camera_ret = get_eg3d_convention_camera_pose_intrinsic({'euler': torch.tensor(coeff_dict['euler']).reshape([1,3]), 'trans': torch.tensor(coeff_dict['trans']).reshape([1,3])})
+            c2w, intrinsics = camera_ret['c2w'], camera_ret['intrinsics']
+            camera = np.concatenate([c2w.reshape([1,16]), intrinsics.reshape([1,9])], axis=-1)
+            coeff_names, distance_matrix = self.camera_selector.find_k_nearest(camera, k=100)
+            coeff_names = coeff_names[0] # squeeze
+            if inp['drv_pose_name'] == 'nearest':
+                inp['drv_pose_name'] = coeff_names[0]
+            else:
+                inp['drv_pose_name'] = random.choice(coeff_names)
+            # inp['drv_pose_name'] = coeff_names[0]
+        elif inp['drv_pose_name'] == 'random':
+            inp['drv_pose_name'] = self.camera_selector.random_select()
+        else:
+            inp['drv_pose_name'] = inp['drv_pose_name']
+
+        print(f"| To extract pose from {inp['drv_pose_name']}")
+
+        # extract camera pose 
+        if inp['drv_pose_name'] == 'static':
+            sample['euler'] = torch.tensor(coeff_dict['euler']).reshape([1,3]).cuda().repeat([t_x//2,1]) # default static pose
+            sample['trans'] = torch.tensor(coeff_dict['trans']).reshape([1,3]).cuda().repeat([t_x//2,1])
+        else: # from file
+            if inp['drv_pose_name'].endswith('.mp4'):
+                # extract coeff from video
+                drv_pose_coeff_dict = fit_3dmm_for_a_video(inp['drv_pose_name'], save=False)
+            else:
+                # load from npy
+                drv_pose_coeff_dict = np.load(inp['drv_pose_name'], allow_pickle=True).tolist()
+            print(f"| Extracted pose from {inp['drv_pose_name']}")
+            eulers = convert_to_tensor(drv_pose_coeff_dict['euler']).reshape([-1,3]).cuda()
+            trans = convert_to_tensor(drv_pose_coeff_dict['trans']).reshape([-1,3]).cuda()
+            len_pose = len(eulers)
+            index_lst = [mirror_index(i, len_pose) for i in range(t_x//2)]
+            sample['euler'] = eulers[index_lst]
+            sample['trans'] = trans[index_lst]
+
+        # fix the z axis
+        sample['trans'][:, -1] = sample['trans'][0:1, -1].repeat([sample['trans'].shape[0]])
+
+        # mapping to the init pose
+        if inp.get("map_to_init_pose", 'False') == 'True':
+            diff_euler = torch.tensor(coeff_dict['euler']).reshape([1,3]).cuda() - sample['euler'][0:1]
+            sample['euler'] = sample['euler'] + diff_euler
+            diff_trans = torch.tensor(coeff_dict['trans']).reshape([1,3]).cuda() - sample['trans'][0:1]
+            sample['trans'] = sample['trans'] + diff_trans
+
+        # prepare camera
+        camera_ret = get_eg3d_convention_camera_pose_intrinsic({'euler':sample['euler'].cpu(), 'trans':sample['trans'].cpu()})
+        c2w, intrinsics = camera_ret['c2w'], camera_ret['intrinsics']
+        # smooth camera
+        camera_smo_ksize = 7
+        camera = np.concatenate([c2w.reshape([-1,16]), intrinsics.reshape([-1,9])], axis=-1)
+        camera = smooth_camera_sequence(camera, kernel_size=camera_smo_ksize) # [T, 25]
+        camera = torch.tensor(camera).cuda().float()
+        sample['camera'] = camera
+
+        return sample
+
+    @torch.no_grad()
+    def get_hubert(self, wav16k_name):
+        from data_gen.utils.process_audio.extract_hubert import get_hubert_from_16k_wav
+        hubert = get_hubert_from_16k_wav(wav16k_name).detach().numpy()
+        len_mel = hubert.shape[0]
+        x_multiply = 8
+        if len_mel % x_multiply == 0:
+            num_to_pad = 0
+        else:
+            num_to_pad = x_multiply - len_mel % x_multiply
+        hubert = np.pad(hubert, pad_width=((0,num_to_pad), (0,0)))
+        return hubert
+
+    def get_mfcc(self, wav16k_name):
+        from utils.audio import librosa_wav2mfcc
+        hparams['fft_size'] = 1200
+        hparams['win_size'] = 1200
+        hparams['hop_size'] = 480
+        hparams['audio_num_mel_bins'] = 80
+        hparams['fmin'] = 80
+        hparams['fmax'] = 12000
+        hparams['audio_sample_rate'] = 24000
+        mfcc = librosa_wav2mfcc(wav16k_name,
+            fft_size=hparams['fft_size'],
+            hop_size=hparams['hop_size'],
+            win_length=hparams['win_size'],
+            num_mels=hparams['audio_num_mel_bins'],
+            fmin=hparams['fmin'],
+            fmax=hparams['fmax'],
+            sample_rate=hparams['audio_sample_rate'],
+            center=True)
+        mfcc = np.array(mfcc).reshape([-1, 13])
+        len_mel = mfcc.shape[0]
+        x_multiply = 8
+        if len_mel % x_multiply == 0:
+            num_to_pad = 0
+        else:
+            num_to_pad = x_multiply - len_mel % x_multiply
+        mfcc = np.pad(mfcc, pad_width=((0,num_to_pad), (0,0)))
+        return mfcc
+
+    @torch.no_grad()
+    def forward_audio2secc(self, batch, inp=None):
+        if inp['drv_audio_name'][-4:] in ['.wav', '.mp3']:
+            from inference.infer_utils import extract_audio_motion_from_ref_video
+            if self.use_icl_audio2motion:
+                self.audio2secc_model.empty_context() # make this function reloadable
+            if self.use_icl_audio2motion and inp['drv_talking_style_name'].endswith(".mp4"):
+                ref_exp, ref_hubert, ref_f0 = extract_audio_motion_from_ref_video(inp['drv_talking_style_name'])
+                self.audio2secc_model.add_sample_to_context(ref_exp, ref_hubert, ref_f0)
+            elif self.use_icl_audio2motion and inp['drv_talking_style_name'].endswith((".png",'.jpg')):
+                style_coeff_dict = fit_3dmm_for_a_image(inp['drv_talking_style_name'])
+                ref_exp = torch.tensor(style_coeff_dict['exp']).reshape([1,1,64]).cuda()
+                self.audio2secc_model.add_sample_to_context(ref_exp.repeat([1, 100, 1]), hubert=None, f0=None)
+            else:
+                print("| WARNING: Not assigned reference talking style, passing...")
+            # audio-to-exp
+            ret = {}
+            # pred = self.audio2secc_model.forward(batch, ret=ret,train=False, ,)
+            pred = self.audio2secc_model.forward(batch, ret=ret,train=False, temperature=inp['temperature'], denoising_steps=inp['denoising_steps'], cond_scale=inp['cfg_scale'])
+
+            print("| audio-to-motion finished")
+            if pred.shape[-1] == 144:
+                id = ret['pred'][0][:,:80]
+                exp = ret['pred'][0][:,80:]
+            else:
+                id = batch['id']
+                exp = ret['pred'][0]
+            if len(id) < len(exp): # happens when use ICL
+                id = torch.cat([id, id[0].unsqueeze(0).repeat([len(exp)-len(id),1])])
+            batch['id'] = id
+            batch['exp'] = exp
+        else:
+            drv_motion_coeff_dict = self.drv_motion_coeff_dict
+            batch['exp'] = torch.FloatTensor(drv_motion_coeff_dict['exp']).cuda()
+
+        batch['id'] = batch['id'][:-4]
+        batch['exp'] = batch['exp'][:-4]
+        batch['euler'] = batch['euler'][:-4]
+        batch['trans'] = batch['trans'][:-4]
+        batch = self.get_driving_motion(batch['id'], batch['exp'], batch['euler'], batch['trans'], batch, inp)
+        if self.use_icl_audio2motion:
+            self.audio2secc_model.empty_context()
+        return batch
+
+    @torch.no_grad()
+    def get_driving_motion(self, id, exp, euler, trans, batch, inp):
+        zero_eulers = torch.zeros([id.shape[0], 3]).to(id.device)
+        zero_trans = torch.zeros([id.shape[0], 3]).to(exp.device)
+        # render the secc given the id,exp
+        with torch.no_grad():
+            chunk_size = 50
+            drv_secc_color_lst = []
+            num_iters = len(id)//chunk_size if len(id)%chunk_size == 0 else len(id)//chunk_size+1
+            for i in tqdm.trange(num_iters, desc="rendering drv secc"):
+                torch.cuda.empty_cache()
+                face_mask, drv_secc_color = self.secc_renderer(id[i*chunk_size:(i+1)*chunk_size], exp[i*chunk_size:(i+1)*chunk_size], zero_eulers[i*chunk_size:(i+1)*chunk_size], zero_trans[i*chunk_size:(i+1)*chunk_size])
+                drv_secc_color_lst.append(drv_secc_color.cpu())
+        drv_secc_colors = torch.cat(drv_secc_color_lst, dim=0)
+        _, src_secc_color = self.secc_renderer(id[0:1], exp[0:1], zero_eulers[0:1], zero_trans[0:1])
+        _, cano_secc_color = self.secc_renderer(id[0:1], exp[0:1]*0, zero_eulers[0:1], zero_trans[0:1])
+        batch['drv_secc'] = drv_secc_colors.cuda()
+        batch['src_secc'] = src_secc_color.cuda()
+        batch['cano_secc'] = cano_secc_color.cuda()
+        
+        # blinking secc
+        if inp['blink_mode'] == 'period':
+            period = 5 # second
+
+            if inp['hold_eye_opened'] == 'True':
+                for i in tqdm.trange(len(drv_secc_colors),desc="opening eye for secc"):
+                    batch['drv_secc'][i] = hold_eye_opened_for_secc(batch['drv_secc'][i])
+
+            for i in tqdm.trange(len(drv_secc_colors),desc="blinking secc"):
+                if i % (25*period) == 0:
+                    blink_dur_frames = random.randint(8, 12)
+                    for offset in range(blink_dur_frames):
+                        j = offset + i
+                        if j >= len(drv_secc_colors)-1: break
+                        def blink_percent_fn(t, T):
+                            return -4/T**2 * t**2 + 4/T * t
+                        blink_percent = blink_percent_fn(offset, blink_dur_frames)
+                        secc = batch['drv_secc'][j]
+                        out_secc = blink_eye_for_secc(secc, blink_percent)
+                        out_secc = out_secc.cuda()
+                        batch['drv_secc'][j] = out_secc
+
+        # get the drv_kp for torso model, using the transformed trajectory
+        drv_kp = self.face3d_helper.reconstruct_lm2d(id, exp, euler, trans) # [T, 68, 2]
+
+        drv_kp = (drv_kp-0.5) / 0.5 # rescale to -1~1
+        batch['drv_kp'] = torch.clamp(drv_kp, -1, 1)
+        return batch
+
+    @torch.no_grad()
+    def forward_secc2video(self, batch, inp=None):
+        num_frames = len(batch['drv_secc'])
+        camera = batch['camera']
+        src_kps = batch['src_kp']
+        drv_kps = batch['drv_kp']
+        cano_secc_color = batch['cano_secc']
+        src_secc_color = batch['src_secc']
+        drv_secc_colors = batch['drv_secc']
+        ref_img_gt = batch['ref_gt_img']
+        ref_img_head = batch['ref_head_img']
+        ref_torso_img = batch['ref_torso_img']
+        bg_img = batch['bg_img']
+        segmap = batch['segmap']
+        
+        # smooth torso drv_kp
+        torso_smo_ksize = 7
+        drv_kps = smooth_features_xd(drv_kps.reshape([-1, 68*2]), kernel_size=torso_smo_ksize).reshape([-1, 68, 2])
+
+        # forward renderer
+        img_raw_lst = []
+        img_lst = []
+        depth_img_lst = []
+        with torch.no_grad():
+            with torch.cuda.amp.autocast(inp['fp16']):
+                for i in tqdm.trange(num_frames, desc="Real3D-Portrait is rendering frames"):
+                    kp_src = torch.cat([src_kps[i:i+1].reshape([1, 68, 2]), torch.zeros([1, 68,1]).to(src_kps.device)],dim=-1)
+                    kp_drv = torch.cat([drv_kps[i:i+1].reshape([1, 68, 2]), torch.zeros([1, 68,1]).to(drv_kps.device)],dim=-1)
+                    cond={'cond_cano': cano_secc_color,'cond_src': src_secc_color, 'cond_tgt': drv_secc_colors[i:i+1].cuda(),
+                            'ref_torso_img': ref_torso_img, 'bg_img': bg_img, 'segmap': segmap,
+                            'kp_s': kp_src, 'kp_d': kp_drv,
+                            'ref_cameras': camera[i:i+1],
+                            }
+                    if i == 0:
+                        gen_output = self.secc2video_model.forward(img=ref_img_head, camera=camera[i:i+1], cond=cond, ret={}, cache_backbone=True, use_cached_backbone=False)
+                    else:
+                        gen_output = self.secc2video_model.forward(img=ref_img_head, camera=camera[i:i+1], cond=cond, ret={}, cache_backbone=False, use_cached_backbone=True)
+                    img_lst.append(gen_output['image'])
+                    img_raw_lst.append(gen_output['image_raw'])
+                    depth_img_lst.append(gen_output['image_depth'])
+
+        # save demo video
+        depth_imgs = torch.cat(depth_img_lst)
+        imgs = torch.cat(img_lst)
+        imgs_raw = torch.cat(img_raw_lst)
+        secc_img = torch.cat([torch.nn.functional.interpolate(drv_secc_colors[i:i+1], (512,512)) for i in range(num_frames)])
+        
+        if inp['out_mode'] == 'concat_debug':
+            secc_img = secc_img.cpu()
+            secc_img = ((secc_img + 1) * 127.5).permute(0, 2, 3, 1).int().numpy()
+
+            depth_img = F.interpolate(depth_imgs, (512,512)).cpu()
+            depth_img = depth_img.repeat([1,3,1,1])
+            depth_img = (depth_img - depth_img.min()) / (depth_img.max() - depth_img.min())
+            depth_img = depth_img * 2 - 1
+            depth_img = depth_img.clamp(-1,1)
+
+            secc_img = secc_img / 127.5 - 1
+            secc_img = torch.from_numpy(secc_img).permute(0, 3, 1, 2)
+            imgs = torch.cat([ref_img_gt.repeat([imgs.shape[0],1,1,1]).cpu(), secc_img, F.interpolate(imgs_raw, (512,512)).cpu(), depth_img, imgs.cpu()], dim=-1)
+        elif inp['out_mode'] == 'final':
+            imgs = imgs.cpu()
+        elif inp['out_mode'] == 'debug':
+            raise NotImplementedError("to do: save separate videos")
+        imgs = imgs.clamp(-1,1)
+
+        import imageio
+        debug_name = 'demo.mp4'
+        out_imgs = ((imgs.permute(0, 2, 3, 1) + 1)/2 * 255).int().cpu().numpy().astype(np.uint8)
+        writer = imageio.get_writer(debug_name, fps=25, format='FFMPEG', codec='h264')
+        
+        for i in tqdm.trange(len(out_imgs), desc="Imageio is saving video"):
+            writer.append_data(out_imgs[i])
+        writer.close()
+        
+        out_fname = 'infer_out/tmp/' + os.path.basename(inp['src_image_name'])[:-4] + '_' + os.path.basename(inp['drv_pose_name'])[:-4] + '.mp4' if inp['out_name'] == '' else inp['out_name']
+        try:
+            os.makedirs(os.path.dirname(out_fname), exist_ok=True)
+        except: pass
+        if inp['drv_audio_name'][-4:] in ['.wav', '.mp3']:
+            # cmd = f"ffmpeg -i {debug_name} -i {self.wav16k_name} -y -shortest {out_fname}"
+            cmd = f"ffmpeg -i {debug_name} -i {self.wav16k_name} -y -v quiet -shortest {out_fname}"
+            print(cmd)
+            os.system(cmd)
+            os.system(f"rm {debug_name}")
+            os.system(f"rm {self.wav16k_name}")
+        else:
+            ret = os.system(f"ffmpeg -i {debug_name} -i {inp['drv_audio_name']} -map 0:v -map 1:a -y -v quiet -shortest {out_fname}")
+            if ret != 0: # 没有成功从drv_audio_name里面提取到音频, 则直接输出无音频轨道的纯视频
+                os.system(f"mv {debug_name} {out_fname}")
+        print(f"Saved at {out_fname}")
+        return out_fname
+        
+    @torch.no_grad()
+    def forward_system(self, batch, inp):
+        self.forward_audio2secc(batch, inp)
+        out_fname = self.forward_secc2video(batch, inp)
+        return out_fname
+
+    @classmethod
+    def example_run(cls, inp=None):
+        inp_tmp = {
+            'drv_audio_name': 'data/raw/val_wavs/zozo.wav',
+            'src_image_name': 'data/raw/val_imgs/Macron.png'
+            }
+        if inp is not None:
+            inp_tmp.update(inp)
+        inp = inp_tmp
+
+        infer_instance = cls(inp['a2m_ckpt'], inp['head_ckpt'], inp['torso_ckpt'], inp=inp)
+        infer_instance.infer_once(inp)
+
+    ##############
+    # IO-related
+    ##############
+    def save_wav16k(self, audio_name):
+        supported_types = ('.wav', '.mp3', '.mp4', '.avi')
+        assert audio_name.endswith(supported_types), f"Now we only support {','.join(supported_types)} as audio source!"
+        import uuid
+        wav16k_name = audio_name[:-4] + f'{uuid.uuid1()}_16k.wav'
+        self.wav16k_name = wav16k_name
+        extract_wav_cmd = f"ffmpeg -i {audio_name} -f wav -ar 16000 -v quiet -y {wav16k_name} -y"
+        # extract_wav_cmd = f"ffmpeg -i {audio_name} -f wav -ar 16000 -y {wav16k_name} -y"
+        print(extract_wav_cmd)
+        os.system(extract_wav_cmd)
+        print(f"Extracted wav file (16khz) from {audio_name} to {wav16k_name}.")
+
+    def get_f0(self, wav16k_name):
+        from data_gen.utils.process_audio.extract_mel_f0 import extract_mel_from_fname, extract_f0_from_wav_and_mel
+        wav, mel = extract_mel_from_fname(self.wav16k_name)
+        f0, f0_coarse = extract_f0_from_wav_and_mel(wav, mel)
+        f0 = f0.reshape([-1,1])
+        return f0
+
+if __name__ == '__main__':
+    import argparse, glob, tqdm
+    parser = argparse.ArgumentParser()
+    # parser.add_argument("--a2m_ckpt", default='checkpoints/240112_audio2secc/icl_audio2secc_vox2_cmlr') # checkpoints/0727_audio2secc/audio2secc_withlm2d100_randomframe
+    parser.add_argument("--a2m_ckpt", default='checkpoints/240126_real3dportrait_orig/audio2secc_vae') # checkpoints/0727_audio2secc/audio2secc_withlm2d100_randomframe
+    parser.add_argument("--head_ckpt", default='') # checkpoints/0729_th1kh/secc_img2plane checkpoints/0720_img2planes/secc_img2plane_two_stage
+    # parser.add_argument("--head_ckpt", default='checkpoints/240210_os_secc2plane/secc2plane_trigridv2_blink0.3_pertubeNone') # checkpoints/0729_th1kh/secc_img2plane checkpoints/0720_img2planes/secc_img2plane_two_stage
+    # parser.add_argument("--torso_ckpt", default='') 
+    # parser.add_argument("--torso_ckpt", default='checkpoints/240209_robust_secc2plane_torso/secc2plane_torso_orig_fuseV1_MulMaskFalse') 
+    parser.add_argument("--torso_ckpt", default='checkpoints/240211_robust_secc2plane_torso/secc2plane_torso_orig_fuseV3_MulMaskTrue') 
+    # parser.add_argument("--torso_ckpt", default='checkpoints/240209_robust_secc2plane_torso/secc2plane_torso_orig_fuseV2_MulMaskTrue') 
+    # parser.add_argument("--src_img", default='data/raw/val_imgs/Macron_img.png')
+    # parser.add_argument("--src_img", default='gf2_iclr_test_data/cross_imgs/Trump.png')
+    parser.add_argument("--src_img", default='data/raw/val_imgs/mercy.png')
+    parser.add_argument("--bg_img", default='') # data/raw/val_imgs/bg3.png
+    parser.add_argument("--drv_aud", default='data/raw/val_wavs/yiwise.wav')
+    parser.add_argument("--drv_pose", default='infer_out/240319_tta/trump.mp4') # nearest | topk | random | static | vid_name
+    # parser.add_argument("--drv_pose", default='nearest') # nearest | topk | random | static | vid_name
+    parser.add_argument("--drv_style", default='') # nearest | topk | random | static | vid_name
+    # parser.add_argument("--drv_style", default='infer_out/240319_tta/trump.mp4') # nearest | topk | random | static | vid_name
+    parser.add_argument("--blink_mode", default='period') # none | period
+    parser.add_argument("--temperature", default=0.2, type=float) # nearest | random
+    parser.add_argument("--denoising_steps", default=20, type=int) # nearest | random
+    parser.add_argument("--cfg_scale", default=2.5, type=float) # nearest | random
+    parser.add_argument("--mouth_amp", default=0.4, type=float) # scale of predicted mouth, enabled in audio-driven
+    parser.add_argument("--min_face_area_percent", default=0.2, type=float) # scale of predicted mouth, enabled in audio-driven
+    parser.add_argument("--head_torso_threshold", default=0.5, type=float, help="0.1~1.0, 如果发现头发有半透明的现象,调小该值,以将小weights的头发直接clamp到weights=1.0; 如果发现头外部有荧光色的虚影,调小这个值. 对不同超参的Nerf也是case-to-case")
+    # parser.add_argument("--head_torso_threshold", default=None, type=float, help="0.1~1.0, 如果发现头发有半透明的现象,调小该值,以将小weights的头发直接clamp到weights=1.0; 如果发现头外部有荧光色的虚影,调小这个值. 对不同超参的Nerf也是case-to-case")
+    parser.add_argument("--out_name", default='') # nearest | random
+    parser.add_argument("--out_mode", default='concat_debug') # concat_debug | debug | final 
+    parser.add_argument("--hold_eye_opened", default='False') # concat_debug | debug | final 
+    parser.add_argument("--map_to_init_pose", default='True') # concat_debug | debug | final 
+    parser.add_argument("--seed", default=None, type=int) # random seed, default None to use time.time()
+    parser.add_argument("--fp16", action='store_true')
+
+    args = parser.parse_args()
+
+    inp = {
+            'a2m_ckpt': args.a2m_ckpt,
+            'head_ckpt': args.head_ckpt,
+            'torso_ckpt': args.torso_ckpt,
+            'src_image_name': args.src_img,
+            'bg_image_name': args.bg_img,
+            'drv_audio_name': args.drv_aud,
+            'drv_pose_name': args.drv_pose,
+            'drv_talking_style_name': args.drv_style,
+            'blink_mode': args.blink_mode,
+            'temperature': args.temperature,
+            'mouth_amp': args.mouth_amp,
+            'out_name': args.out_name,
+            'out_mode': args.out_mode,
+            'map_to_init_pose': args.map_to_init_pose,
+            'hold_eye_opened': args.hold_eye_opened,
+            'head_torso_threshold': args.head_torso_threshold,
+            'min_face_area_percent': args.min_face_area_percent,
+            'denoising_steps': args.denoising_steps,
+            'cfg_scale': args.cfg_scale,
+            'seed': args.seed,
+            'fp16': args.fp16, # 目前的ckpt使用fp16会导致nan，发现是因为i2p模型的layernorm产生了单个nan导致的，在训练阶段也采用fp16可能可以解决这个问题
+            }
+
+    GeneFace2Infer.example_run(inp)

inference/real3dportrait_demo.ipynb

@@ -0,0 +1,287 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "view-in-github"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/yerfor/Real3DPortrait/blob/main/real3dportrait_demo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "QS04K9oO21AW"
+      },
+      "source": [
+        "Check GPU"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1ESQRDb-yVUG"
+      },
+      "outputs": [],
+      "source": [
+        "!nvidia-smi --query-gpu=name,memory.total,memory.free --format=csv,noheader"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "y-ctmIvu3Ei8"
+      },
+      "source": [
+        "Installation"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "gXu76wdDgaxo"
+      },
+      "outputs": [],
+      "source": [
+        "# install pytorch3d, about 15s\n",
+        "import os\n",
+        "import sys\n",
+        "import torch\n",
+        "need_pytorch3d=False\n",
+        "try:\n",
+        "    import pytorch3d\n",
+        "except ModuleNotFoundError:\n",
+        "    need_pytorch3d=True\n",
+        "if need_pytorch3d:\n",
+        "    if torch.__version__.startswith(\"2.1.\") and sys.platform.startswith(\"linux\"):\n",
+        "        # We try to install PyTorch3D via a released wheel.\n",
+        "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
+        "        version_str=\"\".join([\n",
+        "            f\"py3{sys.version_info.minor}_cu\",\n",
+        "            torch.version.cuda.replace(\".\",\"\"),\n",
+        "            f\"_pyt{pyt_version_str}\"\n",
+        "        ])\n",
+        "        !pip install fvcore iopath\n",
+        "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
+        "    else:\n",
+        "        # We try to install PyTorch3D from source.\n",
+        "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "DuUynxmotG_-"
+      },
+      "outputs": [],
+      "source": [
+        "# install dependencies, about 5~10 min\n",
+        "!pip install tensorboard==2.13.0 tensorboardX==2.6.1\n",
+        "!pip install pyspy==0.1.1\n",
+        "!pip install protobuf==3.20.3\n",
+        "!pip install scipy==1.9.1\n",
+        "!pip install kornia==0.5.0\n",
+        "!pip install trimesh==3.22.0\n",
+        "!pip install einops==0.6.1 torchshow==0.5.1\n",
+        "!pip install imageio==2.31.1 imageio-ffmpeg==0.4.8\n",
+        "!pip install scikit-learn==1.3.0 scikit-image==0.21.0\n",
+        "!pip install av==10.0.0 lpips==0.1.4\n",
+        "!pip install timm==0.9.2 librosa==0.9.2\n",
+        "!pip install openmim==0.3.9\n",
+        "!mim install mmcv==2.1.0 # use mim to speed up installation for mmcv\n",
+        "!pip install transformers==4.33.2\n",
+        "!pip install pretrainedmodels==0.7.4\n",
+        "!pip install ninja==1.11.1\n",
+        "!pip install faiss-cpu==1.7.4\n",
+        "!pip install praat-parselmouth==0.4.3 moviepy==1.0.3\n",
+        "!pip install mediapipe==0.10.7\n",
+        "!pip install --upgrade attr\n",
+        "!pip install beartype==0.16.4 gateloop_transformer==0.4.0\n",
+        "!pip install torchode==0.2.0 torchdiffeq==0.2.3\n",
+        "!pip install hydra-core==1.3.2 pandas==2.1.3\n",
+        "!pip install pytorch_lightning==2.1.2\n",
+        "!pip install httpx==0.23.3\n",
+        "!pip install gradio==4.16.0\n",
+        "!pip install gdown\n",
+        "!pip install pyloudnorm webrtcvad pyworld==0.2.1rc0 pypinyin==0.42.0"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0GLEV0HVu8rj"
+      },
+      "outputs": [],
+      "source": [
+        "# RESTART kernel to make sure runtime is correct if you meet runtime errors\n",
+        "# import os\n",
+        "# os.kill(os.getpid(), 9)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5UfKHKrH6kcq"
+      },
+      "source": [
+        "Clone code and download checkpoints"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "-gfRsd9DwIgl"
+      },
+      "outputs": [],
+      "source": [
+        "# clone Real3DPortrait repo from github\n",
+        "!git clone https://github.com/yerfor/Real3DPortrait\n",
+        "%cd Real3DPortrait"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "Yju8dQY7x5OS"
+      },
+      "outputs": [],
+      "source": [
+        "# download pretrained ckpts & third-party ckpts from google drive, about 1 min\n",
+        "!pip install --upgrade --no-cache-dir gdown\n",
+        "%cd deep_3drecon/BFM\n",
+        "!gdown https://drive.google.com/uc?id=1SPM3IHsyNAaVMwqZZGV6QVaV7I2Hly0v\n",
+        "!gdown https://drive.google.com/uc?id=1MSldX9UChKEb3AXLVTPzZQcsbGD4VmGF\n",
+        "!gdown https://drive.google.com/uc?id=180ciTvm16peWrcpl4DOekT9eUQ-lJfMU\n",
+        "!gdown https://drive.google.com/uc?id=1KX9MyGueFB3M-X0Ss152x_johyTXHTfU\n",
+        "!gdown https://drive.google.com/uc?id=19-NyZn_I0_mkF-F5GPyFMwQJ_-WecZIL\n",
+        "!gdown https://drive.google.com/uc?id=11ouQ7Wr2I-JKStp2Fd1afedmWeuifhof\n",
+        "!gdown https://drive.google.com/uc?id=18ICIvQoKX-7feYWP61RbpppzDuYTptCq\n",
+        "!gdown https://drive.google.com/uc?id=1VktuY46m0v_n_d4nvOupauJkK4LF6mHE\n",
+        "%cd ../..\n",
+        "\n",
+        "%cd checkpoints\n",
+        "!gdown https://drive.google.com/uc?id=1gz8A6xestHp__GbZT5qozb43YaybRJhZ\n",
+        "!gdown https://drive.google.com/uc?id=1gSUIw2AkkKnlLJnNfS2FCqtaVw9tw3QF\n",
+        "!unzip 240210_real3dportrait_orig.zip\n",
+        "!unzip pretrained_ckpts.zip\n",
+        "!ls\n",
+        "%cd ..\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "LHzLro206pnA"
+      },
+      "source": [
+        "Inference sample"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "!python inference/real3d_infer.py -h"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2aCDwxNivQoS"
+      },
+      "outputs": [],
+      "source": [
+        "# sample inference, about 3 min\n",
+        "!python inference/real3d_infer.py \\\n",
+        "--src_img data/raw/examples/Macron.png \\\n",
+        "--drv_aud data/raw/examples/Obama_5s.wav \\\n",
+        "--drv_pose data/raw/examples/May_5s.mp4 \\\n",
+        "--bg_img data/raw/examples/bg.png \\\n",
+        "--out_name output.mp4 \\\n",
+        "--out_mode concat_debug \\\n",
+        "--low_memory_usage"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "XL0c54l19mBG"
+      },
+      "source": [
+        "Display output video"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6olmWwZP9Icj"
+      },
+      "outputs": [],
+      "source": [
+        "# borrow code from makeittalk\n",
+        "from IPython.display import HTML\n",
+        "from base64 import b64encode\n",
+        "import os, sys\n",
+        "import glob\n",
+        "\n",
+        "mp4_name = './output.mp4'\n",
+        "\n",
+        "mp4 = open('{}'.format(mp4_name),'rb').read()\n",
+        "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
+        "\n",
+        "print('Display animation: {}'.format(mp4_name), file=sys.stderr)\n",
+        "display(HTML(\"\"\"\n",
+        "  <video width=256 controls>\n",
+        "        <source src=\"%s\" type=\"video/mp4\">\n",
+        "  </video>\n",
+        "  \"\"\" % data_url))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "WebUI"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "\n",
+        "!python inference/app_real3dportrait.py --share"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "authorship_tag": "ABX9TyPu++zOlOS4yKF4xn4FHGtZ",
+      "gpuType": "T4",
+      "include_colab_link": true,
+      "private_outputs": true,
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

inference/train_mimictalk_on_a_video.py

@@ -0,0 +1,608 @@
+"""
+将One-shot的说话人大模型(os_secc2plane or os_secc2plane_torso)在单一说话人(一张照片或一段视频)上overfit, 实现和GeneFace++类似的效果
+"""
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import librosa
+import random
+import time
+import numpy as np
+import importlib
+import tqdm
+import copy
+import cv2
+import glob
+import imageio
+# common utils
+from utils.commons.hparams import hparams, set_hparams
+from utils.commons.tensor_utils import move_to_cuda, convert_to_tensor
+from utils.commons.ckpt_utils import load_ckpt, get_last_checkpoint
+# 3DMM-related utils
+from deep_3drecon.deep_3drecon_models.bfm import ParametricFaceModel
+from data_util.face3d_helper import Face3DHelper
+from data_gen.utils.process_image.fit_3dmm_landmark import fit_3dmm_for_a_image
+from data_gen.utils.process_video.fit_3dmm_landmark import fit_3dmm_for_a_video
+from data_gen.utils.process_video.extract_segment_imgs import decode_segmap_mask_from_image
+from deep_3drecon.secc_renderer import SECC_Renderer
+from data_gen.eg3d.convert_to_eg3d_convention import get_eg3d_convention_camera_pose_intrinsic
+from data_gen.runs.binarizer_nerf import get_lip_rect
+# Face Parsing 
+from data_gen.utils.mp_feature_extractors.mp_segmenter import MediapipeSegmenter
+from data_gen.utils.process_video.extract_segment_imgs import inpaint_torso_job, extract_background
+# other inference utils
+from inference.infer_utils import mirror_index, load_img_to_512_hwc_array, load_img_to_normalized_512_bchw_tensor
+from inference.infer_utils import smooth_camera_sequence, smooth_features_xd
+from inference.edit_secc import blink_eye_for_secc, hold_eye_opened_for_secc
+from modules.commons.loralib.utils import mark_only_lora_as_trainable
+from utils.nn.model_utils import num_params
+import lpips
+from utils.commons.meters import AvgrageMeter
+meter = AvgrageMeter()
+from torch.utils.tensorboard import SummaryWriter
+class LoRATrainer(nn.Module):
+    def __init__(self, inp):
+        super().__init__()
+        self.inp = inp
+        self.lora_args = {'lora_mode': inp['lora_mode'], 'lora_r': inp['lora_r']}
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        head_model_dir = inp['head_ckpt']
+        torso_model_dir = inp['torso_ckpt']
+        model_dir = torso_model_dir if torso_model_dir != '' else head_model_dir
+        cmd = f"cp {os.path.join(model_dir, 'config.yaml')} {self.inp['work_dir']}"
+        print(cmd)
+        os.system(cmd)
+        with open(os.path.join(self.inp['work_dir'], 'config.yaml'), "a") as f:
+            f.write(f"\nlora_r: {inp['lora_r']}")
+            f.write(f"\nlora_mode: {inp['lora_mode']}")
+            f.write(f"\n")
+        self.secc2video_model = self.load_secc2video(model_dir)
+        self.secc2video_model.to(device).eval()
+        self.seg_model = MediapipeSegmenter()
+        self.secc_renderer = SECC_Renderer(512)
+        self.face3d_helper = Face3DHelper(use_gpu=True, keypoint_mode='lm68')
+        self.mp_face3d_helper = Face3DHelper(use_gpu=True, keypoint_mode='mediapipe')
+        # self.camera_selector = KNearestCameraSelector()
+        self.load_training_data(inp)
+    def load_secc2video(self, model_dir):
+        inp = self.inp
+        from modules.real3d.secc_img2plane_torso import OSAvatarSECC_Img2plane, OSAvatarSECC_Img2plane_Torso
+        hp = set_hparams(f"{model_dir}/config.yaml", print_hparams=False, global_hparams=True)
+        hp['htbsr_head_threshold'] = 1.0
+        self.neural_rendering_resolution = hp['neural_rendering_resolution']
+        if 'torso' in hp['task_cls'].lower():
+            self.torso_mode = True
+            model = OSAvatarSECC_Img2plane_Torso(hp=hp, lora_args=self.lora_args)
+        else:
+            self.torso_mode = False
+            model = OSAvatarSECC_Img2plane(hp=hp, lora_args=self.lora_args)
+        mark_only_lora_as_trainable(model, bias='none')
+        lora_ckpt_path = os.path.join(inp['work_dir'], 'checkpoint.ckpt')
+        if os.path.exists(lora_ckpt_path):
+            self.learnable_triplane = nn.Parameter(torch.zeros([1, 3, model.triplane_hid_dim*model.triplane_depth, 256, 256]).float().cuda(), requires_grad=True)
+            model._last_cano_planes = self.learnable_triplane
+            load_ckpt(model, lora_ckpt_path, model_name='model', strict=False)   
+        else:
+            load_ckpt(model, f"{model_dir}", model_name='model', strict=False)   
+            
+        num_params(model)
+        self.model = model 
+        return model
+    def load_training_data(self, inp):
+        video_id = inp['video_id']
+        if video_id.endswith((".mp4", ".png", ".jpg", ".jpeg")):
+            # If input video is not GeneFace training videos, convert it into GeneFace convention
+            video_id_ = video_id
+            video_id = os.path.basename(video_id)[:-4]
+            inp['video_id'] = video_id
+            target_video_path = f'data/raw/videos/{video_id}.mp4'
+            if not os.path.exists(target_video_path):
+                print(f"| Copying video to {target_video_path}")
+                os.makedirs(os.path.dirname(target_video_path), exist_ok=True)
+                cmd = f"ffmpeg -i {video_id_} -vf fps=25,scale=w=512:h=512 -qmin 1 -q:v 1 -y {target_video_path}"
+                print(f"| {cmd}")
+                os.system(cmd)
+        target_video_path = f'data/raw/videos/{video_id}.mp4'
+        print(f"| Copy source video into work dir: {self.inp['work_dir']}")
+        os.system(f"cp {target_video_path} {self.inp['work_dir']}")
+        # check head_img path
+        head_img_pattern = f'data/processed/videos/{video_id}/head_imgs/*.png'
+        head_img_names = sorted(glob.glob(head_img_pattern))
+        if len(head_img_names) == 0:
+            # extract head_imgs
+            head_img_dir = os.path.dirname(head_img_pattern)
+            print(f"| Pre-extracted head_imgs not found, try to extract and save to {head_img_dir}, this may take a while...")
+            gt_img_dir = f"data/processed/videos/{video_id}/gt_imgs"
+            os.makedirs(gt_img_dir, exist_ok=True)
+            target_video_path = f'data/raw/videos/{video_id}.mp4'
+            cmd = f"ffmpeg -i {target_video_path} -vf fps=25,scale=w=512:h=512 -qmin 1 -q:v 1 -start_number 0 -y {gt_img_dir}/%08d.jpg"
+            print(f"| {cmd}")
+            os.system(cmd)
+            # extract image, segmap, and background
+            cmd = f"python data_gen/utils/process_video/extract_segment_imgs.py --ds_name=nerf --vid_dir={target_video_path}" 
+            print(f"| {cmd}")
+            os.system(cmd)
+            print("| Head images Extracted!")
+        num_samples = len(head_img_names)
+        npy_name = f"data/processed/videos/{video_id}/coeff_fit_mp_for_lora.npy"
+        if os.path.exists(npy_name):
+            coeff_dict = np.load(npy_name, allow_pickle=True).tolist()
+        else:
+            print(f"| Pre-extracted 3DMM coefficient not found, try to extract and save to {npy_name}, this may take a while...")
+            coeff_dict = fit_3dmm_for_a_video(f'data/raw/videos/{video_id}.mp4', save=False)
+            os.makedirs(os.path.dirname(npy_name), exist_ok=True)
+            np.save(npy_name, coeff_dict)
+        ids = convert_to_tensor(coeff_dict['id']).reshape([-1,80]).cuda()
+        exps = convert_to_tensor(coeff_dict['exp']).reshape([-1,64]).cuda()
+        eulers = convert_to_tensor(coeff_dict['euler']).reshape([-1,3]).cuda()
+        trans = convert_to_tensor(coeff_dict['trans']).reshape([-1,3]).cuda()
+        WH = 512 # now we only support 512x512
+        lm2ds = WH * self.face3d_helper.reconstruct_lm2d(ids, exps, eulers, trans).cpu().numpy()
+        lip_rects = [get_lip_rect(lm2ds[i], WH, WH) for i in range(len(lm2ds))]
+        kps = self.face3d_helper.reconstruct_lm2d(ids, exps, eulers, trans).cuda()
+        kps = (kps-0.5) / 0.5 # rescale to -1~1
+        kps = torch.cat([kps, torch.zeros([*kps.shape[:-1], 1]).cuda()], dim=-1)
+        camera_ret = get_eg3d_convention_camera_pose_intrinsic({'euler': torch.tensor(coeff_dict['euler']).reshape([-1,3]), 'trans': torch.tensor(coeff_dict['trans']).reshape([-1,3])})
+        c2w, intrinsics = camera_ret['c2w'], camera_ret['intrinsics']
+        cameras = torch.tensor(np.concatenate([c2w.reshape([-1,16]), intrinsics.reshape([-1,9])], axis=-1)).cuda()
+        camera_smo_ksize = 7
+        cameras = smooth_camera_sequence(cameras.cpu().numpy(), kernel_size=camera_smo_ksize) # [T, 25]
+        cameras = torch.tensor(cameras).cuda()
+        zero_eulers = eulers * 0
+        zero_trans = trans * 0
+        _, cano_secc_color = self.secc_renderer(ids[0:1], exps[0:1]*0, zero_eulers[0:1], zero_trans[0:1])
+        src_idx = 0
+        _, src_secc_color = self.secc_renderer(ids[0:1], exps[src_idx:src_idx+1], zero_eulers[0:1], zero_trans[0:1])
+        drv_secc_colors = [None for _ in range(len(exps))]
+        drv_head_imgs = [None for _ in range(len(exps))]
+        drv_torso_imgs = [None for _ in range(len(exps))]
+        drv_com_imgs = [None for _ in range(len(exps))]
+        segmaps = [None for _ in range(len(exps))]
+        img_name = f'data/processed/videos/{video_id}/bg.jpg'
+        bg_img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
+        ds = {
+            'id': ids.cuda().float(),
+            'exps': exps.cuda().float(),
+            'eulers': eulers.cuda().float(),
+            'trans': trans.cuda().float(),
+            'cano_secc_color': cano_secc_color.cuda().float(),
+            'src_secc_color': src_secc_color.cuda().float(),
+            'cameras': cameras.float(),
+            'video_id': video_id,
+            'lip_rects': lip_rects,
+            'head_imgs': drv_head_imgs,
+            'torso_imgs': drv_torso_imgs,
+            'com_imgs': drv_com_imgs,
+            'bg_img': bg_img,
+            'segmaps': segmaps,
+            'kps': kps,
+        }
+        self.ds = ds
+        return ds
+    
+    def training_loop(self, inp):
+        trainer = self
+        video_id = self.ds['video_id']
+        lora_params = [p for k, p in self.secc2video_model.named_parameters() if 'lora_' in k]
+        self.criterion_lpips = lpips.LPIPS(net='alex',lpips=True).cuda()
+        self.logger = SummaryWriter(log_dir=inp['work_dir'])
+        if not hasattr(self, 'learnable_triplane'):
+            src_idx = 0 # init triplane from the first frame's prediction
+            self.learnable_triplane = nn.Parameter(torch.zeros([1, 3, self.secc2video_model.triplane_hid_dim*self.secc2video_model.triplane_depth, 256, 256]).float().cuda(), requires_grad=True)
+            img_name = f'data/processed/videos/{video_id}/head_imgs/{format(src_idx, "08d")}.png'
+            img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float().cuda().float() # [3, H, W]
+            cano_plane = self.secc2video_model.cal_cano_plane(img.unsqueeze(0)) # [1, 3, CD, h, w]
+            self.learnable_triplane.data = cano_plane.data
+            self.secc2video_model._last_cano_planes = self.learnable_triplane
+        if len(lora_params) == 0:
+            self.optimizer = torch.optim.AdamW([self.learnable_triplane], lr=inp['lr_triplane'], weight_decay=0.01, betas=(0.9,0.98))
+        else:
+            self.optimizer = torch.optim.Adam(lora_params, lr=inp['lr'], betas=(0.9,0.98))
+            self.optimizer.add_param_group({
+                'params': [self.learnable_triplane],
+                'lr': inp['lr_triplane'],
+                'betas': (0.9, 0.98)
+            })
+        
+        ids = self.ds['id']
+        exps = self.ds['exps']
+        zero_eulers = self.ds['eulers']*0
+        zero_trans = self.ds['trans']*0
+        num_updates = inp['max_updates']
+        batch_size = inp['batch_size'] # 1 for lower gpu mem usage
+        num_samples = len(self.ds['cameras'])
+        init_plane = self.learnable_triplane.detach().clone()
+        if num_samples <= 5:
+            lambda_reg_triplane = 1.0
+        elif num_samples <= 250:
+            lambda_reg_triplane = 0.1
+        else:
+            lambda_reg_triplane = 0.
+        for i_step in tqdm.trange(num_updates+1,desc="training lora..."):
+            milestone_steps = []
+            # milestone_steps = [100, 200, 500]
+            if i_step % 2000 == 0 or i_step in milestone_steps:
+                trainer.test_loop(inp, step=i_step)
+                if i_step != 0:
+                    filepath = os.path.join(inp['work_dir'], f"model_ckpt_steps_{i_step}.ckpt") 
+                    checkpoint = self.dump_checkpoint(inp)
+                    tmp_path = str(filepath) + ".part"
+                    torch.save(checkpoint, tmp_path, _use_new_zipfile_serialization=False)
+                    os.replace(tmp_path, filepath)
+                
+            drv_idx = [random.randint(0, num_samples-1) for _ in range(batch_size)]
+            drv_secc_colors = []
+            gt_imgs = []
+            head_imgs = []
+            segmaps_0 = []
+            segmaps = []
+            torso_imgs = []
+            drv_lip_rects = []
+            kp_src = []
+            kp_drv = []
+            for di in drv_idx:
+                # 读取target image
+                if self.torso_mode:
+                    if self.ds['com_imgs'][di] is None:
+                        # img_name = f'data/processed/videos/{video_id}/gt_imgs/{format(di, "08d")}.jpg'
+                        img_name = f'data/processed/videos/{video_id}/com_imgs/{format(di, "08d")}.jpg'
+                        img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
+                        self.ds['com_imgs'][di] = img
+                    gt_imgs.append(self.ds['com_imgs'][di])
+                else:
+                    if self.ds['head_imgs'][di] is None:
+                        img_name = f'data/processed/videos/{video_id}/head_imgs/{format(di, "08d")}.png'
+                        img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
+                        self.ds['head_imgs'][di] = img
+                    gt_imgs.append(self.ds['head_imgs'][di])
+                if self.ds['head_imgs'][di] is None:
+                    img_name = f'data/processed/videos/{video_id}/head_imgs/{format(di, "08d")}.png'
+                    img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
+                    self.ds['head_imgs'][di] = img
+                head_imgs.append(self.ds['head_imgs'][di])
+                # 使用第一帧的torso作为face v2v的输入
+                if self.ds['torso_imgs'][0] is None:
+                    img_name = f'data/processed/videos/{video_id}/inpaint_torso_imgs/{format(0, "08d")}.png'
+                    img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
+                    self.ds['torso_imgs'][0] = img
+                torso_imgs.append(self.ds['torso_imgs'][0])
+                # 所以segmap也用第一帧的了
+                if self.ds['segmaps'][0] is None:
+                    img_name = f'data/processed/videos/{video_id}/segmaps/{format(0, "08d")}.png'
+                    seg_img = cv2.imread(img_name)[:,:, ::-1]
+                    segmap = torch.from_numpy(decode_segmap_mask_from_image(seg_img)) # [6, H, W]
+                    self.ds['segmaps'][0] = segmap
+                segmaps_0.append(self.ds['segmaps'][0])
+                if self.ds['segmaps'][di] is None:
+                    img_name = f'data/processed/videos/{video_id}/segmaps/{format(di, "08d")}.png'
+                    seg_img = cv2.imread(img_name)[:,:, ::-1]
+                    segmap = torch.from_numpy(decode_segmap_mask_from_image(seg_img)) # [6, H, W]
+                    self.ds['segmaps'][di] = segmap
+                segmaps.append(self.ds['segmaps'][di])
+                _, secc_color = self.secc_renderer(ids[0:1], exps[di:di+1], zero_eulers[0:1], zero_trans[0:1])
+                drv_secc_colors.append(secc_color)
+                drv_lip_rects.append(self.ds['lip_rects'][di])
+                kp_src.append(self.ds['kps'][0])
+                kp_drv.append(self.ds['kps'][di])
+            bg_img = self.ds['bg_img'].unsqueeze(0).repeat([batch_size, 1, 1, 1]).cuda()
+            ref_torso_imgs = torch.stack(torso_imgs).float().cuda()
+            kp_src = torch.stack(kp_src).float().cuda()
+            kp_drv = torch.stack(kp_drv).float().cuda()
+            segmaps = torch.stack(segmaps).float().cuda()
+            segmaps_0 = torch.stack(segmaps_0).float().cuda()
+            tgt_imgs = torch.stack(gt_imgs).float().cuda()
+            head_imgs = torch.stack(head_imgs).float().cuda()
+            drv_secc_color = torch.cat(drv_secc_colors)
+            cano_secc_color = self.ds['cano_secc_color'].repeat([batch_size, 1, 1, 1])
+            src_secc_color = self.ds['src_secc_color'].repeat([batch_size, 1, 1, 1])
+            cond = {'cond_cano': cano_secc_color,'cond_src': src_secc_color, 'cond_tgt': drv_secc_color,
+                    'ref_torso_img': ref_torso_imgs, 'bg_img': bg_img, 
+                    'segmap': segmaps_0, # v2v使用第一帧的torso作为source image来warp
+                    'kp_s': kp_src, 'kp_d': kp_drv}
+            camera = self.ds['cameras'][drv_idx]
+            gen_output = self.secc2video_model.forward(img=None, camera=camera, cond=cond, ret={}, cache_backbone=False, use_cached_backbone=True)
+            pred_imgs = gen_output['image']
+            pred_imgs_raw = gen_output['image_raw']
+
+            losses = {}
+            loss_weights = {
+                'v2v_occlusion_reg_l1_loss': 0.001, # loss for face_vid2vid-based torso
+                'v2v_occlusion_2_reg_l1_loss': 0.001, # loss for face_vid2vid-based torso
+                'v2v_occlusion_2_weights_entropy_loss': hparams['lam_occlusion_weights_entropy'], # loss for face_vid2vid-based torso
+                'density_weight_l2_loss': 0.01, # supervised density
+                'density_weight_entropy_loss': 0.001, # keep the density change sharp
+                'mse_loss': 1.,
+                'head_mse_loss': 0.2, # loss on neural rendering low-reso pred_img
+                'lip_mse_loss': 1.0,
+                'lpips_loss': 0.5,
+                'head_lpips_loss': 0.1,
+                'lip_lpips_loss': 1.0, # make the teeth more clear
+                'blink_reg_loss': 0.003, # increase it when you find head shake while blinking; decrease it when you find the eye cannot closed.
+                'triplane_reg_loss': lambda_reg_triplane,
+                'secc_reg_loss': 0.01, # used to reduce flicking
+            }
+
+            occlusion_reg_l1 = gen_output.get("losses", {}).get('facev2v/occlusion_reg_l1', 0.)
+            occlusion_2_reg_l1 = gen_output.get("losses", {}).get('facev2v/occlusion_2_reg_l1', 0.)
+            occlusion_2_weights_entropy = gen_output.get("losses", {}).get('facev2v/occlusion_2_weights_entropy', 0.)
+            losses['v2v_occlusion_reg_l1_loss'] = occlusion_reg_l1
+            losses['v2v_occlusion_2_reg_l1_loss'] = occlusion_2_reg_l1
+            losses['v2v_occlusion_2_weights_entropy_loss'] = occlusion_2_weights_entropy
+
+            # Weights Reg loss in torso
+            neural_rendering_reso = self.neural_rendering_resolution
+            alphas = gen_output['weights_img'].clamp(1e-5, 1 - 1e-5)
+            loss_weights_entropy = torch.mean(- alphas * torch.log2(alphas) - (1 - alphas) * torch.log2(1 - alphas))
+            mv_head_masks = segmaps[:, [1,3,5]].sum(dim=1)
+            mv_head_masks_raw = F.interpolate(mv_head_masks.unsqueeze(1), size=(neural_rendering_reso,neural_rendering_reso)).squeeze(1)
+            face_mask = mv_head_masks_raw.bool().unsqueeze(1)
+            nonface_mask = ~ face_mask
+            loss_weights_l2_loss = (alphas[nonface_mask]-0).pow(2).mean() + (alphas[face_mask]-1).pow(2).mean()
+            losses['density_weight_l2_loss'] = loss_weights_l2_loss
+            losses['density_weight_entropy_loss'] = loss_weights_entropy
+
+            mse_loss = (pred_imgs - tgt_imgs).abs().mean()
+            head_mse_loss = (pred_imgs_raw - F.interpolate(head_imgs, size=(neural_rendering_reso,neural_rendering_reso), mode='bilinear', antialias=True)).abs().mean()
+            lpips_loss = self.criterion_lpips(pred_imgs, tgt_imgs).mean()
+            head_lpips_loss = self.criterion_lpips(pred_imgs_raw, F.interpolate(head_imgs, size=(neural_rendering_reso,neural_rendering_reso), mode='bilinear', antialias=True)).mean()
+            lip_mse_loss = 0
+            lip_lpips_loss = 0
+            for i in range(len(drv_idx)):
+                xmin, xmax, ymin, ymax = drv_lip_rects[i]
+                lip_tgt_imgs = tgt_imgs[i:i+1,:, ymin:ymax,xmin:xmax].contiguous()
+                lip_pred_imgs = pred_imgs[i:i+1,:, ymin:ymax,xmin:xmax].contiguous()
+                try:
+                    lip_mse_loss = lip_mse_loss + (lip_pred_imgs - lip_tgt_imgs).abs().mean()
+                    lip_lpips_loss = lip_lpips_loss + self.criterion_lpips(lip_pred_imgs, lip_tgt_imgs).mean()
+                except: pass 
+            losses['mse_loss'] = mse_loss
+            losses['head_mse_loss'] = head_mse_loss
+            losses['lpips_loss'] = lpips_loss
+            losses['head_lpips_loss'] = head_lpips_loss
+            losses['lip_mse_loss'] = lip_mse_loss
+            losses['lip_lpips_loss'] = lip_lpips_loss   
+
+            # eye blink reg loss
+            if i_step % 4 == 0:
+                blink_secc_lst1 = []
+                blink_secc_lst2 = []
+                blink_secc_lst3 = []
+                for i in range(len(drv_secc_color)):
+                    secc = drv_secc_color[i]
+                    blink_percent1 = random.random() * 0.5 # 0~0.5
+                    blink_percent3 = 0.5 + random.random() * 0.5 # 0.5~1.0
+                    blink_percent2 = (blink_percent1 + blink_percent3)/2
+                    try:
+                        out_secc1 = blink_eye_for_secc(secc, blink_percent1).to(secc.device)
+                        out_secc2 = blink_eye_for_secc(secc, blink_percent2).to(secc.device)
+                        out_secc3 = blink_eye_for_secc(secc, blink_percent3).to(secc.device)
+                    except:
+                        print("blink eye for secc failed, use original secc")
+                        out_secc1 = copy.deepcopy(secc)
+                        out_secc2 = copy.deepcopy(secc)
+                        out_secc3 = copy.deepcopy(secc)
+                    blink_secc_lst1.append(out_secc1)
+                    blink_secc_lst2.append(out_secc2)
+                    blink_secc_lst3.append(out_secc3)
+                src_secc_color1 = torch.stack(blink_secc_lst1)
+                src_secc_color2 = torch.stack(blink_secc_lst2)
+                src_secc_color3 = torch.stack(blink_secc_lst3)
+            blink_cond1 = {'cond_cano': cano_secc_color, 'cond_src': src_secc_color, 'cond_tgt': src_secc_color1}
+            blink_cond2 = {'cond_cano': cano_secc_color, 'cond_src': src_secc_color, 'cond_tgt': src_secc_color2}
+            blink_cond3 = {'cond_cano': cano_secc_color, 'cond_src': src_secc_color, 'cond_tgt': src_secc_color3}
+            blink_secc_plane1 = self.model.cal_secc_plane(blink_cond1)
+            blink_secc_plane2 = self.model.cal_secc_plane(blink_cond2)
+            blink_secc_plane3 = self.model.cal_secc_plane(blink_cond3)
+            interpolate_blink_secc_plane = (blink_secc_plane1 + blink_secc_plane3) / 2
+            blink_reg_loss = torch.nn.functional.l1_loss(blink_secc_plane2, interpolate_blink_secc_plane)
+            losses['blink_reg_loss'] = blink_reg_loss
+
+            # Triplane Reg loss
+            triplane_reg_loss = (self.learnable_triplane - init_plane).abs().mean()
+            losses['triplane_reg_loss'] = triplane_reg_loss
+
+
+            ref_id = self.ds['id'][0:1]
+            secc_pertube_randn_scale = hparams['secc_pertube_randn_scale']
+            perturbed_id = ref_id + torch.randn_like(ref_id) * secc_pertube_randn_scale
+            drv_exp = self.ds['exps'][drv_idx]
+            perturbed_exp = drv_exp + torch.randn_like(drv_exp) * secc_pertube_randn_scale
+            zero_euler = torch.zeros([len(drv_idx), 3], device=ref_id.device, dtype=ref_id.dtype)
+            zero_trans = torch.zeros([len(drv_idx), 3], device=ref_id.device, dtype=ref_id.dtype)
+            perturbed_secc = self.secc_renderer(perturbed_id, perturbed_exp, zero_euler, zero_trans)[1]
+            secc_reg_loss = torch.nn.functional.l1_loss(drv_secc_color, perturbed_secc)
+            losses['secc_reg_loss'] = secc_reg_loss
+
+
+            total_loss = sum([loss_weights[k] * v for k, v in losses.items() if isinstance(v, torch.Tensor) and v.requires_grad])
+            # Update weights
+            self.optimizer.zero_grad()
+            total_loss.backward()
+            self.learnable_triplane.grad.data = self.learnable_triplane.grad.data * self.learnable_triplane.numel()
+            self.optimizer.step()
+            meter.update(total_loss.item())
+            if i_step % 10 == 0:
+                log_line = f"Iter {i_step+1}: total_loss={meter.avg} "
+                for k, v in losses.items():
+                    log_line = log_line + f" {k}={v.item()}, "
+                    self.logger.add_scalar(f"train/{k}", v.item(), i_step)
+                print(log_line)
+                meter.reset()
+    @torch.no_grad()
+    def test_loop(self, inp, step=''):
+        self.model.eval()
+        # coeff_dict = np.load('data/processed/videos/Lieu/coeff_fit_mp_for_lora.npy', allow_pickle=True).tolist()
+        # drv_exps = torch.tensor(coeff_dict['exp']).cuda().float()
+        drv_exps = self.ds['exps']
+        zero_eulers = self.ds['eulers']*0
+        zero_trans = self.ds['trans']*0
+        batch_size = 1
+        num_samples = len(self.ds['cameras'])
+        video_writer = imageio.get_writer(os.path.join(inp['work_dir'], f'val_step{step}.mp4'), fps=25)
+        total_iters = min(num_samples, 250)
+        video_id = inp['video_id']
+        for i in tqdm.trange(total_iters,desc="testing lora..."):
+            drv_idx = [i]
+            drv_secc_colors = []
+            gt_imgs = []
+            segmaps = []
+            torso_imgs = []
+            drv_lip_rects = []
+            kp_src = []
+            kp_drv = []
+            for di in drv_idx:
+                # 读取target image
+                if self.torso_mode:
+                    if self.ds['com_imgs'][di] is None:
+                        img_name = f'data/processed/videos/{video_id}/com_imgs/{format(di, "08d")}.jpg'
+                        img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
+                        self.ds['com_imgs'][di] = img
+                    gt_imgs.append(self.ds['com_imgs'][di])
+                else:
+                    if self.ds['head_imgs'][di] is None:
+                        img_name = f'data/processed/videos/{video_id}/head_imgs/{format(di, "08d")}.png'
+                        img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
+                        self.ds['head_imgs'][di] = img
+                    gt_imgs.append(self.ds['head_imgs'][di])
+                # 使用第一帧的torso作为face v2v的输入 
+                if self.ds['torso_imgs'][0] is None:
+                    img_name = f'data/processed/videos/{video_id}/inpaint_torso_imgs/{format(0, "08d")}.png'
+                    img = torch.tensor(cv2.imread(img_name)[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
+                    self.ds['torso_imgs'][0] = img
+                torso_imgs.append(self.ds['torso_imgs'][0])
+                # 所以segmap也用第一帧的了
+                if self.ds['segmaps'][0] is None:
+                    img_name = f'data/processed/videos/{video_id}/segmaps/{format(0, "08d")}.png'
+                    seg_img = cv2.imread(img_name)[:,:, ::-1]
+                    segmap = torch.from_numpy(decode_segmap_mask_from_image(seg_img)) # [6, H, W]
+                    self.ds['segmaps'][0] = segmap
+                segmaps.append(self.ds['segmaps'][0])
+                drv_lip_rects.append(self.ds['lip_rects'][di])
+                kp_src.append(self.ds['kps'][0])
+                kp_drv.append(self.ds['kps'][di])
+            bg_img = self.ds['bg_img'].unsqueeze(0).repeat([batch_size, 1, 1, 1]).cuda()
+            ref_torso_imgs = torch.stack(torso_imgs).float().cuda()
+            kp_src = torch.stack(kp_src).float().cuda()
+            kp_drv = torch.stack(kp_drv).float().cuda()
+            segmaps = torch.stack(segmaps).float().cuda()
+            tgt_imgs = torch.stack(gt_imgs).float().cuda()
+            for di in drv_idx:
+                _, secc_color = self.secc_renderer(self.ds['id'][0:1], drv_exps[di:di+1], zero_eulers[0:1], zero_trans[0:1])
+                drv_secc_colors.append(secc_color)
+            drv_secc_color = torch.cat(drv_secc_colors)
+            cano_secc_color = self.ds['cano_secc_color'].repeat([batch_size, 1, 1, 1])
+            src_secc_color = self.ds['src_secc_color'].repeat([batch_size, 1, 1, 1])
+            cond = {'cond_cano': cano_secc_color,'cond_src': src_secc_color, 'cond_tgt': drv_secc_color,
+                    'ref_torso_img': ref_torso_imgs, 'bg_img': bg_img, 'segmap': segmaps,
+                    'kp_s': kp_src, 'kp_d': kp_drv}
+            camera = self.ds['cameras'][drv_idx]
+            gen_output = self.secc2video_model.forward(img=None, camera=camera, cond=cond, ret={}, cache_backbone=False, use_cached_backbone=True)
+            pred_img = gen_output['image']
+            pred_img = ((pred_img.permute(0, 2, 3, 1) + 1)/2 * 255).int().cpu().numpy().astype(np.uint8)
+            video_writer.append_data(pred_img[0])
+        video_writer.close()
+        self.model.train()
+
+    def masked_error_loss(self, img_pred, img_gt, mask, unmasked_weight=0.1, mode='l1'):
+        # 对raw图像，因为deform的原因背景没法全黑，导致这部分mse过高，我们将其mask掉，只计算人脸部分
+        masked_weight = 1.0
+        weight_mask = mask.float() * masked_weight + (~mask).float() * unmasked_weight
+        if mode == 'l1':
+            error = (img_pred - img_gt).abs().sum(dim=1) * weight_mask
+        else:
+            error = (img_pred - img_gt).pow(2).sum(dim=1) * weight_mask
+        error.clamp_(0, max(0.5, error.quantile(0.8).item())) # clamp掉较高loss的pixel，避免姿态没对齐的pixel导致的异常值占主导影响训练
+        loss = error.mean()
+        return loss
+    
+    def dilate(self, bin_img, ksize=5, mode='max_pool'):
+        """
+        mode: max_pool or avg_pool
+        """
+        # bin_img, [1, h, w]
+        pad = (ksize-1)//2
+        bin_img = F.pad(bin_img, pad=[pad,pad,pad,pad], mode='reflect')
+        if mode == 'max_pool':
+            out = F.max_pool2d(bin_img, kernel_size=ksize, stride=1, padding=0)
+        else:
+            out = F.avg_pool2d(bin_img, kernel_size=ksize, stride=1, padding=0)
+        return out
+        
+    def dilate_mask(self, mask, ksize=21):
+        mask = self.dilate(mask, ksize=ksize, mode='max_pool')
+        return mask
+    
+    def set_unmasked_to_black(self, img, mask):
+        out_img = img * mask.float() - (~mask).float() # -1 denotes black
+        return out_img
+    
+    def dump_checkpoint(self, inp):
+        checkpoint = {}
+        # save optimizers
+        optimizer_states = []
+        self.optimizers = [self.optimizer]
+        for i, optimizer in enumerate(self.optimizers):
+            if optimizer is not None:
+                state_dict = optimizer.state_dict()
+                state_dict = {k.replace('_orig_mod.', ''): v for k, v in state_dict.items()}
+                optimizer_states.append(state_dict)
+        checkpoint['optimizer_states'] = optimizer_states
+        state_dict = {
+            'model': self.model.state_dict(),
+            'learnable_triplane': self.model.state_dict()['_last_cano_planes'],
+        }
+        del state_dict['model']['_last_cano_planes']
+        checkpoint['state_dict'] = state_dict
+        checkpoint['lora_args'] = self.lora_args
+        person_ds = {}
+        video_id = inp['video_id']
+        img_name = f'data/processed/videos/{video_id}/gt_imgs/{format(0, "08d")}.jpg'
+        gt_img = torch.tensor(cv2.resize(cv2.imread(img_name), (512, 512))[..., ::-1] / 127.5 - 1).permute(2,0,1).float() # [3, H, W]
+        person_ds['gt_img'] = gt_img.reshape([1, 3, 512, 512])
+        person_ds['id'] = self.ds['id'].cpu().reshape([1, 80])
+        person_ds['src_kp'] = self.ds['kps'][0].cpu()
+        person_ds['video_id'] = inp['video_id']
+        checkpoint['person_ds'] = person_ds
+        return checkpoint
+if __name__ == '__main__':
+    import argparse, glob, tqdm
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--head_ckpt", default='') # checkpoints/0729_th1kh/secc_img2plane checkpoints/0720_img2planes/secc_img2plane_two_stage
+    # parser.add_argument("--torso_ckpt", default='checkpoints/240210_real3dportrait_orig/secc2plane_torso_orig') # checkpoints/0729_th1kh/secc_img2plane checkpoints/0720_img2planes/secc_img2plane_two_stage
+    parser.add_argument("--torso_ckpt", default='checkpoints/mimictalk_orig/os_secc2plane_torso') # checkpoints/0729_th1kh/secc_img2plane checkpoints/0720_img2planes/secc_img2plane_two_stage
+    parser.add_argument("--video_id", default='data/raw/examples/GER.mp4', help="identity source, we support (1) already processed <video_id> of GeneFace, (2) video path, (3) image path")
+    parser.add_argument("--work_dir", default=None) 
+    parser.add_argument("--max_updates", default=10000, type=int, help="for video, 2000 is good; for an image, 3~10 is good") 
+    parser.add_argument("--test", action='store_true') 
+    parser.add_argument("--batch_size", default=1, type=int, help="batch size during training, 1 needs 8GB, 2 needs 15GB") 
+    parser.add_argument("--lr", default=0.001) 
+    parser.add_argument("--lr_triplane", default=0.005, help="for video, 0.1; for an image, 0.001; for ablation with_triplane, 0.") 
+    parser.add_argument("--lora_r", default=2, type=int, help="width of lora unit") 
+    parser.add_argument("--lora_mode", default='secc2plane_sr', help='for video, full; for an image, none')
+
+    args = parser.parse_args()
+    inp = {
+            'head_ckpt': args.head_ckpt,
+            'torso_ckpt': args.torso_ckpt,
+            'video_id': args.video_id,
+            'work_dir': args.work_dir,
+            'max_updates': args.max_updates,
+            'batch_size': args.batch_size,
+            'test': args.test,
+            'lr': float(args.lr),
+            'lr_triplane': float(args.lr_triplane),
+            'lora_mode': args.lora_mode,
+            'lora_r': args.lora_r,
+            }
+    if inp['work_dir'] == None:
+        video_id = os.path.basename(inp['video_id'])[:-4] if inp['video_id'].endswith((".mp4", ".png", ".jpg", ".jpeg")) else inp['video_id']
+        inp['work_dir'] = f'checkpoints_mimictalk/{video_id}'
+    os.makedirs(inp['work_dir'], exist_ok=True)
+    trainer = LoRATrainer(inp)
+    if inp['test']:
+        trainer.test_loop(inp)
+    else:
+        trainer.training_loop(inp)
+        trainer.test_loop(inp)
+    print(" ")

install_guide.md

@@ -0,0 +1,27 @@
+# Prepare the Environment
+[中文文档](./install_guide-zh.md)
+
+This guide is about building a python environment for MimicTalk with Conda (the same as `Real3D-Portrait`).
+
+The following installation process is verified in A100/V100 + CUDA12.1.
+
+# Install Python Packages & CUDA
+```bash
+cd <MimicTalkRoot>
+source <CondaRoot>/bin/activate
+conda create -n mimictalk python=3.9
+conda activate mimictalk
+
+# MMCV for SegFormer network structure
+# other dependencies
+pip install -r docs/prepare_env/requirements.txt -v
+pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu121
+pip install cython
+pip install openmim==0.3.9
+mim install mmcv==2.1.0 # use mim to speed up installation for mmcv
+## build pytorch3d from Github's source code. 
+## It may take a long time (maybe tens of minutes), Proxy is recommended if encountering the time-out problem
+# Before install pytorch3d, you need to install CUDA-12.1 (https://developer.nvidia.com/cuda-toolkit-archive) and make sure /usr/local/cuda points to the `cuda-12.1` directory
+pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable"
+
+```

requirements.txt

@@ -0,0 +1,79 @@
+Cython
+numpy # ==1.23.0
+numba==0.56.4
+pandas
+transformers
+scipy==1.11.1 # required by cal_fid. https://github.com/mseitzer/pytorch-fid/issues/103
+scikit-learn
+scikit-image
+# tensorflow # you can flexible it, this is gpu version
+tensorboard
+tensorboardX
+python_speech_features
+resampy
+opencv_python
+face_alignment
+matplotlib
+configargparse
+librosa==0.9.2
+praat-parselmouth # ==0.4.3
+trimesh
+kornia==0.5.0
+PyMCubes
+lpips
+setuptools # ==59.5.0
+ffmpeg-python
+moviepy
+dearpygui
+ninja
+pyaudio # for extract esperanto
+mediapipe
+protobuf
+decord
+soundfile
+pillow
+# torch # it's better to install torch with conda 
+av
+timm
+pretrainedmodels
+faiss-cpu # for fast nearest camera pose retriveal
+einops
+# mmcv # use mim install is faster
+
+# conditional flow matching
+beartype
+torchode
+torchdiffeq
+
+# tts
+cython
+textgrid
+pyloudnorm
+websocket-client
+pyworld==0.2.1rc0
+pypinyin==0.42.0
+webrtcvad
+torchshow
+
+# cal spk sim
+# s3prl
+# fire
+
+# cal LMD
+# dlib
+
+# debug
+# ipykernel
+
+# lama
+# hydra-core
+# pytorch_lightning
+# setproctitle
+
+# Gradio GUI
+# httpx==0.23.3
+# gradio==4.16.0
+gradio==4.43.0
+httpx==0.23.3
+# gradio_client==0.8.1
+fastapi==0.112.2