files added

.gitignore

@@ -0,0 +1,45 @@
+# PyCharm files
+.idea/
+
+# macOS dir files
+.DS_Store
+
+# VS Code configuration dir
+.vscode/
+
+# Jupyter Notebook cache files
+.ipynb_checkpoints/
+*.ipynb
+
+# Python cache files
+__pycache__/
+
+# folders
+wandb/
+*debug*
+/debug
+/output
+/validation
+/test
+/models/
+/venv/
+/detect_results/
+/temp
+
+# checkpoint files
+*.safetensors
+*.ckpt
+*.pt
+
+# data files
+*.mp4
+*.avi
+*.wav
+*.png
+*.jpg
+*.jpeg
+*.csv
+
+!/latentsync/utils/mask.png
+/checkpoints/
+!/assets/*

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,12 +1,170 @@
----
-title: LatentSync
-emoji: 🌍
-colorFrom: green
-colorTo: yellow
-sdk: gradio
-sdk_version: 5.12.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# LatentSync: Audio Conditioned Latent Diffusion Models for Lip Sync
+
+<div align="center">
+
+[![arXiv](https://img.shields.io/badge/arXiv_paper-2412.09262-b31b1b)](https://arxiv.org/abs/2412.09262)
+[![arXiv](https://img.shields.io/badge/%F0%9F%A4%97%20space-HuggingFace-yellow)](https://huggingface.co/spaces/fffiloni/LatentSync)
+<a href="https://replicate.com/lucataco/latentsync"><img src="https://replicate.com/lucataco/latentsync/badge" alt="Replicate"></a>
+
+</div>
+
+## 📖 Abstract
+
+We present *LatentSync*, an end-to-end lip sync framework based on audio conditioned latent diffusion models without any intermediate motion representation, diverging from previous diffusion-based lip sync methods based on pixel space diffusion or two-stage generation. Our framework can leverage the powerful capabilities of Stable Diffusion to directly model complex audio-visual correlations. Additionally, we found that the diffusion-based lip sync methods exhibit inferior temporal consistency due to the inconsistency in the diffusion process across different frames. We propose *Temporal REPresentation Alignment (TREPA)* to enhance temporal consistency while preserving lip-sync accuracy. TREPA uses temporal representations extracted by large-scale self-supervised video models to align the generated frames with the ground truth frames.
+
+## 🏗️ Framework
+
+<p align="center">
+<img src="assets/framework.png" width=100%>
+<p>
+
+LatentSync uses the [Whisper](https://github.com/openai/whisper) to convert melspectrogram into audio embeddings, which are then integrated into the U-Net via cross-attention layers. The reference and masked frames are channel-wise concatenated with noised latents as the input of U-Net. In the training process, we use a one-step method to get estimated clean latents from predicted noises, which are then decoded to obtain the estimated clean frames. The TREPA, [LPIPS](https://arxiv.org/abs/1801.03924) and [SyncNet](https://www.robots.ox.ac.uk/~vgg/publications/2016/Chung16a/chung16a.pdf) losses are added in the pixel space.
+
+## 🎬 Demo
+
+<table class="center">
+  <tr style="font-weight: bolder;text-align:center;">
+        <td width="50%"><b>Original video</b></td>
+        <td width="50%"><b>Lip-synced video</b></td>
+  </tr>
+  <tr>
+    <td>
+      <video src=https://github.com/user-attachments/assets/ff3a84da-dc9b-498a-950f-5c54f58dd5c5 controls preload></video>
+    </td>
+    <td>
+      <video src=https://github.com/user-attachments/assets/150e00fd-381e-4421-a478-a9ea3d1212a8 controls preload></video>
+    </td>
+  </tr>
+  <tr>
+    <td>
+      <video src=https://github.com/user-attachments/assets/32c830a9-4d7d-4044-9b33-b184d8e11010 controls preload></video>
+    </td>
+    <td>
+      <video src=https://github.com/user-attachments/assets/84e4fe9d-b108-44a4-8712-13a012348145 controls preload></video>
+    </td>
+  </tr>
+  <tr>
+    <td>
+      <video src=https://github.com/user-attachments/assets/7510a448-255a-44ee-b093-a1b98bd3961d controls preload></video>
+    </td>
+    <td>
+      <video src=https://github.com/user-attachments/assets/6150c453-c559-4ae0-bb00-c565f135ff41 controls preload></video>
+    </td>
+  </tr>
+  <tr>
+    <td width=300px>
+      <video src=https://github.com/user-attachments/assets/0f7f9845-68b2-4165-bd08-c7bbe01a0e52 controls preload></video>
+    </td>
+    <td width=300px>
+      <video src=https://github.com/user-attachments/assets/c34fe89d-0c09-4de3-8601-3d01229a69e3 controls preload></video>
+    </td>
+  </tr>
+  <tr>
+    <td>
+      <video src=https://github.com/user-attachments/assets/7ce04d50-d39f-4154-932a-ec3a590a8f64 controls preload></video>
+    </td>
+    <td>
+      <video src=https://github.com/user-attachments/assets/70bde520-42fa-4a0e-b66c-d3040ae5e065 controls preload></video>
+    </td>
+  </tr>
+</table>
+
+(Photorealistic videos are filmed by contracted models, and anime videos are from [VASA-1](https://www.microsoft.com/en-us/research/project/vasa-1/) and [EMO](https://humanaigc.github.io/emote-portrait-alive/))
+
+## 📑 Open-source Plan
+
+- [x] Inference code and checkpoints
+- [x] Data processing pipeline
+- [x] Training code
+
+## 🔧 Setting up the Environment
+
+Install the required packages and download the checkpoints via:
+
+```bash
+source setup_env.sh
+```
+
+If the download is successful, the checkpoints should appear as follows:
+
+```
+./checkpoints/
+|-- latentsync_unet.pt
+|-- latentsync_syncnet.pt
+|-- whisper
+|   `-- tiny.pt
+|-- auxiliary
+|   |-- 2DFAN4-cd938726ad.zip
+|   |-- i3d_torchscript.pt
+|   |-- koniq_pretrained.pkl
+|   |-- s3fd-619a316812.pth
+|   |-- sfd_face.pth
+|   |-- syncnet_v2.model
+|   |-- vgg16-397923af.pth
+|   `-- vit_g_hybrid_pt_1200e_ssv2_ft.pth
+```
+
+These already include all the checkpoints required for latentsync training and inference. If you just want to try inference, you only need to download `latentsync_unet.pt` and `tiny.pt` from our [HuggingFace repo](https://huggingface.co/chunyu-li/LatentSync)
+
+## 🚀 Inference
+
+### 1. Gradio App
+
+Run the Gradio app for inference, which requires about 6.5 GB GPU memory.
+
+```bash
+python gradio_app.py
+```
+
+### 2. Command Line Interface
+
+Run the script for inference, which requires about 6.5 GB GPU memory.
+
+```bash
+./inference.sh
+```
+
+You can change the parameter `guidance_scale` to 1.5 to improve the lip-sync accuracy.
+
+## 🔄 Data Processing Pipeline
+
+The complete data processing pipeline includes the following steps:
+
+1. Remove the broken video files.
+2. Resample the video FPS to 25, and resample the audio to 16000 Hz.
+3. Scene detect via [PySceneDetect](https://github.com/Breakthrough/PySceneDetect).
+4. Split each video into 5-10 second segments.
+5. Remove videos where the face is smaller than 256 $\times$ 256, as well as videos with more than one face.
+6. Affine transform the faces according to the landmarks detected by [face-alignment](https://github.com/1adrianb/face-alignment), then resize to 256 $\times$ 256.
+7. Remove videos with [sync confidence score](https://www.robots.ox.ac.uk/~vgg/publications/2016/Chung16a/chung16a.pdf) lower than 3, and adjust the audio-visual offset to 0.
+8. Calculate [hyperIQA](https://openaccess.thecvf.com/content_CVPR_2020/papers/Su_Blindly_Assess_Image_Quality_in_the_Wild_Guided_by_a_CVPR_2020_paper.pdf) score, and remove videos with scores lower than 40.
+
+Run the script to execute the data processing pipeline:
+
+```bash
+./data_processing_pipeline.sh
+```
+
+You can change the parameter `input_dir` in the script to specify the data directory to be processed. The processed data will be saved in the `high_visual_quality` directory. Each step will generate a new directory to prevent the need to redo the entire pipeline in case the process is interrupted by an unexpected error.
+
+## 🏋️‍♂️ Training U-Net
+
+Before training, you must process the data as described above and download all the checkpoints. We released a pretrained SyncNet with 94% accuracy on the VoxCeleb2 dataset for the supervision of U-Net training. Note that this SyncNet is trained on affine transformed videos, so when using or evaluating this SyncNet, you need to perform affine transformation on the video first (the code of affine transformation is included in the data processing pipeline).
+
+If all the preparations are complete, you can train the U-Net with the following script:
+
+```bash
+./train_unet.sh
+```
+
+You should change the parameters in U-Net config file to specify the data directory, checkpoint save path, and other training hyperparameters.
+
+## 🏋️‍♂️ Training SyncNet
+
+In case you want to train SyncNet on your own datasets, you can run the following script. The data processing pipeline for SyncNet is the same as U-Net. 
+
+```bash
+./train_syncnet.sh
+```
+
+After `validations_steps` training, the loss charts will be saved in `train_output_dir`. They contain both the training and validation loss.

app.py

@@ -1 +1,161 @@
-print("hello")
+import gradio as gr
+from pathlib import Path
+from scripts.inference import main
+from omegaconf import OmegaConf
+import argparse
+from datetime import datetime
+
+CONFIG_PATH = Path("configs/unet/second_stage.yaml")
+CHECKPOINT_PATH = Path("checkpoints/latentsync_unet.pt")
+
+
+def process_video(
+    video_path,
+    audio_path,
+    guidance_scale,
+    inference_steps,
+    seed,
+):
+    # Create the temp directory if it doesn't exist
+    output_dir = Path("./temp")
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Convert paths to absolute Path objects and normalize them
+    video_file_path = Path(video_path)
+    video_path = video_file_path.absolute().as_posix()
+    audio_path = Path(audio_path).absolute().as_posix()
+
+    current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
+    # Set the output path for the processed video
+    output_path = str(
+        output_dir / f"{video_file_path.stem}_{current_time}.mp4"
+    )  # Change the filename as needed
+
+    config = OmegaConf.load(CONFIG_PATH)
+
+    config["run"].update(
+        {
+            "guidance_scale": guidance_scale,
+            "inference_steps": inference_steps,
+        }
+    )
+
+    # Parse the arguments
+    args = create_args(video_path, audio_path, output_path, guidance_scale, seed)
+
+    try:
+        result = main(
+            config=config,
+            args=args,
+        )
+        print("Processing completed successfully.")
+        return output_path  # Ensure the output path is returned
+    except Exception as e:
+        print(f"Error during processing: {str(e)}")
+        raise gr.Error(f"Error during processing: {str(e)}")
+
+
+def create_args(
+    video_path: str, audio_path: str, output_path: str, guidance_scale: float, seed: int
+) -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--inference_ckpt_path", type=str, required=True)
+    parser.add_argument("--video_path", type=str, required=True)
+    parser.add_argument("--audio_path", type=str, required=True)
+    parser.add_argument("--video_out_path", type=str, required=True)
+    parser.add_argument("--guidance_scale", type=float, default=1.0)
+    parser.add_argument("--seed", type=int, default=1247)
+
+    return parser.parse_args(
+        [
+            "--inference_ckpt_path",
+            CHECKPOINT_PATH.absolute().as_posix(),
+            "--video_path",
+            video_path,
+            "--audio_path",
+            audio_path,
+            "--video_out_path",
+            output_path,
+            "--guidance_scale",
+            str(guidance_scale),
+            "--seed",
+            str(seed),
+        ]
+    )
+
+
+# Create Gradio interface
+with gr.Blocks(title="LatentSync Video Processing") as demo:
+    gr.Markdown(
+        """
+    # LatentSync: Audio Conditioned Latent Diffusion Models for Lip Sync
+    Upload a video and audio file to process with LatentSync model.
+
+    <div align="center">
+        <strong>Chunyu Li1,2  Chao Zhang1  Weikai Xu1  Jinghui Xie1,†  Weiguo Feng1
+        Bingyue Peng1  Weiwei Xing2,†</strong>
+    </div>
+
+    <div align="center">
+        <strong>1ByteDance   2Beijing Jiaotong University</strong>
+    </div>
+
+    <div style="display:flex;justify-content:center;column-gap:4px;">
+        <a href="https://github.com/bytedance/LatentSync">
+            <img src='https://img.shields.io/badge/GitHub-Repo-blue'>
+        </a> 
+        <a href="https://arxiv.org/pdf/2412.09262">
+            <img src='https://img.shields.io/badge/ArXiv-Paper-red'>
+        </a>
+    </div>
+    """
+    )
+
+    with gr.Row():
+        with gr.Column():
+            video_input = gr.Video(label="Input Video")
+            audio_input = gr.Audio(label="Input Audio", type="filepath")
+
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    minimum=0.1,
+                    maximum=3.0,
+                    value=1.0,
+                    step=0.1,
+                    label="Guidance Scale",
+                )
+                inference_steps = gr.Slider(
+                    minimum=1, maximum=50, value=20, step=1, label="Inference Steps"
+                )
+
+            with gr.Row():
+                seed = gr.Number(value=1247, label="Random Seed", precision=0)
+
+            process_btn = gr.Button("Process Video")
+
+        with gr.Column():
+            video_output = gr.Video(label="Output Video")
+
+            gr.Examples(
+                examples=[
+                    ["assets/demo1_video.mp4", "assets/demo1_audio.wav"],
+                    ["assets/demo2_video.mp4", "assets/demo2_audio.wav"],
+                    ["assets/demo3_video.mp4", "assets/demo3_audio.wav"],
+                ],
+                inputs=[video_input, audio_input],
+            )
+
+    process_btn.click(
+        fn=process_video,
+        inputs=[
+            video_input,
+            audio_input,
+            guidance_scale,
+            inference_steps,
+            seed,
+        ],
+        outputs=video_output,
+    )
+
+if __name__ == "__main__":
+    demo.launch(inbrowser=True, share=True)

cog.yaml

@@ -0,0 +1,44 @@
+# Configuration for Cog ⚙️
+# Reference: https://cog.run/yaml
+
+build:
+  gpu: true
+  cuda: "12.1"
+  system_packages:
+    - "ffmpeg"
+    - "libgl1"
+  python_version: "3.10.13"
+  python_packages:
+    - "torch==2.2.2"
+    - "torchvision"
+    - "triton==2.2.0"
+    - "diffusers==0.11.1"
+    - "transformers==4.38.0"
+    - "huggingface-hub==0.25.2"
+    - "imageio==2.27.0"
+    - "decord==0.6.0"
+    - "accelerate==0.26.1"
+    - "einops==0.7.0"
+    - "omegaconf==2.3.0"
+    - "safetensors==0.4.2"
+    - "opencv-python==4.9.0.80"
+    - "mediapipe==0.10.11"
+    - "av==11.0.0"
+    - "torch-fidelity==0.3.0"
+    - "torchmetrics==1.3.1"
+    - "python_speech_features==0.6"
+    - "librosa==0.10.1"
+    - "scenedetect==0.6.1"
+    - "ffmpeg-python==0.2.0"
+    - "lpips==0.1.4"
+    - "face-alignment==1.4.1"
+    - "ninja==1.11.1.1"
+    - "pandas==2.0.3"
+    - "numpy==1.24.4"
+    - "xformers==0.0.26"
+
+  run:
+    - curl -o /usr/local/bin/pget -L "https://github.com/replicate/pget/releases/download/v0.8.2/pget_linux_x86_64" && chmod +x /usr/local/bin/pget
+
+# predict.py defines how predictions are run on your model
+predict: "predict.py:Predictor"

configs/audio.yaml

@@ -0,0 +1,23 @@
+audio:
+  num_mels: 80 # Number of mel-spectrogram channels and local conditioning dimensionality
+  rescale: true # Whether to rescale audio prior to preprocessing
+  rescaling_max: 0.9 # Rescaling value
+  use_lws:
+    false # Use LWS (https://github.com/Jonathan-LeRoux/lws) for STFT and phase reconstruction
+    # It"s preferred to set True to use with https://github.com/r9y9/wavenet_vocoder
+    # Does not work if n_ffit is not multiple of hop_size!!
+  n_fft: 800 # Extra window size is filled with 0 paddings to match this parameter
+  hop_size: 200 # For 16000Hz, 200 = 12.5 ms (0.0125 * sample_rate)
+  win_size: 800 # For 16000Hz, 800 = 50 ms (If None, win_size = n_fft) (0.05 * sample_rate)
+  sample_rate: 16000 # 16000Hz (corresponding to librispeech) (sox --i <filename>)
+  frame_shift_ms: null
+  signal_normalization: true
+  allow_clipping_in_normalization: true
+  symmetric_mels: true
+  max_abs_value: 4.0
+  preemphasize: true # whether to apply filter
+  preemphasis: 0.97 # filter coefficient.
+  min_level_db: -100
+  ref_level_db: 20
+  fmin: 55
+  fmax: 7600

configs/scheduler_config.json

@@ -0,0 +1,13 @@
+{
+  "_class_name": "DDIMScheduler",
+  "_diffusers_version": "0.6.0.dev0",
+  "beta_end": 0.012,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.00085,
+  "clip_sample": false,
+  "num_train_timesteps": 1000,
+  "set_alpha_to_one": false,
+  "steps_offset": 1,
+  "trained_betas": null,
+  "skip_prk_steps": true
+}

configs/syncnet/syncnet_16_latent.yaml

@@ -0,0 +1,46 @@
+model:
+  audio_encoder: # input (1, 80, 52)
+    in_channels: 1
+    block_out_channels: [32, 64, 128, 256, 512, 1024]
+    downsample_factors: [[2, 1], 2, 2, 2, 2, [2, 3]]
+    attn_blocks: [0, 0, 0, 0, 0, 0]
+    dropout: 0.0
+  visual_encoder: # input (64, 32, 32)
+    in_channels: 64
+    block_out_channels: [64, 128, 256, 256, 512, 1024]
+    downsample_factors: [2, 2, 2, 1, 2, 2]
+    attn_blocks: [0, 0, 0, 0, 0, 0]
+    dropout: 0.0
+
+ckpt:
+  resume_ckpt_path: ""
+  inference_ckpt_path: ""
+  save_ckpt_steps: 2500
+
+data:
+  train_output_dir: output/syncnet
+  num_val_samples: 1200
+  batch_size: 120 # 40
+  num_workers: 11 # 11
+  latent_space: true
+  num_frames: 16
+  resolution: 256
+  train_fileslist: ""
+  train_data_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/VoxCeleb2/high_visual_quality/train
+  val_fileslist: ""
+  val_data_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/VoxCeleb2/high_visual_quality/val
+  audio_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/mel_new
+  lower_half: false
+  pretrained_audio_model_path: facebook/wav2vec2-large-xlsr-53
+  audio_sample_rate: 16000
+  video_fps: 25
+
+optimizer:
+  lr: 1e-5
+  max_grad_norm: 1.0
+
+run:
+  max_train_steps: 10000000
+  validation_steps: 2500
+  mixed_precision_training: true
+  seed: 42

configs/syncnet/syncnet_16_pixel.yaml

@@ -0,0 +1,45 @@
+model:
+  audio_encoder: # input (1, 80, 52)
+    in_channels: 1
+    block_out_channels: [32, 64, 128, 256, 512, 1024, 2048]
+    downsample_factors: [[2, 1], 2, 2, 1, 2, 2, [2, 3]]
+    attn_blocks: [0, 0, 0, 0, 0, 0, 0]
+    dropout: 0.0
+  visual_encoder: # input (48, 128, 256)
+    in_channels: 48
+    block_out_channels: [64, 128, 256, 256, 512, 1024, 2048, 2048]
+    downsample_factors: [[1, 2], 2, 2, 2, 2, 2, 2, 2]
+    attn_blocks: [0, 0, 0, 0, 0, 0, 0, 0]
+    dropout: 0.0
+
+ckpt:
+  resume_ckpt_path: ""
+  inference_ckpt_path: checkpoints/latentsync_syncnet.pt
+  save_ckpt_steps: 2500
+
+data:
+  train_output_dir: debug/syncnet
+  num_val_samples: 2048
+  batch_size: 128 # 128
+  num_workers: 11 # 11
+  latent_space: false
+  num_frames: 16
+  resolution: 256
+  train_fileslist: /mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/all_data_v6.txt
+  train_data_dir: ""
+  val_fileslist: ""
+  val_data_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/VoxCeleb2/high_visual_quality/val
+  audio_mel_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/mel_new
+  lower_half: true
+  audio_sample_rate: 16000
+  video_fps: 25
+
+optimizer:
+  lr: 1e-5
+  max_grad_norm: 1.0
+
+run:
+  max_train_steps: 10000000
+  validation_steps: 2500
+  mixed_precision_training: true
+  seed: 42

configs/syncnet/syncnet_25_pixel.yaml

@@ -0,0 +1,45 @@
+model:
+  audio_encoder: # input (1, 80, 80)
+    in_channels: 1
+    block_out_channels: [64, 128, 256, 256, 512, 1024]
+    downsample_factors: [2, 2, 2, 2, 2, 2]
+    dropout: 0.0
+  visual_encoder: # input (75, 128, 256)
+    in_channels: 75
+    block_out_channels: [128, 128, 256, 256, 512, 512, 1024, 1024]
+    downsample_factors: [[1, 2], 2, 2, 2, 2, 2, 2, 2]
+    dropout: 0.0
+
+ckpt:
+  resume_ckpt_path: ""
+  inference_ckpt_path: ""
+  save_ckpt_steps: 2500
+
+data:
+  train_output_dir: debug/syncnet
+  num_val_samples: 2048
+  batch_size: 64 # 64
+  num_workers: 11 # 11
+  latent_space: false
+  num_frames: 25
+  resolution: 256
+  train_fileslist: /mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/hdtf_vox_avatars_ads_affine.txt
+  # /mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/hdtf_voxceleb_avatars_affine.txt
+  train_data_dir: ""
+  val_fileslist: /mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/vox_affine_val.txt
+  # /mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/voxceleb_val.txt
+  val_data_dir: ""
+  audio_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/mel
+  lower_half: true
+  pretrained_audio_model_path: facebook/wav2vec2-large-xlsr-53
+  audio_sample_rate: 16000
+  video_fps: 25
+
+optimizer:
+  lr: 1e-5
+  max_grad_norm: 1.0
+
+run:
+  max_train_steps: 10000000
+  mixed_precision_training: true
+  seed: 42

configs/unet/first_stage.yaml

@@ -0,0 +1,103 @@
+data:
+  syncnet_config_path: configs/syncnet/syncnet_16_pixel.yaml
+  train_output_dir: debug/unet
+  train_fileslist: /mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/all_data_v6.txt
+  train_data_dir: ""
+  audio_embeds_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/whisper_new
+  audio_mel_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/mel_new
+
+  val_video_path: assets/demo1_video.mp4
+  val_audio_path: assets/demo1_audio.wav
+  batch_size: 8 # 8
+  num_workers: 11 # 11
+  num_frames: 16
+  resolution: 256
+  mask: fix_mask
+  audio_sample_rate: 16000
+  video_fps: 25
+
+ckpt:
+  resume_ckpt_path: checkpoints/latentsync_unet.pt
+  save_ckpt_steps: 5000
+
+run:
+  pixel_space_supervise: false
+  use_syncnet: false
+  sync_loss_weight: 0.05 # 1/283
+  perceptual_loss_weight: 0.1 # 0.1
+  recon_loss_weight: 1 # 1
+  guidance_scale: 1.0 # 1.5 or 1.0
+  trepa_loss_weight: 10
+  inference_steps: 20
+  seed: 1247
+  use_mixed_noise: true
+  mixed_noise_alpha: 1 # 1
+  mixed_precision_training: true
+  enable_gradient_checkpointing: false
+  enable_xformers_memory_efficient_attention: true
+  max_train_steps: 10000000
+  max_train_epochs: -1
+
+optimizer:
+  lr: 1e-5
+  scale_lr: false
+  max_grad_norm: 1.0
+  lr_scheduler: constant
+  lr_warmup_steps: 0
+
+model:
+  act_fn: silu
+  add_audio_layer: true
+  custom_audio_layer: false
+  audio_condition_method: cross_attn # Choose between [cross_attn, group_norm]
+  attention_head_dim: 8
+  block_out_channels: [320, 640, 1280, 1280]
+  center_input_sample: false
+  cross_attention_dim: 384
+  down_block_types:
+    [
+      "CrossAttnDownBlock3D",
+      "CrossAttnDownBlock3D",
+      "CrossAttnDownBlock3D",
+      "DownBlock3D",
+    ]
+  mid_block_type: UNetMidBlock3DCrossAttn
+  up_block_types:
+    [
+      "UpBlock3D",
+      "CrossAttnUpBlock3D",
+      "CrossAttnUpBlock3D",
+      "CrossAttnUpBlock3D",
+    ]
+  downsample_padding: 1
+  flip_sin_to_cos: true
+  freq_shift: 0
+  in_channels: 13 # 49
+  layers_per_block: 2
+  mid_block_scale_factor: 1
+  norm_eps: 1e-5
+  norm_num_groups: 32
+  out_channels: 4 # 16
+  sample_size: 64
+  resnet_time_scale_shift: default # Choose between [default, scale_shift]
+  unet_use_cross_frame_attention: false
+  unet_use_temporal_attention: false
+
+  # Actually we don't use the motion module in the final version of LatentSync
+  # When we started the project, we used the codebase of AnimateDiff and tried motion module, the results are poor
+  # We decied to leave the code here for possible future usage
+  use_motion_module: false
+  motion_module_resolutions: [1, 2, 4, 8]
+  motion_module_mid_block: false
+  motion_module_decoder_only: false
+  motion_module_type: Vanilla
+  motion_module_kwargs:
+    num_attention_heads: 8
+    num_transformer_block: 1
+    attention_block_types:
+      - Temporal_Self
+      - Temporal_Self
+    temporal_position_encoding: true
+    temporal_position_encoding_max_len: 16
+    temporal_attention_dim_div: 1
+    zero_initialize: true

configs/unet/second_stage.yaml

@@ -0,0 +1,103 @@
+data:
+  syncnet_config_path: configs/syncnet/syncnet_16_pixel.yaml
+  train_output_dir: debug/unet
+  train_fileslist: /mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/all_data_v6.txt
+  train_data_dir: ""
+  audio_embeds_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/whisper_new
+  audio_mel_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/mel_new
+  
+  val_video_path: assets/demo1_video.mp4
+  val_audio_path: assets/demo1_audio.wav
+  batch_size: 2 # 8
+  num_workers: 11 # 11
+  num_frames: 16
+  resolution: 256
+  mask: fix_mask
+  audio_sample_rate: 16000
+  video_fps: 25
+
+ckpt:
+  resume_ckpt_path: checkpoints/latentsync_unet.pt
+  save_ckpt_steps: 5000
+
+run:
+  pixel_space_supervise: true
+  use_syncnet: true
+  sync_loss_weight: 0.05 # 1/283
+  perceptual_loss_weight: 0.1 # 0.1
+  recon_loss_weight: 1 # 1
+  guidance_scale: 1.0 # 1.5 or 1.0
+  trepa_loss_weight: 10
+  inference_steps: 20
+  seed: 1247
+  use_mixed_noise: true
+  mixed_noise_alpha: 1 # 1
+  mixed_precision_training: true
+  enable_gradient_checkpointing: false
+  enable_xformers_memory_efficient_attention: true
+  max_train_steps: 10000000
+  max_train_epochs: -1
+
+optimizer:
+  lr: 1e-5
+  scale_lr: false
+  max_grad_norm: 1.0
+  lr_scheduler: constant
+  lr_warmup_steps: 0
+
+model:
+  act_fn: silu
+  add_audio_layer: true
+  custom_audio_layer: false
+  audio_condition_method: cross_attn # Choose between [cross_attn, group_norm]
+  attention_head_dim: 8
+  block_out_channels: [320, 640, 1280, 1280]
+  center_input_sample: false
+  cross_attention_dim: 384
+  down_block_types:
+    [
+      "CrossAttnDownBlock3D",
+      "CrossAttnDownBlock3D",
+      "CrossAttnDownBlock3D",
+      "DownBlock3D",
+    ]
+  mid_block_type: UNetMidBlock3DCrossAttn
+  up_block_types:
+    [
+      "UpBlock3D",
+      "CrossAttnUpBlock3D",
+      "CrossAttnUpBlock3D",
+      "CrossAttnUpBlock3D",
+    ]
+  downsample_padding: 1
+  flip_sin_to_cos: true
+  freq_shift: 0
+  in_channels: 13 # 49
+  layers_per_block: 2
+  mid_block_scale_factor: 1
+  norm_eps: 1e-5
+  norm_num_groups: 32
+  out_channels: 4 # 16
+  sample_size: 64
+  resnet_time_scale_shift: default # Choose between [default, scale_shift]
+  unet_use_cross_frame_attention: false
+  unet_use_temporal_attention: false
+
+  # Actually we don't use the motion module in the final version of LatentSync
+  # When we started the project, we used the codebase of AnimateDiff and tried motion module, the results are poor
+  # We decied to leave the code here for possible future usage
+  use_motion_module: false
+  motion_module_resolutions: [1, 2, 4, 8]
+  motion_module_mid_block: false
+  motion_module_decoder_only: false
+  motion_module_type: Vanilla
+  motion_module_kwargs:
+    num_attention_heads: 8
+    num_transformer_block: 1
+    attention_block_types:
+      - Temporal_Self
+      - Temporal_Self
+    temporal_position_encoding: true
+    temporal_position_encoding_max_len: 16
+    temporal_attention_dim_div: 1
+    zero_initialize: true

data/syncnet_dataset.py

@@ -0,0 +1,153 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import numpy as np
+from torch.utils.data import Dataset
+import torch
+import random
+from ..utils.util import gather_video_paths_recursively
+from ..utils.image_processor import ImageProcessor
+from ..utils.audio import melspectrogram
+import math
+
+from decord import AudioReader, VideoReader, cpu
+
+
+class SyncNetDataset(Dataset):
+    def __init__(self, data_dir: str, fileslist: str, config):
+        if fileslist != "":
+            with open(fileslist) as file:
+                self.video_paths = [line.rstrip() for line in file]
+        elif data_dir != "":
+            self.video_paths = gather_video_paths_recursively(data_dir)
+        else:
+            raise ValueError("data_dir and fileslist cannot be both empty")
+
+        self.resolution = config.data.resolution
+        self.num_frames = config.data.num_frames
+
+        self.mel_window_length = math.ceil(self.num_frames / 5 * 16)
+
+        self.audio_sample_rate = config.data.audio_sample_rate
+        self.video_fps = config.data.video_fps
+        self.audio_samples_length = int(
+            config.data.audio_sample_rate // config.data.video_fps * config.data.num_frames
+        )
+        self.image_processor = ImageProcessor(resolution=config.data.resolution, mask="half")
+        self.audio_mel_cache_dir = config.data.audio_mel_cache_dir
+        os.makedirs(self.audio_mel_cache_dir, exist_ok=True)
+
+    def __len__(self):
+        return len(self.video_paths)
+
+    def read_audio(self, video_path: str):
+        ar = AudioReader(video_path, ctx=cpu(self.worker_id), sample_rate=self.audio_sample_rate)
+        original_mel = melspectrogram(ar[:].asnumpy().squeeze(0))
+        return torch.from_numpy(original_mel)
+
+    def crop_audio_window(self, original_mel, start_index):
+        start_idx = int(80.0 * (start_index / float(self.video_fps)))
+        end_idx = start_idx + self.mel_window_length
+        return original_mel[:, start_idx:end_idx].unsqueeze(0)
+
+    def get_frames(self, video_reader: VideoReader):
+        total_num_frames = len(video_reader)
+
+        start_idx = random.randint(0, total_num_frames - self.num_frames)
+        frames_index = np.arange(start_idx, start_idx + self.num_frames, dtype=int)
+
+        while True:
+            wrong_start_idx = random.randint(0, total_num_frames - self.num_frames)
+            # wrong_start_idx = random.randint(
+            #     max(0, start_idx - 25), min(total_num_frames - self.num_frames, start_idx + 25)
+            # )
+            if wrong_start_idx == start_idx:
+                continue
+            # if wrong_start_idx >= start_idx - self.num_frames and wrong_start_idx <= start_idx + self.num_frames:
+            #     continue
+            wrong_frames_index = np.arange(wrong_start_idx, wrong_start_idx + self.num_frames, dtype=int)
+            break
+
+        frames = video_reader.get_batch(frames_index).asnumpy()
+        wrong_frames = video_reader.get_batch(wrong_frames_index).asnumpy()
+
+        return frames, wrong_frames, start_idx
+
+    def worker_init_fn(self, worker_id):
+        # Initialize the face mesh object in each worker process,
+        # because the face mesh object cannot be called in subprocesses
+        self.worker_id = worker_id
+        # setattr(self, f"image_processor_{worker_id}", ImageProcessor(self.resolution, self.mask))
+
+    def __getitem__(self, idx):
+        # image_processor = getattr(self, f"image_processor_{self.worker_id}")
+        while True:
+            try:
+                idx = random.randint(0, len(self) - 1)
+
+                # Get video file path
+                video_path = self.video_paths[idx]
+
+                vr = VideoReader(video_path, ctx=cpu(self.worker_id))
+
+                if len(vr) < 2 * self.num_frames:
+                    continue
+
+                frames, wrong_frames, start_idx = self.get_frames(vr)
+
+                mel_cache_path = os.path.join(
+                    self.audio_mel_cache_dir, os.path.basename(video_path).replace(".mp4", "_mel.pt")
+                )
+
+                if os.path.isfile(mel_cache_path):
+                    try:
+                        original_mel = torch.load(mel_cache_path)
+                    except Exception as e:
+                        print(f"{type(e).__name__} - {e} - {mel_cache_path}")
+                        os.remove(mel_cache_path)
+                        original_mel = self.read_audio(video_path)
+                        torch.save(original_mel, mel_cache_path)
+                else:
+                    original_mel = self.read_audio(video_path)
+                    torch.save(original_mel, mel_cache_path)
+
+                mel = self.crop_audio_window(original_mel, start_idx)
+
+                if mel.shape[-1] != self.mel_window_length:
+                    continue
+
+                if random.choice([True, False]):
+                    y = torch.ones(1).float()
+                    chosen_frames = frames
+                else:
+                    y = torch.zeros(1).float()
+                    chosen_frames = wrong_frames
+
+                chosen_frames = self.image_processor.process_images(chosen_frames)
+                # chosen_frames, _, _ = image_processor.prepare_masks_and_masked_images(
+                #     chosen_frames, affine_transform=True
+                # )
+
+                vr.seek(0)  # avoid memory leak
+                break
+
+            except Exception as e:  # Handle the exception of face not detcted
+                print(f"{type(e).__name__} - {e} - {video_path}")
+                if "vr" in locals():
+                    vr.seek(0)  # avoid memory leak
+
+        sample = dict(frames=chosen_frames, audio_samples=mel, y=y)
+
+        return sample

data/unet_dataset.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import numpy as np
+from torch.utils.data import Dataset
+import torch
+import random
+import cv2
+from ..utils.image_processor import ImageProcessor, load_fixed_mask
+from ..utils.audio import melspectrogram
+from decord import AudioReader, VideoReader, cpu
+
+
+class UNetDataset(Dataset):
+    def __init__(self, train_data_dir: str, config):
+        if config.data.train_fileslist != "":
+            with open(config.data.train_fileslist) as file:
+                self.video_paths = [line.rstrip() for line in file]
+        elif train_data_dir != "":
+            self.video_paths = []
+            for file in os.listdir(train_data_dir):
+                if file.endswith(".mp4"):
+                    self.video_paths.append(os.path.join(train_data_dir, file))
+        else:
+            raise ValueError("data_dir and fileslist cannot be both empty")
+
+        self.resolution = config.data.resolution
+        self.num_frames = config.data.num_frames
+
+        if self.num_frames == 16:
+            self.mel_window_length = 52
+        elif self.num_frames == 5:
+            self.mel_window_length = 16
+        else:
+            raise NotImplementedError("Only support 16 and 5 frames now")
+
+        self.audio_sample_rate = config.data.audio_sample_rate
+        self.video_fps = config.data.video_fps
+        self.mask = config.data.mask
+        self.mask_image = load_fixed_mask(self.resolution)
+        self.load_audio_data = config.model.add_audio_layer and config.run.use_syncnet
+        self.audio_mel_cache_dir = config.data.audio_mel_cache_dir
+        os.makedirs(self.audio_mel_cache_dir, exist_ok=True)
+
+    def __len__(self):
+        return len(self.video_paths)
+
+    def read_audio(self, video_path: str):
+        ar = AudioReader(video_path, ctx=cpu(self.worker_id), sample_rate=self.audio_sample_rate)
+        original_mel = melspectrogram(ar[:].asnumpy().squeeze(0))
+        return torch.from_numpy(original_mel)
+
+    def crop_audio_window(self, original_mel, start_index):
+        start_idx = int(80.0 * (start_index / float(self.video_fps)))
+        end_idx = start_idx + self.mel_window_length
+        return original_mel[:, start_idx:end_idx].unsqueeze(0)
+
+    def get_frames(self, video_reader: VideoReader):
+        total_num_frames = len(video_reader)
+
+        start_idx = random.randint(self.num_frames // 2, total_num_frames - self.num_frames - self.num_frames // 2)
+        frames_index = np.arange(start_idx, start_idx + self.num_frames, dtype=int)
+
+        while True:
+            wrong_start_idx = random.randint(0, total_num_frames - self.num_frames)
+            if wrong_start_idx > start_idx - self.num_frames and wrong_start_idx < start_idx + self.num_frames:
+                continue
+            wrong_frames_index = np.arange(wrong_start_idx, wrong_start_idx + self.num_frames, dtype=int)
+            break
+
+        frames = video_reader.get_batch(frames_index).asnumpy()
+        wrong_frames = video_reader.get_batch(wrong_frames_index).asnumpy()
+
+        return frames, wrong_frames, start_idx
+
+    def worker_init_fn(self, worker_id):
+        # Initialize the face mesh object in each worker process,
+        # because the face mesh object cannot be called in subprocesses
+        self.worker_id = worker_id
+        setattr(
+            self,
+            f"image_processor_{worker_id}",
+            ImageProcessor(self.resolution, self.mask, mask_image=self.mask_image),
+        )
+
+    def __getitem__(self, idx):
+        image_processor = getattr(self, f"image_processor_{self.worker_id}")
+        while True:
+            try:
+                idx = random.randint(0, len(self) - 1)
+
+                # Get video file path
+                video_path = self.video_paths[idx]
+
+                vr = VideoReader(video_path, ctx=cpu(self.worker_id))
+
+                if len(vr) < 3 * self.num_frames:
+                    continue
+
+                continuous_frames, ref_frames, start_idx = self.get_frames(vr)
+
+                if self.load_audio_data:
+                    mel_cache_path = os.path.join(
+                        self.audio_mel_cache_dir, os.path.basename(video_path).replace(".mp4", "_mel.pt")
+                    )
+
+                    if os.path.isfile(mel_cache_path):
+                        try:
+                            original_mel = torch.load(mel_cache_path)
+                        except Exception as e:
+                            print(f"{type(e).__name__} - {e} - {mel_cache_path}")
+                            os.remove(mel_cache_path)
+                            original_mel = self.read_audio(video_path)
+                            torch.save(original_mel, mel_cache_path)
+                    else:
+                        original_mel = self.read_audio(video_path)
+                        torch.save(original_mel, mel_cache_path)
+
+                    mel = self.crop_audio_window(original_mel, start_idx)
+
+                    if mel.shape[-1] != self.mel_window_length:
+                        continue
+                else:
+                    mel = []
+
+                gt, masked_gt, mask = image_processor.prepare_masks_and_masked_images(
+                    continuous_frames, affine_transform=False
+                )
+
+                if self.mask == "fix_mask":
+                    ref, _, _ = image_processor.prepare_masks_and_masked_images(ref_frames, affine_transform=False)
+                else:
+                    ref = image_processor.process_images(ref_frames)
+                vr.seek(0)  # avoid memory leak
+                break
+
+            except Exception as e:  # Handle the exception of face not detcted
+                print(f"{type(e).__name__} - {e} - {video_path}")
+                if "vr" in locals():
+                    vr.seek(0)  # avoid memory leak
+
+        sample = dict(
+            gt=gt,
+            masked_gt=masked_gt,
+            ref=ref,
+            mel=mel,
+            mask=mask,
+            video_path=video_path,
+            start_idx=start_idx,
+        )
+
+        return sample

data_processing_pipeline.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+python -m preprocess.data_processing_pipeline \
+    --total_num_workers 20 \
+    --per_gpu_num_workers 10 \
+    --resolution 256 \
+    --sync_conf_threshold 3 \
+    --temp_dir temp \
+    --input_dir /mnt/bn/maliva-gen-ai-v2/chunyu.li/VoxCeleb2/raw

eval/detectors/README.md

@@ -0,0 +1,3 @@
+# Face detector
+
+This face detector is adapted from `https://github.com/cs-giung/face-detection-pytorch`.

eval/detectors/__init__.py

@@ -0,0 +1 @@
+from .s3fd import S3FD

eval/detectors/s3fd/__init__.py

@@ -0,0 +1,61 @@
+import time
+import numpy as np
+import cv2
+import torch
+from torchvision import transforms
+from .nets import S3FDNet
+from .box_utils import nms_
+
+PATH_WEIGHT = 'checkpoints/auxiliary/sfd_face.pth'
+img_mean = np.array([104., 117., 123.])[:, np.newaxis, np.newaxis].astype('float32')
+
+
+class S3FD():
+
+    def __init__(self, device='cuda'):
+
+        tstamp = time.time()
+        self.device = device
+
+        print('[S3FD] loading with', self.device)
+        self.net = S3FDNet(device=self.device).to(self.device)
+        state_dict = torch.load(PATH_WEIGHT, map_location=self.device)
+        self.net.load_state_dict(state_dict)
+        self.net.eval()
+        print('[S3FD] finished loading (%.4f sec)' % (time.time() - tstamp))
+    
+    def detect_faces(self, image, conf_th=0.8, scales=[1]):
+
+        w, h = image.shape[1], image.shape[0]
+
+        bboxes = np.empty(shape=(0, 5))
+
+        with torch.no_grad():
+            for s in scales:
+                scaled_img = cv2.resize(image, dsize=(0, 0), fx=s, fy=s, interpolation=cv2.INTER_LINEAR)
+
+                scaled_img = np.swapaxes(scaled_img, 1, 2)
+                scaled_img = np.swapaxes(scaled_img, 1, 0)
+                scaled_img = scaled_img[[2, 1, 0], :, :]
+                scaled_img = scaled_img.astype('float32')
+                scaled_img -= img_mean
+                scaled_img = scaled_img[[2, 1, 0], :, :]
+                x = torch.from_numpy(scaled_img).unsqueeze(0).to(self.device)
+                y = self.net(x)
+
+                detections = y.data
+                scale = torch.Tensor([w, h, w, h])
+
+                for i in range(detections.size(1)):
+                    j = 0
+                    while detections[0, i, j, 0] > conf_th:
+                        score = detections[0, i, j, 0]
+                        pt = (detections[0, i, j, 1:] * scale).cpu().numpy()
+                        bbox = (pt[0], pt[1], pt[2], pt[3], score)
+                        bboxes = np.vstack((bboxes, bbox))
+                        j += 1
+
+            keep = nms_(bboxes, 0.1)
+            bboxes = bboxes[keep]
+
+        return bboxes

eval/detectors/s3fd/box_utils.py

@@ -0,0 +1,221 @@
+import numpy as np
+from itertools import product as product
+import torch
+from torch.autograd import Function
+import warnings
+
+
+def nms_(dets, thresh):
+    """
+    Courtesy of Ross Girshick
+    [https://github.com/rbgirshick/py-faster-rcnn/blob/master/lib/nms/py_cpu_nms.py]
+    """
+    x1 = dets[:, 0]
+    y1 = dets[:, 1]
+    x2 = dets[:, 2]
+    y2 = dets[:, 3]
+    scores = dets[:, 4]
+
+    areas = (x2 - x1) * (y2 - y1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(int(i))
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1)
+        h = np.maximum(0.0, yy2 - yy1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+        inds = np.where(ovr <= thresh)[0]
+        order = order[inds + 1]
+
+    return np.array(keep).astype(np.int32)
+
+
+def decode(loc, priors, variances):
+    """Decode locations from predictions using priors to undo
+    the encoding we did for offset regression at train time.
+    Args:
+        loc (tensor): location predictions for loc layers,
+            Shape: [num_priors,4]
+        priors (tensor): Prior boxes in center-offset form.
+            Shape: [num_priors,4].
+        variances: (list[float]) Variances of priorboxes
+    Return:
+        decoded bounding box predictions
+    """
+
+    boxes = torch.cat((
+        priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:],
+        priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1)
+    boxes[:, :2] -= boxes[:, 2:] / 2
+    boxes[:, 2:] += boxes[:, :2]
+    return boxes
+
+
+def nms(boxes, scores, overlap=0.5, top_k=200):
+    """Apply non-maximum suppression at test time to avoid detecting too many
+    overlapping bounding boxes for a given object.
+    Args:
+        boxes: (tensor) The location preds for the img, Shape: [num_priors,4].
+        scores: (tensor) The class predscores for the img, Shape:[num_priors].
+        overlap: (float) The overlap thresh for suppressing unnecessary boxes.
+        top_k: (int) The Maximum number of box preds to consider.
+    Return:
+        The indices of the kept boxes with respect to num_priors.
+    """
+
+    keep = scores.new(scores.size(0)).zero_().long()
+    if boxes.numel() == 0:
+        return keep, 0
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+    area = torch.mul(x2 - x1, y2 - y1)
+    v, idx = scores.sort(0)  # sort in ascending order
+    # I = I[v >= 0.01]
+    idx = idx[-top_k:]  # indices of the top-k largest vals
+    xx1 = boxes.new()
+    yy1 = boxes.new()
+    xx2 = boxes.new()
+    yy2 = boxes.new()
+    w = boxes.new()
+    h = boxes.new()
+
+    # keep = torch.Tensor()
+    count = 0
+    while idx.numel() > 0:
+        i = idx[-1]  # index of current largest val
+        # keep.append(i)
+        keep[count] = i
+        count += 1
+        if idx.size(0) == 1:
+            break
+        idx = idx[:-1]  # remove kept element from view
+        # load bboxes of next highest vals
+        with warnings.catch_warnings():
+            # Ignore UserWarning within this block
+            warnings.simplefilter("ignore", category=UserWarning)
+            torch.index_select(x1, 0, idx, out=xx1)
+            torch.index_select(y1, 0, idx, out=yy1)
+            torch.index_select(x2, 0, idx, out=xx2)
+            torch.index_select(y2, 0, idx, out=yy2)
+        # store element-wise max with next highest score
+        xx1 = torch.clamp(xx1, min=x1[i])
+        yy1 = torch.clamp(yy1, min=y1[i])
+        xx2 = torch.clamp(xx2, max=x2[i])
+        yy2 = torch.clamp(yy2, max=y2[i])
+        w.resize_as_(xx2)
+        h.resize_as_(yy2)
+        w = xx2 - xx1
+        h = yy2 - yy1
+        # check sizes of xx1 and xx2.. after each iteration
+        w = torch.clamp(w, min=0.0)
+        h = torch.clamp(h, min=0.0)
+        inter = w * h
+        # IoU = i / (area(a) + area(b) - i)
+        rem_areas = torch.index_select(area, 0, idx)  # load remaining areas)
+        union = (rem_areas - inter) + area[i]
+        IoU = inter / union  # store result in iou
+        # keep only elements with an IoU <= overlap
+        idx = idx[IoU.le(overlap)]
+    return keep, count
+
+
+class Detect(object):
+
+    def __init__(self, num_classes=2,
+                    top_k=750, nms_thresh=0.3, conf_thresh=0.05,
+                    variance=[0.1, 0.2], nms_top_k=5000):
+        
+        self.num_classes = num_classes
+        self.top_k = top_k
+        self.nms_thresh = nms_thresh
+        self.conf_thresh = conf_thresh
+        self.variance = variance
+        self.nms_top_k = nms_top_k
+
+    def forward(self, loc_data, conf_data, prior_data):
+
+        num = loc_data.size(0)
+        num_priors = prior_data.size(0)
+
+        conf_preds = conf_data.view(num, num_priors, self.num_classes).transpose(2, 1)
+        batch_priors = prior_data.view(-1, num_priors, 4).expand(num, num_priors, 4)
+        batch_priors = batch_priors.contiguous().view(-1, 4)
+
+        decoded_boxes = decode(loc_data.view(-1, 4), batch_priors, self.variance)
+        decoded_boxes = decoded_boxes.view(num, num_priors, 4)
+
+        output = torch.zeros(num, self.num_classes, self.top_k, 5)
+
+        for i in range(num):
+            boxes = decoded_boxes[i].clone()
+            conf_scores = conf_preds[i].clone()
+
+            for cl in range(1, self.num_classes):
+                c_mask = conf_scores[cl].gt(self.conf_thresh)
+                scores = conf_scores[cl][c_mask]
+                
+                if scores.dim() == 0:
+                    continue
+                l_mask = c_mask.unsqueeze(1).expand_as(boxes)
+                boxes_ = boxes[l_mask].view(-1, 4)
+                ids, count = nms(boxes_, scores, self.nms_thresh, self.nms_top_k)
+                count = count if count < self.top_k else self.top_k
+
+                output[i, cl, :count] = torch.cat((scores[ids[:count]].unsqueeze(1), boxes_[ids[:count]]), 1)
+
+        return output
+
+
+class PriorBox(object):
+
+    def __init__(self, input_size, feature_maps,
+                    variance=[0.1, 0.2],
+                    min_sizes=[16, 32, 64, 128, 256, 512],
+                    steps=[4, 8, 16, 32, 64, 128],
+                    clip=False):
+
+        super(PriorBox, self).__init__()
+
+        self.imh = input_size[0]
+        self.imw = input_size[1]
+        self.feature_maps = feature_maps
+
+        self.variance = variance
+        self.min_sizes = min_sizes
+        self.steps = steps
+        self.clip = clip
+
+    def forward(self):
+        mean = []
+        for k, fmap in enumerate(self.feature_maps):
+            feath = fmap[0]
+            featw = fmap[1]
+            for i, j in product(range(feath), range(featw)):
+                f_kw = self.imw / self.steps[k]
+                f_kh = self.imh / self.steps[k]
+
+                cx = (j + 0.5) / f_kw
+                cy = (i + 0.5) / f_kh
+
+                s_kw = self.min_sizes[k] / self.imw
+                s_kh = self.min_sizes[k] / self.imh
+
+                mean += [cx, cy, s_kw, s_kh]
+
+        output = torch.FloatTensor(mean).view(-1, 4)
+        
+        if self.clip:
+            output.clamp_(max=1, min=0)
+        
+        return output

eval/detectors/s3fd/nets.py

@@ -0,0 +1,174 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.nn.init as init
+from .box_utils import Detect, PriorBox
+
+
+class L2Norm(nn.Module):
+
+    def __init__(self, n_channels, scale):
+        super(L2Norm, self).__init__()
+        self.n_channels = n_channels
+        self.gamma = scale or None
+        self.eps = 1e-10
+        self.weight = nn.Parameter(torch.Tensor(self.n_channels))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        init.constant_(self.weight, self.gamma)
+
+    def forward(self, x):
+        norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps
+        x = torch.div(x, norm)
+        out = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand_as(x) * x
+        return out
+
+
+class S3FDNet(nn.Module):
+
+    def __init__(self, device='cuda'):
+        super(S3FDNet, self).__init__()
+        self.device = device
+
+        self.vgg = nn.ModuleList([
+            nn.Conv2d(3, 64, 3, 1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(64, 64, 3, 1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+
+            nn.Conv2d(64, 128, 3, 1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, 3, 1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+            
+            nn.Conv2d(128, 256, 3, 1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, 3, 1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, 3, 1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2, ceil_mode=True),
+            
+            nn.Conv2d(256, 512, 3, 1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(512, 512, 3, 1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(512, 512, 3, 1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+
+            nn.Conv2d(512, 512, 3, 1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(512, 512, 3, 1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(512, 512, 3, 1, padding=1),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(2, 2),
+
+            nn.Conv2d(512, 1024, 3, 1, padding=6, dilation=6),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(1024, 1024, 1, 1),
+            nn.ReLU(inplace=True),
+        ])
+
+        self.L2Norm3_3 = L2Norm(256, 10)
+        self.L2Norm4_3 = L2Norm(512, 8)
+        self.L2Norm5_3 = L2Norm(512, 5)
+
+        self.extras = nn.ModuleList([
+            nn.Conv2d(1024, 256, 1, 1),
+            nn.Conv2d(256, 512, 3, 2, padding=1),
+            nn.Conv2d(512, 128, 1, 1),
+            nn.Conv2d(128, 256, 3, 2, padding=1),
+        ])
+        
+        self.loc = nn.ModuleList([
+            nn.Conv2d(256, 4, 3, 1, padding=1),
+            nn.Conv2d(512, 4, 3, 1, padding=1),
+            nn.Conv2d(512, 4, 3, 1, padding=1),
+            nn.Conv2d(1024, 4, 3, 1, padding=1),
+            nn.Conv2d(512, 4, 3, 1, padding=1),
+            nn.Conv2d(256, 4, 3, 1, padding=1),
+        ])
+
+        self.conf = nn.ModuleList([
+            nn.Conv2d(256, 4, 3, 1, padding=1),
+            nn.Conv2d(512, 2, 3, 1, padding=1),
+            nn.Conv2d(512, 2, 3, 1, padding=1),
+            nn.Conv2d(1024, 2, 3, 1, padding=1),
+            nn.Conv2d(512, 2, 3, 1, padding=1),
+            nn.Conv2d(256, 2, 3, 1, padding=1),
+        ])
+
+        self.softmax = nn.Softmax(dim=-1)
+        self.detect = Detect()
+
+    def forward(self, x):
+        size = x.size()[2:]
+        sources = list()
+        loc = list()
+        conf = list()
+
+        for k in range(16):
+            x = self.vgg[k](x)
+        s = self.L2Norm3_3(x)
+        sources.append(s)
+
+        for k in range(16, 23):
+            x = self.vgg[k](x)
+        s = self.L2Norm4_3(x)
+        sources.append(s)
+
+        for k in range(23, 30):
+            x = self.vgg[k](x)
+        s = self.L2Norm5_3(x)
+        sources.append(s)
+
+        for k in range(30, len(self.vgg)):
+            x = self.vgg[k](x)
+        sources.append(x)
+        
+        # apply extra layers and cache source layer outputs
+        for k, v in enumerate(self.extras):
+            x = F.relu(v(x), inplace=True)
+            if k % 2 == 1:
+                sources.append(x)
+
+        # apply multibox head to source layers
+        loc_x = self.loc[0](sources[0])
+        conf_x = self.conf[0](sources[0])
+
+        max_conf, _ = torch.max(conf_x[:, 0:3, :, :], dim=1, keepdim=True)
+        conf_x = torch.cat((max_conf, conf_x[:, 3:, :, :]), dim=1)
+
+        loc.append(loc_x.permute(0, 2, 3, 1).contiguous())
+        conf.append(conf_x.permute(0, 2, 3, 1).contiguous())
+
+        for i in range(1, len(sources)):
+            x = sources[i]
+            conf.append(self.conf[i](x).permute(0, 2, 3, 1).contiguous())
+            loc.append(self.loc[i](x).permute(0, 2, 3, 1).contiguous())
+
+        features_maps = []
+        for i in range(len(loc)):
+            feat = []
+            feat += [loc[i].size(1), loc[i].size(2)]
+            features_maps += [feat]
+
+        loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1)
+        conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)
+
+        with torch.no_grad():
+            self.priorbox = PriorBox(size, features_maps)
+            self.priors = self.priorbox.forward()
+
+        output = self.detect.forward(
+            loc.view(loc.size(0), -1, 4),
+            self.softmax(conf.view(conf.size(0), -1, 2)),
+            self.priors.type(type(x.data)).to(self.device)
+        )
+
+        return output

eval/draw_syncnet_lines.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import matplotlib.pyplot as plt
+
+
+class Chart:
+    def __init__(self):
+        self.loss_list = []
+
+    def add_ckpt(self, ckpt_path, line_name):
+        ckpt = torch.load(ckpt_path, map_location="cpu")
+        train_step_list = ckpt["train_step_list"]
+        train_loss_list = ckpt["train_loss_list"]
+        val_step_list = ckpt["val_step_list"]
+        val_loss_list = ckpt["val_loss_list"]
+        val_step_list = [val_step_list[0]] + val_step_list[4::5]
+        val_loss_list = [val_loss_list[0]] + val_loss_list[4::5]
+        self.loss_list.append((line_name, train_step_list, train_loss_list, val_step_list, val_loss_list))
+
+    def draw(self, save_path, plot_val=True):
+        # Global settings
+        plt.rcParams["font.size"] = 14
+        plt.rcParams["font.family"] = "serif"
+        plt.rcParams["font.sans-serif"] = ["Arial", "DejaVu Sans", "Lucida Grande"]
+        plt.rcParams["font.serif"] = ["Times New Roman", "DejaVu Serif"]
+
+        # Creating the plot
+        plt.figure(figsize=(7.766, 4.8)) # Golden ratio
+        for loss in self.loss_list:
+            if plot_val:
+                (line,) = plt.plot(loss[1], loss[2], label=loss[0], linewidth=0.5, alpha=0.5)
+                line_color = line.get_color()
+                plt.plot(loss[3], loss[4], linewidth=1.5, color=line_color)
+            else:
+                plt.plot(loss[1], loss[2], label=loss[0], linewidth=1)
+        plt.xlabel("Step")
+        plt.ylabel("Loss")
+        legend = plt.legend()
+        # legend = plt.legend(loc='upper right', bbox_to_anchor=(1, 0.82))
+
+        # Adjust the linewidth of legend
+        for line in legend.get_lines():
+            line.set_linewidth(2)
+
+        plt.savefig(save_path, transparent=True)
+        plt.close()
+
+
+if __name__ == "__main__":
+    chart = Chart()
+    # chart.add_ckpt("output/syncnet/train-2024_10_25-18:14:43/checkpoints/checkpoint-10000.pt", "w/ self-attn")
+    # chart.add_ckpt("output/syncnet/train-2024_10_25-18:21:59/checkpoints/checkpoint-10000.pt", "w/o self-attn")
+    chart.add_ckpt("output/syncnet/train-2024_10_24-21:03:11/checkpoints/checkpoint-10000.pt", "Dim 512")
+    chart.add_ckpt("output/syncnet/train-2024_10_25-18:21:59/checkpoints/checkpoint-10000.pt", "Dim 2048")
+    chart.add_ckpt("output/syncnet/train-2024_10_24-22:37:04/checkpoints/checkpoint-10000.pt", "Dim 4096")
+    chart.add_ckpt("output/syncnet/train-2024_10_25-02:30:17/checkpoints/checkpoint-10000.pt", "Dim 6144")
+    chart.draw("ablation.pdf", plot_val=True)

eval/eval_fvd.py

@@ -0,0 +1,96 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import mediapipe as mp
+import cv2
+from decord import VideoReader
+from einops import rearrange
+import os
+import numpy as np
+import torch
+import tqdm
+from eval.fvd import compute_our_fvd
+
+
+class FVD:
+    def __init__(self, resolution=(224, 224)):
+        self.face_detector = mp.solutions.face_detection.FaceDetection(model_selection=0, min_detection_confidence=0.5)
+        self.resolution = resolution
+
+    def detect_face(self, image):
+        height, width = image.shape[:2]
+        # Process the image and detect faces.
+        results = self.face_detector.process(image)
+
+        if not results.detections:  # Face not detected
+            raise Exception("Face not detected")
+
+        detection = results.detections[0]  # Only use the first face in the image
+        bounding_box = detection.location_data.relative_bounding_box
+        xmin = int(bounding_box.xmin * width)
+        ymin = int(bounding_box.ymin * height)
+        face_width = int(bounding_box.width * width)
+        face_height = int(bounding_box.height * height)
+
+        # Crop the image to the bounding box.
+        xmin = max(0, xmin)
+        ymin = max(0, ymin)
+        xmax = min(width, xmin + face_width)
+        ymax = min(height, ymin + face_height)
+        image = image[ymin:ymax, xmin:xmax]
+
+        return image
+
+    def detect_video(self, video_path, real: bool = True):
+        vr = VideoReader(video_path)
+        video_frames = vr[20:36].asnumpy()  # Use one frame per second
+        vr.seek(0)  # avoid memory leak
+        faces = []
+        for frame in video_frames:
+            face = self.detect_face(frame)
+            face = cv2.resize(face, (self.resolution[1], self.resolution[0]), interpolation=cv2.INTER_AREA)
+            faces.append(face)
+
+        if len(faces) != 16:
+            return None
+        faces = np.stack(faces, axis=0)  # (f, h, w, c)
+        faces = torch.from_numpy(faces)
+        return faces
+
+
+def eval_fvd(real_videos_dir, fake_videos_dir):
+    fvd = FVD()
+    real_features_list = []
+    fake_features_list = []
+    for file in tqdm.tqdm(os.listdir(fake_videos_dir)):
+        if file.endswith(".mp4"):
+            real_video_path = os.path.join(real_videos_dir, file.replace("_out.mp4", ".mp4"))
+            fake_video_path = os.path.join(fake_videos_dir, file)
+            real_features = fvd.detect_video(real_video_path, real=True)
+            fake_features = fvd.detect_video(fake_video_path, real=False)
+            if real_features is None or fake_features is None:
+                continue
+            real_features_list.append(real_features)
+            fake_features_list.append(fake_features)
+
+    real_features = torch.stack(real_features_list) / 255.0
+    fake_features = torch.stack(fake_features_list) / 255.0
+    print(compute_our_fvd(real_features, fake_features, device="cpu"))
+
+
+if __name__ == "__main__":
+    real_videos_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/VoxCeleb2/segmented/cross"
+    fake_videos_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/VoxCeleb2/segmented/latentsync_cross"
+
+    eval_fvd(real_videos_dir, fake_videos_dir)

eval/eval_sync_conf.py

@@ -0,0 +1,77 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+import tqdm
+from statistics import fmean
+from eval.syncnet import SyncNetEval
+from eval.syncnet_detect import SyncNetDetector
+from latentsync.utils.util import red_text
+import torch
+
+
+def syncnet_eval(syncnet, syncnet_detector, video_path, temp_dir, detect_results_dir="detect_results"):
+    syncnet_detector(video_path=video_path, min_track=50)
+    crop_videos = os.listdir(os.path.join(detect_results_dir, "crop"))
+    if crop_videos == []:
+        raise Exception(red_text(f"Face not detected in {video_path}"))
+    av_offset_list = []
+    conf_list = []
+    for video in crop_videos:
+        av_offset, _, conf = syncnet.evaluate(
+            video_path=os.path.join(detect_results_dir, "crop", video), temp_dir=temp_dir
+        )
+        av_offset_list.append(av_offset)
+        conf_list.append(conf)
+    av_offset = int(fmean(av_offset_list))
+    conf = fmean(conf_list)
+    print(f"Input video: {video_path}\nSyncNet confidence: {conf:.2f}\nAV offset: {av_offset}")
+    return av_offset, conf
+
+
+def main():
+    parser = argparse.ArgumentParser(description="SyncNet")
+    parser.add_argument("--initial_model", type=str, default="checkpoints/auxiliary/syncnet_v2.model", help="")
+    parser.add_argument("--video_path", type=str, default=None, help="")
+    parser.add_argument("--videos_dir", type=str, default="/root/processed")
+    parser.add_argument("--temp_dir", type=str, default="temp", help="")
+
+    args = parser.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    syncnet = SyncNetEval(device=device)
+    syncnet.loadParameters(args.initial_model)
+
+    syncnet_detector = SyncNetDetector(device=device, detect_results_dir="detect_results")
+
+    if args.video_path is not None:
+        syncnet_eval(syncnet, syncnet_detector, args.video_path, args.temp_dir)
+    else:
+        sync_conf_list = []
+        video_names = sorted([f for f in os.listdir(args.videos_dir) if f.endswith(".mp4")])
+        for video_name in tqdm.tqdm(video_names):
+            try:
+                _, conf = syncnet_eval(
+                    syncnet, syncnet_detector, os.path.join(args.videos_dir, video_name), args.temp_dir
+                )
+                sync_conf_list.append(conf)
+            except Exception as e:
+                print(e)
+        print(f"The average sync confidence is {fmean(sync_conf_list):.02f}")
+
+
+if __name__ == "__main__":
+    main()

eval/eval_sync_conf.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+python -m eval.eval_sync_conf --video_path "RD_Radio1_000_006_out.mp4"

eval/eval_syncnet_acc.py

@@ -0,0 +1,118 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+from tqdm.auto import tqdm
+import torch
+import torch.nn as nn
+from einops import rearrange
+from latentsync.models.syncnet import SyncNet
+from latentsync.data.syncnet_dataset import SyncNetDataset
+from diffusers import AutoencoderKL
+from omegaconf import OmegaConf
+from accelerate.utils import set_seed
+
+
+def main(config):
+    set_seed(config.run.seed)
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if config.data.latent_space:
+        vae = AutoencoderKL.from_pretrained(
+            "runwayml/stable-diffusion-inpainting", subfolder="vae", revision="fp16", torch_dtype=torch.float16
+        )
+        vae.requires_grad_(False)
+        vae.to(device)
+
+    # Dataset and Dataloader setup
+    dataset = SyncNetDataset(config.data.val_data_dir, config.data.val_fileslist, config)
+
+    test_dataloader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=config.data.batch_size,
+        shuffle=False,
+        num_workers=config.data.num_workers,
+        drop_last=False,
+        worker_init_fn=dataset.worker_init_fn,
+    )
+
+    # Model
+    syncnet = SyncNet(OmegaConf.to_container(config.model)).to(device)
+
+    print(f"Load checkpoint from: {config.ckpt.inference_ckpt_path}")
+    checkpoint = torch.load(config.ckpt.inference_ckpt_path, map_location=device)
+
+    syncnet.load_state_dict(checkpoint["state_dict"])
+    syncnet.to(dtype=torch.float16)
+    syncnet.requires_grad_(False)
+    syncnet.eval()
+
+    global_step = 0
+    num_val_batches = config.data.num_val_samples // config.data.batch_size
+    progress_bar = tqdm(range(0, num_val_batches), initial=0, desc="Testing accuracy")
+
+    num_correct_preds = 0
+    num_total_preds = 0
+
+    while True:
+        for step, batch in enumerate(test_dataloader):
+            ### >>>> Test >>>> ###
+
+            frames = batch["frames"].to(device, dtype=torch.float16)
+            audio_samples = batch["audio_samples"].to(device, dtype=torch.float16)
+            y = batch["y"].to(device, dtype=torch.float16).squeeze(1)
+
+            if config.data.latent_space:
+                frames = rearrange(frames, "b f c h w -> (b f) c h w")
+
+                with torch.no_grad():
+                    frames = vae.encode(frames).latent_dist.sample() * 0.18215
+
+                frames = rearrange(frames, "(b f) c h w -> b (f c) h w", f=config.data.num_frames)
+            else:
+                frames = rearrange(frames, "b f c h w -> b (f c) h w")
+
+            if config.data.lower_half:
+                height = frames.shape[2]
+                frames = frames[:, :, height // 2 :, :]
+
+            with torch.no_grad():
+                vision_embeds, audio_embeds = syncnet(frames, audio_samples)
+
+            sims = nn.functional.cosine_similarity(vision_embeds, audio_embeds)
+
+            preds = (sims > 0.5).to(dtype=torch.float16)
+            num_correct_preds += (preds == y).sum().item()
+            num_total_preds += len(sims)
+
+            progress_bar.update(1)
+            global_step += 1
+
+            if global_step >= num_val_batches:
+                progress_bar.close()
+                print(f"Accuracy score: {num_correct_preds / num_total_preds*100:.2f}%")
+                return
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Code to test the accuracy of expert lip-sync discriminator")
+
+    parser.add_argument("--config_path", type=str, default="configs/syncnet/syncnet_16_latent.yaml")
+    args = parser.parse_args()
+
+    # Load a configuration file
+    config = OmegaConf.load(args.config_path)
+
+    main(config)

eval/eval_syncnet_acc.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python -m eval.eval_syncnet_acc --config_path "configs/syncnet/syncnet_16_pixel.yaml"

eval/fvd.py

@@ -0,0 +1,56 @@
+# Adapted from https://github.com/universome/fvd-comparison/blob/master/our_fvd.py
+
+from typing import Tuple
+import scipy
+import numpy as np
+import torch
+
+
+def compute_fvd(feats_fake: np.ndarray, feats_real: np.ndarray) -> float:
+    mu_gen, sigma_gen = compute_stats(feats_fake)
+    mu_real, sigma_real = compute_stats(feats_real)
+
+    m = np.square(mu_gen - mu_real).sum()
+    s, _ = scipy.linalg.sqrtm(np.dot(sigma_gen, sigma_real), disp=False)  # pylint: disable=no-member
+    fid = np.real(m + np.trace(sigma_gen + sigma_real - s * 2))
+
+    return float(fid)
+
+
+def compute_stats(feats: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    mu = feats.mean(axis=0)  # [d]
+    sigma = np.cov(feats, rowvar=False)  # [d, d]
+
+    return mu, sigma
+
+
+@torch.no_grad()
+def compute_our_fvd(videos_fake: np.ndarray, videos_real: np.ndarray, device: str = "cuda") -> float:
+    i3d_path = "checkpoints/auxiliary/i3d_torchscript.pt"
+    i3d_kwargs = dict(
+        rescale=False, resize=False, return_features=True
+    )  # Return raw features before the softmax layer.
+
+    with open(i3d_path, "rb") as f:
+        i3d_model = torch.jit.load(f).eval().to(device)
+
+    videos_fake = videos_fake.permute(0, 4, 1, 2, 3).to(device)
+    videos_real = videos_real.permute(0, 4, 1, 2, 3).to(device)
+
+    feats_fake = i3d_model(videos_fake, **i3d_kwargs).cpu().numpy()
+    feats_real = i3d_model(videos_real, **i3d_kwargs).cpu().numpy()
+
+    return compute_fvd(feats_fake, feats_real)
+
+
+def main():
+    # input shape: (b, f, h, w, c)
+    videos_fake = torch.rand(10, 16, 224, 224, 3)
+    videos_real = torch.rand(10, 16, 224, 224, 3)
+
+    our_fvd_result = compute_our_fvd(videos_fake, videos_real)
+    print(f"[FVD scores] Ours: {our_fvd_result}")
+
+
+if __name__ == "__main__":
+    main()

eval/hyper_iqa.py

@@ -0,0 +1,343 @@
+# Adapted from https://github.com/SSL92/hyperIQA/blob/master/models.py
+
+import torch as torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.nn import init
+import math
+import torch.utils.model_zoo as model_zoo
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
+}
+
+
+class HyperNet(nn.Module):
+    """
+    Hyper network for learning perceptual rules.
+
+    Args:
+        lda_out_channels: local distortion aware module output size.
+        hyper_in_channels: input feature channels for hyper network.
+        target_in_size: input vector size for target network.
+        target_fc(i)_size: fully connection layer size of target network.
+        feature_size: input feature map width/height for hyper network.
+
+    Note:
+        For size match, input args must satisfy: 'target_fc(i)_size * target_fc(i+1)_size' is divisible by 'feature_size ^ 2'.
+
+    """
+    def __init__(self, lda_out_channels, hyper_in_channels, target_in_size, target_fc1_size, target_fc2_size, target_fc3_size, target_fc4_size, feature_size):
+        super(HyperNet, self).__init__()
+
+        self.hyperInChn = hyper_in_channels
+        self.target_in_size = target_in_size
+        self.f1 = target_fc1_size
+        self.f2 = target_fc2_size
+        self.f3 = target_fc3_size
+        self.f4 = target_fc4_size
+        self.feature_size = feature_size
+
+        self.res = resnet50_backbone(lda_out_channels, target_in_size, pretrained=True)
+
+        self.pool = nn.AdaptiveAvgPool2d((1, 1))
+
+        # Conv layers for resnet output features
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(2048, 1024, 1, padding=(0, 0)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(1024, 512, 1, padding=(0, 0)),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(512, self.hyperInChn, 1, padding=(0, 0)),
+            nn.ReLU(inplace=True)
+        )
+
+        # Hyper network part, conv for generating target fc weights, fc for generating target fc biases
+        self.fc1w_conv = nn.Conv2d(self.hyperInChn, int(self.target_in_size * self.f1 / feature_size ** 2), 3,  padding=(1, 1))
+        self.fc1b_fc = nn.Linear(self.hyperInChn, self.f1)
+
+        self.fc2w_conv = nn.Conv2d(self.hyperInChn, int(self.f1 * self.f2 / feature_size ** 2), 3, padding=(1, 1))
+        self.fc2b_fc = nn.Linear(self.hyperInChn, self.f2)
+
+        self.fc3w_conv = nn.Conv2d(self.hyperInChn, int(self.f2 * self.f3 / feature_size ** 2), 3, padding=(1, 1))
+        self.fc3b_fc = nn.Linear(self.hyperInChn, self.f3)
+
+        self.fc4w_conv = nn.Conv2d(self.hyperInChn, int(self.f3 * self.f4 / feature_size ** 2), 3, padding=(1, 1))
+        self.fc4b_fc = nn.Linear(self.hyperInChn, self.f4)
+
+        self.fc5w_fc = nn.Linear(self.hyperInChn, self.f4)
+        self.fc5b_fc = nn.Linear(self.hyperInChn, 1)
+
+        # initialize
+        for i, m_name in enumerate(self._modules):
+            if i > 2:
+                nn.init.kaiming_normal_(self._modules[m_name].weight.data)
+
+    def forward(self, img):
+        feature_size = self.feature_size
+
+        res_out = self.res(img)
+
+        # input vector for target net
+        target_in_vec = res_out['target_in_vec'].reshape(-1, self.target_in_size, 1, 1)
+
+        # input features for hyper net
+        hyper_in_feat = self.conv1(res_out['hyper_in_feat']).reshape(-1, self.hyperInChn, feature_size, feature_size)
+
+        # generating target net weights & biases
+        target_fc1w = self.fc1w_conv(hyper_in_feat).reshape(-1, self.f1, self.target_in_size, 1, 1)
+        target_fc1b = self.fc1b_fc(self.pool(hyper_in_feat).squeeze()).reshape(-1, self.f1)
+
+        target_fc2w = self.fc2w_conv(hyper_in_feat).reshape(-1, self.f2, self.f1, 1, 1)
+        target_fc2b = self.fc2b_fc(self.pool(hyper_in_feat).squeeze()).reshape(-1, self.f2)
+
+        target_fc3w = self.fc3w_conv(hyper_in_feat).reshape(-1, self.f3, self.f2, 1, 1)
+        target_fc3b = self.fc3b_fc(self.pool(hyper_in_feat).squeeze()).reshape(-1, self.f3)
+
+        target_fc4w = self.fc4w_conv(hyper_in_feat).reshape(-1, self.f4, self.f3, 1, 1)
+        target_fc4b = self.fc4b_fc(self.pool(hyper_in_feat).squeeze()).reshape(-1, self.f4)
+
+        target_fc5w = self.fc5w_fc(self.pool(hyper_in_feat).squeeze()).reshape(-1, 1, self.f4, 1, 1)
+        target_fc5b = self.fc5b_fc(self.pool(hyper_in_feat).squeeze()).reshape(-1, 1)
+
+        out = {}
+        out['target_in_vec'] = target_in_vec
+        out['target_fc1w'] = target_fc1w
+        out['target_fc1b'] = target_fc1b
+        out['target_fc2w'] = target_fc2w
+        out['target_fc2b'] = target_fc2b
+        out['target_fc3w'] = target_fc3w
+        out['target_fc3b'] = target_fc3b
+        out['target_fc4w'] = target_fc4w
+        out['target_fc4b'] = target_fc4b
+        out['target_fc5w'] = target_fc5w
+        out['target_fc5b'] = target_fc5b
+
+        return out
+
+
+class TargetNet(nn.Module):
+    """
+    Target network for quality prediction.
+    """
+    def __init__(self, paras):
+        super(TargetNet, self).__init__()
+        self.l1 = nn.Sequential(
+            TargetFC(paras['target_fc1w'], paras['target_fc1b']),
+            nn.Sigmoid(),
+        )
+        self.l2 = nn.Sequential(
+            TargetFC(paras['target_fc2w'], paras['target_fc2b']),
+            nn.Sigmoid(),
+        )
+
+        self.l3 = nn.Sequential(
+            TargetFC(paras['target_fc3w'], paras['target_fc3b']),
+            nn.Sigmoid(),
+        )
+
+        self.l4 = nn.Sequential(
+            TargetFC(paras['target_fc4w'], paras['target_fc4b']),
+            nn.Sigmoid(),
+            TargetFC(paras['target_fc5w'], paras['target_fc5b']),
+        )
+
+    def forward(self, x):
+        q = self.l1(x)
+        # q = F.dropout(q)
+        q = self.l2(q)
+        q = self.l3(q)
+        q = self.l4(q).squeeze()
+        return q
+
+
+class TargetFC(nn.Module):
+    """
+    Fully connection operations for target net
+
+    Note:
+        Weights & biases are different for different images in a batch,
+        thus here we use group convolution for calculating images in a batch with individual weights & biases.
+    """
+    def __init__(self, weight, bias):
+        super(TargetFC, self).__init__()
+        self.weight = weight
+        self.bias = bias
+
+    def forward(self, input_):
+
+        input_re = input_.reshape(-1, input_.shape[0] * input_.shape[1], input_.shape[2], input_.shape[3])
+        weight_re = self.weight.reshape(self.weight.shape[0] * self.weight.shape[1], self.weight.shape[2], self.weight.shape[3], self.weight.shape[4])
+        bias_re = self.bias.reshape(self.bias.shape[0] * self.bias.shape[1])
+        out = F.conv2d(input=input_re, weight=weight_re, bias=bias_re, groups=self.weight.shape[0])
+
+        return out.reshape(input_.shape[0], self.weight.shape[1], input_.shape[2], input_.shape[3])
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class ResNetBackbone(nn.Module):
+
+    def __init__(self, lda_out_channels, in_chn, block, layers, num_classes=1000):
+        super(ResNetBackbone, self).__init__()
+        self.inplanes = 64
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+
+        # local distortion aware module
+        self.lda1_pool = nn.Sequential(
+            nn.Conv2d(256, 16, kernel_size=1, stride=1, padding=0, bias=False),
+            nn.AvgPool2d(7, stride=7),
+        )
+        self.lda1_fc = nn.Linear(16 * 64, lda_out_channels)
+
+        self.lda2_pool = nn.Sequential(
+            nn.Conv2d(512, 32, kernel_size=1, stride=1, padding=0, bias=False),
+            nn.AvgPool2d(7, stride=7),
+        )
+        self.lda2_fc = nn.Linear(32 * 16, lda_out_channels)
+
+        self.lda3_pool = nn.Sequential(
+            nn.Conv2d(1024, 64, kernel_size=1, stride=1, padding=0, bias=False),
+            nn.AvgPool2d(7, stride=7),
+        )
+        self.lda3_fc = nn.Linear(64 * 4, lda_out_channels)
+
+        self.lda4_pool = nn.AvgPool2d(7, stride=7)
+        self.lda4_fc = nn.Linear(2048, in_chn - lda_out_channels * 3)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+        # initialize
+        nn.init.kaiming_normal_(self.lda1_pool._modules['0'].weight.data)
+        nn.init.kaiming_normal_(self.lda2_pool._modules['0'].weight.data)
+        nn.init.kaiming_normal_(self.lda3_pool._modules['0'].weight.data)
+        nn.init.kaiming_normal_(self.lda1_fc.weight.data)
+        nn.init.kaiming_normal_(self.lda2_fc.weight.data)
+        nn.init.kaiming_normal_(self.lda3_fc.weight.data)
+        nn.init.kaiming_normal_(self.lda4_fc.weight.data)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+        x = self.layer1(x)
+
+        # the same effect as lda operation in the paper, but save much more memory
+        lda_1 = self.lda1_fc(self.lda1_pool(x).reshape(x.size(0), -1))
+        x = self.layer2(x)
+        lda_2 = self.lda2_fc(self.lda2_pool(x).reshape(x.size(0), -1))
+        x = self.layer3(x)
+        lda_3 = self.lda3_fc(self.lda3_pool(x).reshape(x.size(0), -1))
+        x = self.layer4(x)
+        lda_4 = self.lda4_fc(self.lda4_pool(x).reshape(x.size(0), -1))
+
+        vec = torch.cat((lda_1, lda_2, lda_3, lda_4), 1)
+
+        out = {}
+        out['hyper_in_feat'] = x
+        out['target_in_vec'] = vec
+
+        return out
+
+
+def resnet50_backbone(lda_out_channels, in_chn, pretrained=False, **kwargs):
+    """Constructs a ResNet-50 model_hyper.
+
+    Args:
+        pretrained (bool): If True, returns a model_hyper pre-trained on ImageNet
+    """
+    model = ResNetBackbone(lda_out_channels, in_chn, Bottleneck, [3, 4, 6, 3], **kwargs)
+    if pretrained:
+        save_model = model_zoo.load_url(model_urls['resnet50'])
+        model_dict = model.state_dict()
+        state_dict = {k: v for k, v in save_model.items() if k in model_dict.keys()}
+        model_dict.update(state_dict)
+        model.load_state_dict(model_dict)
+    else:
+        model.apply(weights_init_xavier)
+    return model
+
+
+def weights_init_xavier(m):
+    classname = m.__class__.__name__
+    # print(classname)
+    # if isinstance(m, nn.Conv2d):
+    if classname.find('Conv') != -1:
+        init.kaiming_normal_(m.weight.data)
+    elif classname.find('Linear') != -1:
+        init.kaiming_normal_(m.weight.data)
+    elif classname.find('BatchNorm2d') != -1:
+        init.uniform_(m.weight.data, 1.0, 0.02)
+        init.constant_(m.bias.data, 0.0)

eval/inference_videos.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import subprocess
+from tqdm import tqdm
+
+
+def inference_video_from_dir(input_dir, output_dir, unet_config_path, ckpt_path):
+    os.makedirs(output_dir, exist_ok=True)
+    video_names = sorted([f for f in os.listdir(input_dir) if f.endswith(".mp4")])
+    for video_name in tqdm(video_names):
+        video_path = os.path.join(input_dir, video_name)
+        audio_path = os.path.join(input_dir, video_name.replace(".mp4", "_audio.wav"))
+        video_out_path = os.path.join(output_dir, video_name.replace(".mp4", "_out.mp4"))
+        inference_command = f"python inference.py --unet_config_path {unet_config_path} --video_path {video_path} --audio_path {audio_path} --video_out_path {video_out_path} --inference_ckpt_path {ckpt_path} --seed 1247"
+        subprocess.run(inference_command, shell=True)
+
+
+if __name__ == "__main__":
+    input_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/HDTF/segmented/cross"
+    output_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/HDTF/segmented/latentsync_cross"
+    unet_config_path = "configs/unet/unet_latent_16_diffusion.yaml"
+    ckpt_path = "output/unet/train-2024_10_08-16:23:43/checkpoints/checkpoint-1920000.pt"
+
+    inference_video_from_dir(input_dir, output_dir, unet_config_path, ckpt_path)

eval/syncnet/__init__.py

@@ -0,0 +1 @@
+from .syncnet_eval import SyncNetEval

eval/syncnet/syncnet.py

@@ -0,0 +1,113 @@
+# https://github.com/joonson/syncnet_python/blob/master/SyncNetModel.py
+
+import torch
+import torch.nn as nn
+
+
+def save(model, filename):
+    with open(filename, "wb") as f:
+        torch.save(model, f)
+        print("%s saved." % filename)
+
+
+def load(filename):
+    net = torch.load(filename)
+    return net
+
+
+class S(nn.Module):
+    def __init__(self, num_layers_in_fc_layers=1024):
+        super(S, self).__init__()
+
+        self.__nFeatures__ = 24
+        self.__nChs__ = 32
+        self.__midChs__ = 32
+
+        self.netcnnaud = nn.Sequential(
+            nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=(1, 1), stride=(1, 1)),
+            nn.Conv2d(64, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
+            nn.BatchNorm2d(192),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=(3, 3), stride=(1, 2)),
+            nn.Conv2d(192, 384, kernel_size=(3, 3), padding=(1, 1)),
+            nn.BatchNorm2d(384),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(384, 256, kernel_size=(3, 3), padding=(1, 1)),
+            nn.BatchNorm2d(256),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 256, kernel_size=(3, 3), padding=(1, 1)),
+            nn.BatchNorm2d(256),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=(3, 3), stride=(2, 2)),
+            nn.Conv2d(256, 512, kernel_size=(5, 4), padding=(0, 0)),
+            nn.BatchNorm2d(512),
+            nn.ReLU(),
+        )
+
+        self.netfcaud = nn.Sequential(
+            nn.Linear(512, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Linear(512, num_layers_in_fc_layers),
+        )
+
+        self.netfclip = nn.Sequential(
+            nn.Linear(512, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Linear(512, num_layers_in_fc_layers),
+        )
+
+        self.netcnnlip = nn.Sequential(
+            nn.Conv3d(3, 96, kernel_size=(5, 7, 7), stride=(1, 2, 2), padding=0),
+            nn.BatchNorm3d(96),
+            nn.ReLU(inplace=True),
+            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2)),
+            nn.Conv3d(96, 256, kernel_size=(1, 5, 5), stride=(1, 2, 2), padding=(0, 1, 1)),
+            nn.BatchNorm3d(256),
+            nn.ReLU(inplace=True),
+            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2), padding=(0, 1, 1)),
+            nn.Conv3d(256, 256, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
+            nn.BatchNorm3d(256),
+            nn.ReLU(inplace=True),
+            nn.Conv3d(256, 256, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
+            nn.BatchNorm3d(256),
+            nn.ReLU(inplace=True),
+            nn.Conv3d(256, 256, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
+            nn.BatchNorm3d(256),
+            nn.ReLU(inplace=True),
+            nn.MaxPool3d(kernel_size=(1, 3, 3), stride=(1, 2, 2)),
+            nn.Conv3d(256, 512, kernel_size=(1, 6, 6), padding=0),
+            nn.BatchNorm3d(512),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward_aud(self, x):
+
+        mid = self.netcnnaud(x)
+        # N x ch x 24 x M
+        mid = mid.view((mid.size()[0], -1))
+        # N x (ch x 24)
+        out = self.netfcaud(mid)
+
+        return out
+
+    def forward_lip(self, x):
+
+        mid = self.netcnnlip(x)
+        mid = mid.view((mid.size()[0], -1))
+        # N x (ch x 24)
+        out = self.netfclip(mid)
+
+        return out
+
+    def forward_lipfeat(self, x):
+
+        mid = self.netcnnlip(x)
+        out = mid.view((mid.size()[0], -1))
+        # N x (ch x 24)
+
+        return out

eval/syncnet/syncnet_eval.py

@@ -0,0 +1,220 @@
+# Adapted from https://github.com/joonson/syncnet_python/blob/master/SyncNetInstance.py
+
+import torch
+import numpy
+import time, pdb, argparse, subprocess, os, math, glob
+import cv2
+import python_speech_features
+
+from scipy import signal
+from scipy.io import wavfile
+from .syncnet import S
+from shutil import rmtree
+
+
+# ==================== Get OFFSET ====================
+
+# Video 25 FPS, Audio 16000HZ
+
+
+def calc_pdist(feat1, feat2, vshift=10):
+    win_size = vshift * 2 + 1
+
+    feat2p = torch.nn.functional.pad(feat2, (0, 0, vshift, vshift))
+
+    dists = []
+
+    for i in range(0, len(feat1)):
+
+        dists.append(
+            torch.nn.functional.pairwise_distance(feat1[[i], :].repeat(win_size, 1), feat2p[i : i + win_size, :])
+        )
+
+    return dists
+
+
+# ==================== MAIN DEF ====================
+
+
+class SyncNetEval(torch.nn.Module):
+    def __init__(self, dropout=0, num_layers_in_fc_layers=1024, device="cpu"):
+        super().__init__()
+
+        self.__S__ = S(num_layers_in_fc_layers=num_layers_in_fc_layers).to(device)
+        self.device = device
+
+    def evaluate(self, video_path, temp_dir="temp", batch_size=20, vshift=15):
+
+        self.__S__.eval()
+
+        # ========== ==========
+        # Convert files
+        # ========== ==========
+
+        if os.path.exists(temp_dir):
+            rmtree(temp_dir)
+
+        os.makedirs(temp_dir)
+
+        # temp_video_path = os.path.join(temp_dir, "temp.mp4")
+        # command = f"ffmpeg -loglevel error -nostdin -y -i {video_path} -vf scale='224:224' {temp_video_path}"
+        # subprocess.call(command, shell=True)
+
+        command = (
+            f"ffmpeg -loglevel error -nostdin -y -i {video_path} -f image2 {os.path.join(temp_dir, '%06d.jpg')}"
+        )
+        subprocess.call(command, shell=True, stdout=None)
+
+        command = f"ffmpeg -loglevel error -nostdin -y -i {video_path} -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 {os.path.join(temp_dir, 'audio.wav')}"
+        subprocess.call(command, shell=True, stdout=None)
+
+        # ========== ==========
+        # Load video
+        # ========== ==========
+
+        images = []
+
+        flist = glob.glob(os.path.join(temp_dir, "*.jpg"))
+        flist.sort()
+
+        for fname in flist:
+            img_input = cv2.imread(fname)
+            img_input = cv2.resize(img_input, (224, 224))  # HARD CODED, CHANGE BEFORE RELEASE
+            images.append(img_input)
+
+        im = numpy.stack(images, axis=3)
+        im = numpy.expand_dims(im, axis=0)
+        im = numpy.transpose(im, (0, 3, 4, 1, 2))
+
+        imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float())
+
+        # ========== ==========
+        # Load audio
+        # ========== ==========
+
+        sample_rate, audio = wavfile.read(os.path.join(temp_dir, "audio.wav"))
+        mfcc = zip(*python_speech_features.mfcc(audio, sample_rate))
+        mfcc = numpy.stack([numpy.array(i) for i in mfcc])
+
+        cc = numpy.expand_dims(numpy.expand_dims(mfcc, axis=0), axis=0)
+        cct = torch.autograd.Variable(torch.from_numpy(cc.astype(float)).float())
+
+        # ========== ==========
+        # Check audio and video input length
+        # ========== ==========
+
+        # if (float(len(audio)) / 16000) != (float(len(images)) / 25):
+        #     print(
+        #         "WARNING: Audio (%.4fs) and video (%.4fs) lengths are different."
+        #         % (float(len(audio)) / 16000, float(len(images)) / 25)
+        #     )
+
+        min_length = min(len(images), math.floor(len(audio) / 640))
+
+        # ========== ==========
+        # Generate video and audio feats
+        # ========== ==========
+
+        lastframe = min_length - 5
+        im_feat = []
+        cc_feat = []
+
+        tS = time.time()
+        for i in range(0, lastframe, batch_size):
+
+            im_batch = [imtv[:, :, vframe : vframe + 5, :, :] for vframe in range(i, min(lastframe, i + batch_size))]
+            im_in = torch.cat(im_batch, 0)
+            im_out = self.__S__.forward_lip(im_in.to(self.device))
+            im_feat.append(im_out.data.cpu())
+
+            cc_batch = [
+                cct[:, :, :, vframe * 4 : vframe * 4 + 20] for vframe in range(i, min(lastframe, i + batch_size))
+            ]
+            cc_in = torch.cat(cc_batch, 0)
+            cc_out = self.__S__.forward_aud(cc_in.to(self.device))
+            cc_feat.append(cc_out.data.cpu())
+
+        im_feat = torch.cat(im_feat, 0)
+        cc_feat = torch.cat(cc_feat, 0)
+
+        # ========== ==========
+        # Compute offset
+        # ========== ==========
+
+        dists = calc_pdist(im_feat, cc_feat, vshift=vshift)
+        mean_dists = torch.mean(torch.stack(dists, 1), 1)
+
+        min_dist, minidx = torch.min(mean_dists, 0)
+
+        av_offset = vshift - minidx
+        conf = torch.median(mean_dists) - min_dist
+
+        fdist = numpy.stack([dist[minidx].numpy() for dist in dists])
+        # fdist   = numpy.pad(fdist, (3,3), 'constant', constant_values=15)
+        fconf = torch.median(mean_dists).numpy() - fdist
+        framewise_conf = signal.medfilt(fconf, kernel_size=9)
+
+        # numpy.set_printoptions(formatter={"float": "{: 0.3f}".format})
+        rmtree(temp_dir)
+        return av_offset.item(), min_dist.item(), conf.item()
+
+    def extract_feature(self, opt, videofile):
+
+        self.__S__.eval()
+
+        # ========== ==========
+        # Load video
+        # ========== ==========
+        cap = cv2.VideoCapture(videofile)
+
+        frame_num = 1
+        images = []
+        while frame_num:
+            frame_num += 1
+            ret, image = cap.read()
+            if ret == 0:
+                break
+
+            images.append(image)
+
+        im = numpy.stack(images, axis=3)
+        im = numpy.expand_dims(im, axis=0)
+        im = numpy.transpose(im, (0, 3, 4, 1, 2))
+
+        imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float())
+
+        # ========== ==========
+        # Generate video feats
+        # ========== ==========
+
+        lastframe = len(images) - 4
+        im_feat = []
+
+        tS = time.time()
+        for i in range(0, lastframe, opt.batch_size):
+
+            im_batch = [
+                imtv[:, :, vframe : vframe + 5, :, :] for vframe in range(i, min(lastframe, i + opt.batch_size))
+            ]
+            im_in = torch.cat(im_batch, 0)
+            im_out = self.__S__.forward_lipfeat(im_in.to(self.device))
+            im_feat.append(im_out.data.cpu())
+
+        im_feat = torch.cat(im_feat, 0)
+
+        # ========== ==========
+        # Compute offset
+        # ========== ==========
+
+        print("Compute time %.3f sec." % (time.time() - tS))
+
+        return im_feat
+
+    def loadParameters(self, path):
+        loaded_state = torch.load(path, map_location=lambda storage, loc: storage)
+
+        self_state = self.__S__.state_dict()
+
+        for name, param in loaded_state.items():
+
+            self_state[name].copy_(param)

eval/syncnet_detect.py

@@ -0,0 +1,251 @@
+# Adapted from https://github.com/joonson/syncnet_python/blob/master/run_pipeline.py
+
+import os, pdb, subprocess, glob, cv2
+import numpy as np
+from shutil import rmtree
+import torch
+
+from scenedetect.video_manager import VideoManager
+from scenedetect.scene_manager import SceneManager
+from scenedetect.stats_manager import StatsManager
+from scenedetect.detectors import ContentDetector
+
+from scipy.interpolate import interp1d
+from scipy.io import wavfile
+from scipy import signal
+
+from eval.detectors import S3FD
+
+
+class SyncNetDetector:
+    def __init__(self, device, detect_results_dir="detect_results"):
+        self.s3f_detector = S3FD(device=device)
+        self.detect_results_dir = detect_results_dir
+
+    def __call__(self, video_path: str, min_track=50, scale=False):
+        crop_dir = os.path.join(self.detect_results_dir, "crop")
+        video_dir = os.path.join(self.detect_results_dir, "video")
+        frames_dir = os.path.join(self.detect_results_dir, "frames")
+        temp_dir = os.path.join(self.detect_results_dir, "temp")
+
+        # ========== DELETE EXISTING DIRECTORIES ==========
+        if os.path.exists(crop_dir):
+            rmtree(crop_dir)
+
+        if os.path.exists(video_dir):
+            rmtree(video_dir)
+
+        if os.path.exists(frames_dir):
+            rmtree(frames_dir)
+
+        if os.path.exists(temp_dir):
+            rmtree(temp_dir)
+
+        # ========== MAKE NEW DIRECTORIES ==========
+
+        os.makedirs(crop_dir)
+        os.makedirs(video_dir)
+        os.makedirs(frames_dir)
+        os.makedirs(temp_dir)
+
+        # ========== CONVERT VIDEO AND EXTRACT FRAMES ==========
+
+        if scale:
+            scaled_video_path = os.path.join(video_dir, "scaled.mp4")
+            command = f"ffmpeg -loglevel error -y -nostdin -i {video_path} -vf scale='224:224' {scaled_video_path}"
+            subprocess.run(command, shell=True)
+            video_path = scaled_video_path
+
+        command = f"ffmpeg -y -nostdin -loglevel error -i {video_path} -qscale:v 2 -async 1 -r 25 {os.path.join(video_dir, 'video.mp4')}"
+        subprocess.run(command, shell=True, stdout=None)
+
+        command = f"ffmpeg -y -nostdin -loglevel error -i {os.path.join(video_dir, 'video.mp4')} -qscale:v 2 -f image2 {os.path.join(frames_dir, '%06d.jpg')}"
+        subprocess.run(command, shell=True, stdout=None)
+
+        command = f"ffmpeg -y -nostdin -loglevel error -i {os.path.join(video_dir, 'video.mp4')} -ac 1 -vn -acodec pcm_s16le -ar 16000 {os.path.join(video_dir, 'audio.wav')}"
+        subprocess.run(command, shell=True, stdout=None)
+
+        faces = self.detect_face(frames_dir)
+
+        scene = self.scene_detect(video_dir)
+
+        # Face tracking
+        alltracks = []
+
+        for shot in scene:
+            if shot[1].frame_num - shot[0].frame_num >= min_track:
+                alltracks.extend(self.track_face(faces[shot[0].frame_num : shot[1].frame_num], min_track=min_track))
+
+        # Face crop
+        for ii, track in enumerate(alltracks):
+            self.crop_video(track, os.path.join(crop_dir, "%05d" % ii), frames_dir, 25, temp_dir, video_dir)
+
+        rmtree(temp_dir)
+
+    def scene_detect(self, video_dir):
+        video_manager = VideoManager([os.path.join(video_dir, "video.mp4")])
+        stats_manager = StatsManager()
+        scene_manager = SceneManager(stats_manager)
+        # Add ContentDetector algorithm (constructor takes detector options like threshold).
+        scene_manager.add_detector(ContentDetector())
+        base_timecode = video_manager.get_base_timecode()
+
+        video_manager.set_downscale_factor()
+
+        video_manager.start()
+
+        scene_manager.detect_scenes(frame_source=video_manager)
+
+        scene_list = scene_manager.get_scene_list(base_timecode)
+
+        if scene_list == []:
+            scene_list = [(video_manager.get_base_timecode(), video_manager.get_current_timecode())]
+
+        return scene_list
+
+    def track_face(self, scenefaces, num_failed_det=25, min_track=50, min_face_size=100):
+
+        iouThres = 0.5  # Minimum IOU between consecutive face detections
+        tracks = []
+
+        while True:
+            track = []
+            for framefaces in scenefaces:
+                for face in framefaces:
+                    if track == []:
+                        track.append(face)
+                        framefaces.remove(face)
+                    elif face["frame"] - track[-1]["frame"] <= num_failed_det:
+                        iou = bounding_box_iou(face["bbox"], track[-1]["bbox"])
+                        if iou > iouThres:
+                            track.append(face)
+                            framefaces.remove(face)
+                            continue
+                    else:
+                        break
+
+            if track == []:
+                break
+            elif len(track) > min_track:
+
+                framenum = np.array([f["frame"] for f in track])
+                bboxes = np.array([np.array(f["bbox"]) for f in track])
+
+                frame_i = np.arange(framenum[0], framenum[-1] + 1)
+
+                bboxes_i = []
+                for ij in range(0, 4):
+                    interpfn = interp1d(framenum, bboxes[:, ij])
+                    bboxes_i.append(interpfn(frame_i))
+                bboxes_i = np.stack(bboxes_i, axis=1)
+
+                if (
+                    max(np.mean(bboxes_i[:, 2] - bboxes_i[:, 0]), np.mean(bboxes_i[:, 3] - bboxes_i[:, 1]))
+                    > min_face_size
+                ):
+                    tracks.append({"frame": frame_i, "bbox": bboxes_i})
+
+        return tracks
+
+    def detect_face(self, frames_dir, facedet_scale=0.25):
+        flist = glob.glob(os.path.join(frames_dir, "*.jpg"))
+        flist.sort()
+
+        dets = []
+
+        for fidx, fname in enumerate(flist):
+            image = cv2.imread(fname)
+
+            image_np = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            bboxes = self.s3f_detector.detect_faces(image_np, conf_th=0.9, scales=[facedet_scale])
+
+            dets.append([])
+            for bbox in bboxes:
+                dets[-1].append({"frame": fidx, "bbox": (bbox[:-1]).tolist(), "conf": bbox[-1]})
+
+        return dets
+
+    def crop_video(self, track, cropfile, frames_dir, frame_rate, temp_dir, video_dir, crop_scale=0.4):
+
+        flist = glob.glob(os.path.join(frames_dir, "*.jpg"))
+        flist.sort()
+
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        vOut = cv2.VideoWriter(cropfile + "t.mp4", fourcc, frame_rate, (224, 224))
+
+        dets = {"x": [], "y": [], "s": []}
+
+        for det in track["bbox"]:
+
+            dets["s"].append(max((det[3] - det[1]), (det[2] - det[0])) / 2)
+            dets["y"].append((det[1] + det[3]) / 2)  # crop center x
+            dets["x"].append((det[0] + det[2]) / 2)  # crop center y
+
+        # Smooth detections
+        dets["s"] = signal.medfilt(dets["s"], kernel_size=13)
+        dets["x"] = signal.medfilt(dets["x"], kernel_size=13)
+        dets["y"] = signal.medfilt(dets["y"], kernel_size=13)
+
+        for fidx, frame in enumerate(track["frame"]):
+
+            cs = crop_scale
+
+            bs = dets["s"][fidx]  # Detection box size
+            bsi = int(bs * (1 + 2 * cs))  # Pad videos by this amount
+
+            image = cv2.imread(flist[frame])
+
+            frame = np.pad(image, ((bsi, bsi), (bsi, bsi), (0, 0)), "constant", constant_values=(110, 110))
+            my = dets["y"][fidx] + bsi  # BBox center Y
+            mx = dets["x"][fidx] + bsi  # BBox center X
+
+            face = frame[int(my - bs) : int(my + bs * (1 + 2 * cs)), int(mx - bs * (1 + cs)) : int(mx + bs * (1 + cs))]
+
+            vOut.write(cv2.resize(face, (224, 224)))
+
+        audiotmp = os.path.join(temp_dir, "audio.wav")
+        audiostart = (track["frame"][0]) / frame_rate
+        audioend = (track["frame"][-1] + 1) / frame_rate
+
+        vOut.release()
+
+        # ========== CROP AUDIO FILE ==========
+
+        command = "ffmpeg -y -nostdin -loglevel error -i %s -ss %.3f -to %.3f %s" % (
+            os.path.join(video_dir, "audio.wav"),
+            audiostart,
+            audioend,
+            audiotmp,
+        )
+        output = subprocess.run(command, shell=True, stdout=None)
+
+        sample_rate, audio = wavfile.read(audiotmp)
+
+        # ========== COMBINE AUDIO AND VIDEO FILES ==========
+
+        command = "ffmpeg -y -nostdin -loglevel error -i %st.mp4 -i %s -c:v copy -c:a aac %s.mp4" % (
+            cropfile,
+            audiotmp,
+            cropfile,
+        )
+        output = subprocess.run(command, shell=True, stdout=None)
+
+        os.remove(cropfile + "t.mp4")
+
+        return {"track": track, "proc_track": dets}
+
+
+def bounding_box_iou(boxA, boxB):
+    xA = max(boxA[0], boxB[0])
+    yA = max(boxA[1], boxB[1])
+    xB = min(boxA[2], boxB[2])
+    yB = min(boxA[3], boxB[3])
+
+    interArea = max(0, xB - xA) * max(0, yB - yA)
+
+    boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
+    boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
+
+    iou = interArea / float(boxAArea + boxBArea - interArea)
+
+    return iou

inference.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+python -m scripts.inference \
+    --unet_config_path "configs/unet/second_stage.yaml" \
+    --inference_ckpt_path "checkpoints/latentsync_unet.pt" \
+    --guidance_scale 1.0 \
+    --video_path "assets/demo1_video.mp4" \
+    --audio_path "assets/demo1_audio.wav" \
+    --video_out_path "video_out.mp4"

pipelines/lipsync_pipeline.py

@@ -0,0 +1,470 @@
+# Adapted from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/pipelines/pipeline_animation.py
+
+import inspect
+import os
+import shutil
+from typing import Callable, List, Optional, Union
+import subprocess
+
+import numpy as np
+import torch
+import torchvision
+
+from diffusers.utils import is_accelerate_available
+from packaging import version
+
+from diffusers.configuration_utils import FrozenDict
+from diffusers.models import AutoencoderKL
+from diffusers.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import (
+    DDIMScheduler,
+    DPMSolverMultistepScheduler,
+    EulerAncestralDiscreteScheduler,
+    EulerDiscreteScheduler,
+    LMSDiscreteScheduler,
+    PNDMScheduler,
+)
+from diffusers.utils import deprecate, logging
+
+from einops import rearrange
+
+from ..models.unet import UNet3DConditionModel
+from ..utils.image_processor import ImageProcessor
+from ..utils.util import read_video, read_audio, write_video
+from ..whisper.audio2feature import Audio2Feature
+import tqdm
+import soundfile as sf
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class LipsyncPipeline(DiffusionPipeline):
+    _optional_components = []
+
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        audio_encoder: Audio2Feature,
+        unet: UNet3DConditionModel,
+        scheduler: Union[
+            DDIMScheduler,
+            PNDMScheduler,
+            LMSDiscreteScheduler,
+            EulerDiscreteScheduler,
+            EulerAncestralDiscreteScheduler,
+            DPMSolverMultistepScheduler,
+        ],
+    ):
+        super().__init__()
+
+        if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
+                f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure "
+                "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
+                " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
+                " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
+                " file"
+            )
+            deprecate("steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["steps_offset"] = 1
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True:
+            deprecation_message = (
+                f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
+                " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
+                " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
+                " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
+                " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
+            )
+            deprecate("clip_sample not set", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(scheduler.config)
+            new_config["clip_sample"] = False
+            scheduler._internal_dict = FrozenDict(new_config)
+
+        is_unet_version_less_0_9_0 = hasattr(unet.config, "_diffusers_version") and version.parse(
+            version.parse(unet.config._diffusers_version).base_version
+        ) < version.parse("0.9.0.dev0")
+        is_unet_sample_size_less_64 = hasattr(unet.config, "sample_size") and unet.config.sample_size < 64
+        if is_unet_version_less_0_9_0 and is_unet_sample_size_less_64:
+            deprecation_message = (
+                "The configuration file of the unet has set the default `sample_size` to smaller than"
+                " 64 which seems highly unlikely. If your checkpoint is a fine-tuned version of any of the"
+                " following: \n- CompVis/stable-diffusion-v1-4 \n- CompVis/stable-diffusion-v1-3 \n-"
+                " CompVis/stable-diffusion-v1-2 \n- CompVis/stable-diffusion-v1-1 \n- runwayml/stable-diffusion-v1-5"
+                " \n- runwayml/stable-diffusion-inpainting \n you should change 'sample_size' to 64 in the"
+                " configuration file. Please make sure to update the config accordingly as leaving `sample_size=32`"
+                " in the config might lead to incorrect results in future versions. If you have downloaded this"
+                " checkpoint from the Hugging Face Hub, it would be very nice if you could open a Pull request for"
+                " the `unet/config.json` file"
+            )
+            deprecate("sample_size<64", "1.0.0", deprecation_message, standard_warn=False)
+            new_config = dict(unet.config)
+            new_config["sample_size"] = 64
+            unet._internal_dict = FrozenDict(new_config)
+
+        self.register_modules(
+            vae=vae,
+            audio_encoder=audio_encoder,
+            unet=unet,
+            scheduler=scheduler,
+        )
+
+        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+
+        self.set_progress_bar_config(desc="Steps")
+
+    def enable_vae_slicing(self):
+        self.vae.enable_slicing()
+
+    def disable_vae_slicing(self):
+        self.vae.disable_slicing()
+
+    def enable_sequential_cpu_offload(self, gpu_id=0):
+        if is_accelerate_available():
+            from accelerate import cpu_offload
+        else:
+            raise ImportError("Please install accelerate via `pip install accelerate`")
+
+        device = torch.device(f"cuda:{gpu_id}")
+
+        for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
+            if cpu_offloaded_model is not None:
+                cpu_offload(cpu_offloaded_model, device)
+
+    @property
+    def _execution_device(self):
+        if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"):
+            return self.device
+        for module in self.unet.modules():
+            if (
+                hasattr(module, "_hf_hook")
+                and hasattr(module._hf_hook, "execution_device")
+                and module._hf_hook.execution_device is not None
+            ):
+                return torch.device(module._hf_hook.execution_device)
+        return self.device
+
+    def decode_latents(self, latents):
+        latents = latents / self.vae.config.scaling_factor + self.vae.config.shift_factor
+        latents = rearrange(latents, "b c f h w -> (b f) c h w")
+        decoded_latents = self.vae.decode(latents).sample
+        return decoded_latents
+
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def check_inputs(self, height, width, callback_steps):
+        assert height == width, "Height and width must be equal"
+
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+    def prepare_latents(self, batch_size, num_frames, num_channels_latents, height, width, dtype, device, generator):
+        shape = (
+            batch_size,
+            num_channels_latents,
+            1,
+            height // self.vae_scale_factor,
+            width // self.vae_scale_factor,
+        )
+        rand_device = "cpu" if device.type == "mps" else device
+        latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device)
+        latents = latents.repeat(1, 1, num_frames, 1, 1)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+        return latents
+
+    def prepare_mask_latents(
+        self, mask, masked_image, height, width, dtype, device, generator, do_classifier_free_guidance
+    ):
+        # resize the mask to latents shape as we concatenate the mask to the latents
+        # we do that before converting to dtype to avoid breaking in case we're using cpu_offload
+        # and half precision
+        mask = torch.nn.functional.interpolate(
+            mask, size=(height // self.vae_scale_factor, width // self.vae_scale_factor)
+        )
+        masked_image = masked_image.to(device=device, dtype=dtype)
+
+        # encode the mask image into latents space so we can concatenate it to the latents
+        masked_image_latents = self.vae.encode(masked_image).latent_dist.sample(generator=generator)
+        masked_image_latents = (masked_image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+
+        # aligning device to prevent device errors when concating it with the latent model input
+        masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)
+        mask = mask.to(device=device, dtype=dtype)
+
+        # assume batch size = 1
+        mask = rearrange(mask, "f c h w -> 1 c f h w")
+        masked_image_latents = rearrange(masked_image_latents, "f c h w -> 1 c f h w")
+
+        mask = torch.cat([mask] * 2) if do_classifier_free_guidance else mask
+        masked_image_latents = (
+            torch.cat([masked_image_latents] * 2) if do_classifier_free_guidance else masked_image_latents
+        )
+        return mask, masked_image_latents
+
+    def prepare_image_latents(self, images, device, dtype, generator, do_classifier_free_guidance):
+        images = images.to(device=device, dtype=dtype)
+        image_latents = self.vae.encode(images).latent_dist.sample(generator=generator)
+        image_latents = (image_latents - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+        image_latents = rearrange(image_latents, "f c h w -> 1 c f h w")
+        image_latents = torch.cat([image_latents] * 2) if do_classifier_free_guidance else image_latents
+
+        return image_latents
+
+    def set_progress_bar_config(self, **kwargs):
+        if not hasattr(self, "_progress_bar_config"):
+            self._progress_bar_config = {}
+        self._progress_bar_config.update(kwargs)
+
+    @staticmethod
+    def paste_surrounding_pixels_back(decoded_latents, pixel_values, masks, device, weight_dtype):
+        # Paste the surrounding pixels back, because we only want to change the mouth region
+        pixel_values = pixel_values.to(device=device, dtype=weight_dtype)
+        masks = masks.to(device=device, dtype=weight_dtype)
+        combined_pixel_values = decoded_latents * masks + pixel_values * (1 - masks)
+        return combined_pixel_values
+
+    @staticmethod
+    def pixel_values_to_images(pixel_values: torch.Tensor):
+        pixel_values = rearrange(pixel_values, "f c h w -> f h w c")
+        pixel_values = (pixel_values / 2 + 0.5).clamp(0, 1)
+        images = (pixel_values * 255).to(torch.uint8)
+        images = images.cpu().numpy()
+        return images
+
+    def affine_transform_video(self, video_path):
+        video_frames = read_video(video_path, use_decord=False)
+        faces = []
+        boxes = []
+        affine_matrices = []
+        print(f"Affine transforming {len(video_frames)} faces...")
+        for frame in tqdm.tqdm(video_frames):
+            face, box, affine_matrix = self.image_processor.affine_transform(frame)
+            faces.append(face)
+            boxes.append(box)
+            affine_matrices.append(affine_matrix)
+
+        faces = torch.stack(faces)
+        return faces, video_frames, boxes, affine_matrices
+
+    def restore_video(self, faces, video_frames, boxes, affine_matrices):
+        video_frames = video_frames[: faces.shape[0]]
+        out_frames = []
+        for index, face in enumerate(faces):
+            x1, y1, x2, y2 = boxes[index]
+            height = int(y2 - y1)
+            width = int(x2 - x1)
+            face = torchvision.transforms.functional.resize(face, size=(height, width), antialias=True)
+            face = rearrange(face, "c h w -> h w c")
+            face = (face / 2 + 0.5).clamp(0, 1)
+            face = (face * 255).to(torch.uint8).cpu().numpy()
+            out_frame = self.image_processor.restorer.restore_img(video_frames[index], face, affine_matrices[index])
+            out_frames.append(out_frame)
+        return np.stack(out_frames, axis=0)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        video_path: str,
+        audio_path: str,
+        video_out_path: str,
+        video_mask_path: str = None,
+        num_frames: int = 16,
+        video_fps: int = 25,
+        audio_sample_rate: int = 16000,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 20,
+        guidance_scale: float = 1.5,
+        weight_dtype: Optional[torch.dtype] = torch.float16,
+        eta: float = 0.0,
+        mask: str = "fix_mask",
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: Optional[int] = 1,
+        **kwargs,
+    ):
+        is_train = self.unet.training
+        self.unet.eval()
+
+        # 0. Define call parameters
+        batch_size = 1
+        device = self._execution_device
+        self.image_processor = ImageProcessor(height, mask=mask, device="cuda")
+        self.set_progress_bar_config(desc=f"Sample frames: {num_frames}")
+
+        video_frames, original_video_frames, boxes, affine_matrices = self.affine_transform_video(video_path)
+        audio_samples = read_audio(audio_path)
+
+        # 1. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+
+        # 2. Check inputs
+        self.check_inputs(height, width, callback_steps)
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. set timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        # 4. Prepare extra step kwargs.
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        self.video_fps = video_fps
+
+        if self.unet.add_audio_layer:
+            whisper_feature = self.audio_encoder.audio2feat(audio_path)
+            whisper_chunks = self.audio_encoder.feature2chunks(feature_array=whisper_feature, fps=video_fps)
+
+            num_inferences = min(len(video_frames), len(whisper_chunks)) // num_frames
+        else:
+            num_inferences = len(video_frames) // num_frames
+
+        synced_video_frames = []
+        masked_video_frames = []
+
+        num_channels_latents = self.vae.config.latent_channels
+
+        # Prepare latent variables
+        all_latents = self.prepare_latents(
+            batch_size,
+            num_frames * num_inferences,
+            num_channels_latents,
+            height,
+            width,
+            weight_dtype,
+            device,
+            generator,
+        )
+
+        for i in tqdm.tqdm(range(num_inferences), desc="Doing inference..."):
+            if self.unet.add_audio_layer:
+                audio_embeds = torch.stack(whisper_chunks[i * num_frames : (i + 1) * num_frames])
+                audio_embeds = audio_embeds.to(device, dtype=weight_dtype)
+                if do_classifier_free_guidance:
+                    empty_audio_embeds = torch.zeros_like(audio_embeds)
+                    audio_embeds = torch.cat([empty_audio_embeds, audio_embeds])
+            else:
+                audio_embeds = None
+            inference_video_frames = video_frames[i * num_frames : (i + 1) * num_frames]
+            latents = all_latents[:, :, i * num_frames : (i + 1) * num_frames]
+            pixel_values, masked_pixel_values, masks = self.image_processor.prepare_masks_and_masked_images(
+                inference_video_frames, affine_transform=False
+            )
+
+            # 7. Prepare mask latent variables
+            mask_latents, masked_image_latents = self.prepare_mask_latents(
+                masks,
+                masked_pixel_values,
+                height,
+                width,
+                weight_dtype,
+                device,
+                generator,
+                do_classifier_free_guidance,
+            )
+
+            # 8. Prepare image latents
+            image_latents = self.prepare_image_latents(
+                pixel_values,
+                device,
+                weight_dtype,
+                generator,
+                do_classifier_free_guidance,
+            )
+
+            # 9. Denoising loop
+            num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+            with self.progress_bar(total=num_inference_steps) as progress_bar:
+                for j, t in enumerate(timesteps):
+                    # expand the latents if we are doing classifier free guidance
+                    latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+
+                    # concat latents, mask, masked_image_latents in the channel dimension
+                    latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+                    latent_model_input = torch.cat(
+                        [latent_model_input, mask_latents, masked_image_latents, image_latents], dim=1
+                    )
+
+                    # predict the noise residual
+                    noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=audio_embeds).sample
+
+                    # perform guidance
+                    if do_classifier_free_guidance:
+                        noise_pred_uncond, noise_pred_audio = noise_pred.chunk(2)
+                        noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_audio - noise_pred_uncond)
+
+                    # compute the previous noisy sample x_t -> x_t-1
+                    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
+
+                    # call the callback, if provided
+                    if j == len(timesteps) - 1 or ((j + 1) > num_warmup_steps and (j + 1) % self.scheduler.order == 0):
+                        progress_bar.update()
+                        if callback is not None and j % callback_steps == 0:
+                            callback(j, t, latents)
+
+            # Recover the pixel values
+            decoded_latents = self.decode_latents(latents)
+            decoded_latents = self.paste_surrounding_pixels_back(
+                decoded_latents, pixel_values, 1 - masks, device, weight_dtype
+            )
+            synced_video_frames.append(decoded_latents)
+            masked_video_frames.append(masked_pixel_values)
+
+        synced_video_frames = self.restore_video(
+            torch.cat(synced_video_frames), original_video_frames, boxes, affine_matrices
+        )
+        masked_video_frames = self.restore_video(
+            torch.cat(masked_video_frames), original_video_frames, boxes, affine_matrices
+        )
+
+        audio_samples_remain_length = int(synced_video_frames.shape[0] / video_fps * audio_sample_rate)
+        audio_samples = audio_samples[:audio_samples_remain_length].cpu().numpy()
+
+        if is_train:
+            self.unet.train()
+
+        temp_dir = "temp"
+        if os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir)
+        os.makedirs(temp_dir, exist_ok=True)
+
+        write_video(os.path.join(temp_dir, "video.mp4"), synced_video_frames, fps=25)
+        # write_video(video_mask_path, masked_video_frames, fps=25)
+
+        sf.write(os.path.join(temp_dir, "audio.wav"), audio_samples, audio_sample_rate)
+
+        command = f"ffmpeg -y -loglevel error -nostdin -i {os.path.join(temp_dir, 'video.mp4')} -i {os.path.join(temp_dir, 'audio.wav')} -c:v libx264 -c:a aac -q:v 0 -q:a 0 {video_out_path}"
+        subprocess.run(command, shell=True)

predict.py

@@ -0,0 +1,60 @@
+# Prediction interface for Cog ⚙️
+# https://cog.run/python
+
+from cog import BasePredictor, Input, Path
+import os
+import time
+import subprocess
+
+MODEL_CACHE = "checkpoints"
+MODEL_URL = "https://weights.replicate.delivery/default/chunyu-li/LatentSync/model.tar"
+
+def download_weights(url, dest):
+    start = time.time()
+    print("downloading url: ", url)
+    print("downloading to: ", dest)
+    subprocess.check_call(["pget", "-xf", url, dest], close_fds=False)
+    print("downloading took: ", time.time() - start)
+
+class Predictor(BasePredictor):
+    def setup(self) -> None:
+        """Load the model into memory to make running multiple predictions efficient"""
+        # Download the model weights
+        if not os.path.exists(MODEL_CACHE):
+            download_weights(MODEL_URL, MODEL_CACHE)
+
+        # Soft links for the auxiliary models
+        os.system("mkdir -p ~/.cache/torch/hub/checkpoints")
+        os.system("ln -s $(pwd)/checkpoints/auxiliary/2DFAN4-cd938726ad.zip ~/.cache/torch/hub/checkpoints/2DFAN4-cd938726ad.zip")
+        os.system("ln -s $(pwd)/checkpoints/auxiliary/s3fd-619a316812.pth ~/.cache/torch/hub/checkpoints/s3fd-619a316812.pth")
+        os.system("ln -s $(pwd)/checkpoints/auxiliary/vgg16-397923af.pth ~/.cache/torch/hub/checkpoints/vgg16-397923af.pth")
+
+    def predict(
+        self,
+        video: Path = Input(
+            description="Input video", default=None
+        ),
+        audio: Path = Input(
+            description="Input audio to ", default=None
+        ),
+        guidance_scale: float = Input(
+            description="Guidance scale", ge=0, le=10, default=1.0
+        ),
+        seed: int = Input(
+            description="Set to 0 for Random seed", default=0
+        )
+    ) -> Path:
+        """Run a single prediction on the model"""
+        if seed <= 0:
+            seed = int.from_bytes(os.urandom(2), "big")
+        print(f"Using seed: {seed}")
+
+        video_path = str(video)
+        audio_path = str(audio)
+        config_path = "configs/unet/second_stage.yaml"
+        ckpt_path = "checkpoints/latentsync_unet.pt"
+        output_path = "/tmp/video_out.mp4"
+    
+        # Run the following command:
+        os.system(f"python -m scripts.inference --unet_config_path {config_path} --inference_ckpt_path {ckpt_path} --guidance_scale {str(guidance_scale)} --video_path {video_path} --audio_path {audio_path} --video_out_path {output_path} --seed {seed}")    
+        return Path(output_path)

preprocess/affine_transform.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from latentsync.utils.util import read_video, write_video
+from latentsync.utils.image_processor import ImageProcessor
+import torch
+from einops import rearrange
+import os
+import tqdm
+import subprocess
+from multiprocessing import Process
+import shutil
+
+paths = []
+
+
+def gather_video_paths(input_dir, output_dir):
+    for video in sorted(os.listdir(input_dir)):
+        if video.endswith(".mp4"):
+            video_input = os.path.join(input_dir, video)
+            video_output = os.path.join(output_dir, video)
+            if os.path.isfile(video_output):
+                continue
+            paths.append((video_input, video_output))
+        elif os.path.isdir(os.path.join(input_dir, video)):
+            gather_video_paths(os.path.join(input_dir, video), os.path.join(output_dir, video))
+
+
+class FaceDetector:
+    def __init__(self, resolution: int = 512, device: str = "cpu"):
+        self.image_processor = ImageProcessor(resolution, "fix_mask", device)
+
+    def affine_transform_video(self, video_path):
+        video_frames = read_video(video_path, change_fps=False)
+        results = []
+        for frame in video_frames:
+            frame, _, _ = self.image_processor.affine_transform(frame)
+            results.append(frame)
+        results = torch.stack(results)
+
+        results = rearrange(results, "f c h w -> f h w c").numpy()
+        return results
+
+    def close(self):
+        self.image_processor.close()
+
+
+def combine_video_audio(video_frames, video_input_path, video_output_path, process_temp_dir):
+    video_name = os.path.basename(video_input_path)[:-4]
+    audio_temp = os.path.join(process_temp_dir, f"{video_name}_temp.wav")
+    video_temp = os.path.join(process_temp_dir, f"{video_name}_temp.mp4")
+
+    write_video(video_temp, video_frames, fps=25)
+
+    command = f"ffmpeg -y -loglevel error -i {video_input_path} -q:a 0 -map a {audio_temp}"
+    subprocess.run(command, shell=True)
+
+    os.makedirs(os.path.dirname(video_output_path), exist_ok=True)
+    command = f"ffmpeg -y -loglevel error -i {video_temp} -i {audio_temp} -c:v libx264 -c:a aac -map 0:v -map 1:a -q:v 0 -q:a 0 {video_output_path}"
+    subprocess.run(command, shell=True)
+
+    os.remove(audio_temp)
+    os.remove(video_temp)
+
+
+def func(paths, process_temp_dir, device_id, resolution):
+    os.makedirs(process_temp_dir, exist_ok=True)
+    face_detector = FaceDetector(resolution, f"cuda:{device_id}")
+
+    for video_input, video_output in paths:
+        if os.path.isfile(video_output):
+            continue
+        try:
+            video_frames = face_detector.affine_transform_video(video_input)
+        except Exception as e:  # Handle the exception of face not detcted
+            print(f"Exception: {e} - {video_input}")
+            continue
+
+        os.makedirs(os.path.dirname(video_output), exist_ok=True)
+        combine_video_audio(video_frames, video_input, video_output, process_temp_dir)
+        print(f"Saved: {video_output}")
+
+    face_detector.close()
+
+
+def split(a, n):
+    k, m = divmod(len(a), n)
+    return (a[i * k + min(i, m) : (i + 1) * k + min(i + 1, m)] for i in range(n))
+
+
+def affine_transform_multi_gpus(input_dir, output_dir, temp_dir, resolution, num_workers):
+    print(f"Recursively gathering video paths of {input_dir} ...")
+    gather_video_paths(input_dir, output_dir)
+    num_devices = torch.cuda.device_count()
+    if num_devices == 0:
+        raise RuntimeError("No GPUs found")
+
+    if os.path.exists(temp_dir):
+        shutil.rmtree(temp_dir)
+    os.makedirs(temp_dir, exist_ok=True)
+
+    split_paths = list(split(paths, num_workers * num_devices))
+
+    processes = []
+
+    for i in range(num_devices):
+        for j in range(num_workers):
+            process_index = i * num_workers + j
+            process = Process(
+                target=func, args=(split_paths[process_index], os.path.join(temp_dir, f"process_{i}"), i, resolution)
+            )
+            process.start()
+            processes.append(process)
+
+    for process in processes:
+        process.join()
+
+
+if __name__ == "__main__":
+    input_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/avatars/resampled/train"
+    output_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/avatars/affine_transformed/train"
+    temp_dir = "temp"
+    resolution = 256
+    num_workers = 10  # How many processes per device
+
+    affine_transform_multi_gpus(input_dir, output_dir, temp_dir, resolution, num_workers)

preprocess/data_processing_pipeline.py

@@ -0,0 +1,85 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+from preprocess.affine_transform import affine_transform_multi_gpus
+from preprocess.remove_broken_videos import remove_broken_videos_multiprocessing
+from preprocess.detect_shot import detect_shot_multiprocessing
+from preprocess.filter_high_resolution import filter_high_resolution_multiprocessing
+from preprocess.resample_fps_hz import resample_fps_hz_multiprocessing
+from preprocess.segment_videos import segment_videos_multiprocessing
+from preprocess.sync_av import sync_av_multi_gpus
+from preprocess.filter_visual_quality import filter_visual_quality_multi_gpus
+from preprocess.remove_incorrect_affined import remove_incorrect_affined_multiprocessing
+
+
+def data_processing_pipeline(
+    total_num_workers, per_gpu_num_workers, resolution, sync_conf_threshold, temp_dir, input_dir
+):
+    print("Removing broken videos...")
+    remove_broken_videos_multiprocessing(input_dir, total_num_workers)
+
+    print("Resampling FPS hz...")
+    resampled_dir = os.path.join(os.path.dirname(input_dir), "resampled")
+    resample_fps_hz_multiprocessing(input_dir, resampled_dir, total_num_workers)
+
+    print("Detecting shot...")
+    shot_dir = os.path.join(os.path.dirname(input_dir), "shot")
+    detect_shot_multiprocessing(resampled_dir, shot_dir, total_num_workers)
+
+    print("Segmenting videos...")
+    segmented_dir = os.path.join(os.path.dirname(input_dir), "segmented")
+    segment_videos_multiprocessing(shot_dir, segmented_dir, total_num_workers)
+
+    print("Filtering high resolution...")
+    high_resolution_dir = os.path.join(os.path.dirname(input_dir), "high_resolution")
+    filter_high_resolution_multiprocessing(segmented_dir, high_resolution_dir, resolution, total_num_workers)
+
+    print("Affine transforming videos...")
+    affine_transformed_dir = os.path.join(os.path.dirname(input_dir), "affine_transformed")
+    affine_transform_multi_gpus(
+        high_resolution_dir, affine_transformed_dir, temp_dir, resolution, per_gpu_num_workers // 2
+    )
+
+    print("Removing incorrect affined videos...")
+    remove_incorrect_affined_multiprocessing(affine_transformed_dir, total_num_workers)
+
+    print("Syncing audio and video...")
+    av_synced_dir = os.path.join(os.path.dirname(input_dir), f"av_synced_{sync_conf_threshold}")
+    sync_av_multi_gpus(affine_transformed_dir, av_synced_dir, temp_dir, per_gpu_num_workers, sync_conf_threshold)
+
+    print("Filtering visual quality...")
+    high_visual_quality_dir = os.path.join(os.path.dirname(input_dir), "high_visual_quality")
+    filter_visual_quality_multi_gpus(av_synced_dir, high_visual_quality_dir, per_gpu_num_workers)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--total_num_workers", type=int, default=100)
+    parser.add_argument("--per_gpu_num_workers", type=int, default=20)
+    parser.add_argument("--resolution", type=int, default=256)
+    parser.add_argument("--sync_conf_threshold", type=int, default=3)
+    parser.add_argument("--temp_dir", type=str, default="temp")
+    parser.add_argument("--input_dir", type=str, required=True)
+    args = parser.parse_args()
+
+    data_processing_pipeline(
+        args.total_num_workers,
+        args.per_gpu_num_workers,
+        args.resolution,
+        args.sync_conf_threshold,
+        args.temp_dir,
+        args.input_dir,
+    )

preprocess/detect_shot.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import subprocess
+import tqdm
+from multiprocessing import Pool
+
+paths = []
+
+
+def gather_paths(input_dir, output_dir):
+    for video in sorted(os.listdir(input_dir)):
+        if video.endswith(".mp4"):
+            video_input = os.path.join(input_dir, video)
+            video_output = os.path.join(output_dir, video)
+            if os.path.isfile(video_output):
+                continue
+            paths.append([video_input, output_dir])
+        elif os.path.isdir(os.path.join(input_dir, video)):
+            gather_paths(os.path.join(input_dir, video), os.path.join(output_dir, video))
+
+
+def detect_shot(video_input, output_dir):
+    os.makedirs(output_dir, exist_ok=True)
+    video = os.path.basename(video_input)[:-4]
+    command = f"scenedetect --quiet -i {video_input} detect-adaptive --threshold 2 split-video --filename '{video}_shot_$SCENE_NUMBER' --output {output_dir}"
+    # command = f"scenedetect --quiet -i {video_input} detect-adaptive --threshold 2 split-video --high-quality --filename '{video}_shot_$SCENE_NUMBER' --output {output_dir}"
+    subprocess.run(command, shell=True)
+
+
+def multi_run_wrapper(args):
+    return detect_shot(*args)
+
+
+def detect_shot_multiprocessing(input_dir, output_dir, num_workers):
+    print(f"Recursively gathering video paths of {input_dir} ...")
+    gather_paths(input_dir, output_dir)
+
+    print(f"Detecting shot of {input_dir} ...")
+    with Pool(num_workers) as pool:
+        for _ in tqdm.tqdm(pool.imap_unordered(multi_run_wrapper, paths), total=len(paths)):
+            pass
+
+
+if __name__ == "__main__":
+    input_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/ads/high-resolution"
+    output_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/ads/shot"
+    num_workers = 50
+
+    detect_shot_multiprocessing(input_dir, output_dir, num_workers)

preprocess/filter_high_resolution.py

@@ -0,0 +1,112 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import mediapipe as mp
+from latentsync.utils.util import read_video
+import os
+import tqdm
+import shutil
+from multiprocessing import Pool
+
+paths = []
+
+
+def gather_video_paths(input_dir, output_dir, resolution):
+    for video in sorted(os.listdir(input_dir)):
+        if video.endswith(".mp4"):
+            video_input = os.path.join(input_dir, video)
+            video_output = os.path.join(output_dir, video)
+            if os.path.isfile(video_output):
+                continue
+            paths.append([video_input, video_output, resolution])
+        elif os.path.isdir(os.path.join(input_dir, video)):
+            gather_video_paths(os.path.join(input_dir, video), os.path.join(output_dir, video), resolution)
+
+
+class FaceDetector:
+    def __init__(self, resolution=256):
+        self.face_detection = mp.solutions.face_detection.FaceDetection(
+            model_selection=0, min_detection_confidence=0.5
+        )
+        self.resolution = resolution
+
+    def detect_face(self, image):
+        height, width = image.shape[:2]
+        # Process the image and detect faces.
+        results = self.face_detection.process(image)
+
+        if not results.detections:  # Face not detected
+            raise Exception("Face not detected")
+
+        if len(results.detections) != 1:
+            return False
+        detection = results.detections[0]  # Only use the first face in the image
+
+        bounding_box = detection.location_data.relative_bounding_box
+        face_width = int(bounding_box.width * width)
+        face_height = int(bounding_box.height * height)
+        if face_width < self.resolution or face_height < self.resolution:
+            return False
+        return True
+
+    def detect_video(self, video_path):
+        video_frames = read_video(video_path, change_fps=False)
+        if len(video_frames) == 0:
+            return False
+        for frame in video_frames:
+            if not self.detect_face(frame):
+                return False
+        return True
+
+    def close(self):
+        self.face_detection.close()
+
+
+def filter_video(video_input, video_out, resolution):
+    if os.path.isfile(video_out):
+        return
+    face_detector = FaceDetector(resolution)
+    try:
+        save = face_detector.detect_video(video_input)
+    except Exception as e:
+        # print(f"Exception: {e} Input video: {video_input}")
+        face_detector.close()
+        return
+    if save:
+        os.makedirs(os.path.dirname(video_out), exist_ok=True)
+        shutil.copy(video_input, video_out)
+    face_detector.close()
+
+
+def multi_run_wrapper(args):
+    return filter_video(*args)
+
+
+def filter_high_resolution_multiprocessing(input_dir, output_dir, resolution, num_workers):
+    print(f"Recursively gathering video paths of {input_dir} ...")
+    gather_video_paths(input_dir, output_dir, resolution)
+
+    print(f"Filtering high resolution videos in {input_dir} ...")
+    with Pool(num_workers) as pool:
+        for _ in tqdm.tqdm(pool.imap_unordered(multi_run_wrapper, paths), total=len(paths)):
+            pass
+
+
+if __name__ == "__main__":
+    input_dir = "/mnt/bn/maliva-gen-ai/lichunyu/HDTF/original/train"
+    output_dir = "/mnt/bn/maliva-gen-ai/lichunyu/HDTF/detected/train"
+    resolution = 256
+    num_workers = 50
+
+    filter_high_resolution_multiprocessing(input_dir, output_dir, resolution, num_workers)

preprocess/filter_visual_quality.py

@@ -0,0 +1,127 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import tqdm
+import torch
+import torchvision
+import shutil
+from multiprocessing import Process
+import numpy as np
+from decord import VideoReader
+from einops import rearrange
+from eval.hyper_iqa import HyperNet, TargetNet
+
+
+paths = []
+
+
+def gather_paths(input_dir, output_dir):
+    # os.makedirs(output_dir, exist_ok=True)
+
+    for video in tqdm.tqdm(sorted(os.listdir(input_dir))):
+        if video.endswith(".mp4"):
+            video_input = os.path.join(input_dir, video)
+            video_output = os.path.join(output_dir, video)
+            if os.path.isfile(video_output):
+                continue
+            paths.append((video_input, video_output))
+        elif os.path.isdir(os.path.join(input_dir, video)):
+            gather_paths(os.path.join(input_dir, video), os.path.join(output_dir, video))
+
+
+def read_video(video_path: str):
+    vr = VideoReader(video_path)
+    first_frame = vr[0].asnumpy()
+    middle_frame = vr[len(vr) // 2].asnumpy()
+    last_frame = vr[-1].asnumpy()
+    vr.seek(0)
+    video_frames = np.stack([first_frame, middle_frame, last_frame], axis=0)
+    video_frames = torch.from_numpy(rearrange(video_frames, "b h w c -> b c h w"))
+    video_frames = video_frames / 255.0
+    return video_frames
+
+
+def func(paths, device_id):
+    device = f"cuda:{device_id}"
+
+    model_hyper = HyperNet(16, 112, 224, 112, 56, 28, 14, 7).to(device)
+    model_hyper.train(False)
+
+    # load the pre-trained model on the koniq-10k dataset
+    model_hyper.load_state_dict((torch.load("checkpoints/auxiliary/koniq_pretrained.pkl", map_location=device)))
+
+    transforms = torchvision.transforms.Compose(
+        [
+            torchvision.transforms.CenterCrop(size=224),
+            torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+        ]
+    )
+
+    for video_input, video_output in paths:
+        try:
+            video_frames = read_video(video_input)
+            video_frames = transforms(video_frames)
+            video_frames = video_frames.clone().detach().to(device)
+            paras = model_hyper(video_frames)  # 'paras' contains the network weights conveyed to target network
+
+            # Building target network
+            model_target = TargetNet(paras).to(device)
+            for param in model_target.parameters():
+                param.requires_grad = False
+
+            # Quality prediction
+            pred = model_target(paras["target_in_vec"])  # 'paras['target_in_vec']' is the input to target net
+
+            # quality score ranges from 0-100, a higher score indicates a better quality
+            quality_score = pred.mean().item()
+            print(f"Input video: {video_input}\nVisual quality score: {quality_score:.2f}")
+
+            if quality_score >= 40:
+                os.makedirs(os.path.dirname(video_output), exist_ok=True)
+                shutil.copy(video_input, video_output)
+        except Exception as e:
+            print(e)
+
+
+def split(a, n):
+    k, m = divmod(len(a), n)
+    return (a[i * k + min(i, m) : (i + 1) * k + min(i + 1, m)] for i in range(n))
+
+
+def filter_visual_quality_multi_gpus(input_dir, output_dir, num_workers):
+    gather_paths(input_dir, output_dir)
+    num_devices = torch.cuda.device_count()
+    if num_devices == 0:
+        raise RuntimeError("No GPUs found")
+    split_paths = list(split(paths, num_workers * num_devices))
+    processes = []
+
+    for i in range(num_devices):
+        for j in range(num_workers):
+            process_index = i * num_workers + j
+            process = Process(target=func, args=(split_paths[process_index], i))
+            process.start()
+            processes.append(process)
+
+    for process in processes:
+        process.join()
+
+
+if __name__ == "__main__":
+    input_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/VoxCeleb2/av_synced_high"
+    output_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/VoxCeleb2/high_visual_quality"
+    num_workers = 20  # How many processes per device
+
+    filter_visual_quality_multi_gpus(input_dir, output_dir, num_workers)

preprocess/remove_broken_videos.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+from multiprocessing import Pool
+import tqdm
+
+from latentsync.utils.av_reader import AVReader
+from latentsync.utils.util import gather_video_paths_recursively
+
+
+def remove_broken_video(video_path):
+    try:
+        AVReader(video_path)
+    except Exception:
+        os.remove(video_path)
+
+
+def remove_broken_videos_multiprocessing(input_dir, num_workers):
+    video_paths = gather_video_paths_recursively(input_dir)
+
+    print("Removing broken videos...")
+    with Pool(num_workers) as pool:
+        for _ in tqdm.tqdm(pool.imap_unordered(remove_broken_video, video_paths), total=len(video_paths)):
+            pass
+
+
+if __name__ == "__main__":
+    input_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/multilingual/affine_transformed"
+    num_workers = 50
+
+    remove_broken_videos_multiprocessing(input_dir, num_workers)

preprocess/remove_incorrect_affined.py

@@ -0,0 +1,81 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import mediapipe as mp
+from latentsync.utils.util import read_video, gather_video_paths_recursively
+import os
+import tqdm
+from multiprocessing import Pool
+
+
+class FaceDetector:
+    def __init__(self):
+        self.face_detection = mp.solutions.face_detection.FaceDetection(
+            model_selection=0, min_detection_confidence=0.5
+        )
+
+    def detect_face(self, image):
+        # Process the image and detect faces.
+        results = self.face_detection.process(image)
+
+        if not results.detections:  # Face not detected
+            return False
+
+        if len(results.detections) != 1:
+            return False
+        return True
+
+    def detect_video(self, video_path):
+        try:
+            video_frames = read_video(video_path, change_fps=False)
+        except Exception as e:
+            print(f"Exception: {e} - {video_path}")
+            return False
+        if len(video_frames) == 0:
+            return False
+        for frame in video_frames:
+            if not self.detect_face(frame):
+                return False
+        return True
+
+    def close(self):
+        self.face_detection.close()
+
+
+def remove_incorrect_affined(video_path):
+    if not os.path.isfile(video_path):
+        return
+    face_detector = FaceDetector()
+    has_face = face_detector.detect_video(video_path)
+    if not has_face:
+        os.remove(video_path)
+        print(f"Removed: {video_path}")
+    face_detector.close()
+
+
+def remove_incorrect_affined_multiprocessing(input_dir, num_workers):
+    video_paths = gather_video_paths_recursively(input_dir)
+    print(f"Total videos: {len(video_paths)}")
+
+    print(f"Removing incorrect affined videos in {input_dir} ...")
+    with Pool(num_workers) as pool:
+        for _ in tqdm.tqdm(pool.imap_unordered(remove_incorrect_affined, video_paths), total=len(video_paths)):
+            pass
+
+
+if __name__ == "__main__":
+    input_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/multilingual_dcc/high_visual_quality"
+    num_workers = 50
+
+    remove_incorrect_affined_multiprocessing(input_dir, num_workers)

preprocess/resample_fps_hz.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import subprocess
+import tqdm
+from multiprocessing import Pool
+import cv2
+
+paths = []
+
+
+def gather_paths(input_dir, output_dir):
+    for video in sorted(os.listdir(input_dir)):
+        if video.endswith(".mp4"):
+            video_input = os.path.join(input_dir, video)
+            video_output = os.path.join(output_dir, video)
+            if os.path.isfile(video_output):
+                continue
+            paths.append([video_input, video_output])
+        elif os.path.isdir(os.path.join(input_dir, video)):
+            gather_paths(os.path.join(input_dir, video), os.path.join(output_dir, video))
+
+
+def get_video_fps(video_path: str):
+    cam = cv2.VideoCapture(video_path)
+    fps = cam.get(cv2.CAP_PROP_FPS)
+    return fps
+
+
+def resample_fps_hz(video_input, video_output):
+    os.makedirs(os.path.dirname(video_output), exist_ok=True)
+    if get_video_fps(video_input) == 25:
+        command = f"ffmpeg -loglevel error -y -i {video_input} -c:v copy -ar 16000 -q:a 0 {video_output}"
+    else:
+        command = f"ffmpeg -loglevel error -y -i {video_input} -r 25 -ar 16000 -q:a 0 {video_output}"
+    subprocess.run(command, shell=True)
+
+
+def multi_run_wrapper(args):
+    return resample_fps_hz(*args)
+
+
+def resample_fps_hz_multiprocessing(input_dir, output_dir, num_workers):
+    print(f"Recursively gathering video paths of {input_dir} ...")
+    gather_paths(input_dir, output_dir)
+
+    print(f"Resampling FPS and Hz of {input_dir} ...")
+    with Pool(num_workers) as pool:
+        for _ in tqdm.tqdm(pool.imap_unordered(multi_run_wrapper, paths), total=len(paths)):
+            pass
+
+
+if __name__ == "__main__":
+    input_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/HDTF/segmented/train"
+    output_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/HDTF/resampled_test"
+    num_workers = 20
+
+    resample_fps_hz_multiprocessing(input_dir, output_dir, num_workers)

preprocess/segment_videos.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import subprocess
+import tqdm
+from multiprocessing import Pool
+
+paths = []
+
+
+def gather_paths(input_dir, output_dir):
+    for video in sorted(os.listdir(input_dir)):
+        if video.endswith(".mp4"):
+            video_basename = video[:-4]
+            video_input = os.path.join(input_dir, video)
+            video_output = os.path.join(output_dir, f"{video_basename}_%03d.mp4")
+            if os.path.isfile(video_output):
+                continue
+            paths.append([video_input, video_output])
+        elif os.path.isdir(os.path.join(input_dir, video)):
+            gather_paths(os.path.join(input_dir, video), os.path.join(output_dir, video))
+
+
+def segment_video(video_input, video_output):
+    os.makedirs(os.path.dirname(video_output), exist_ok=True)
+    command = f"ffmpeg -loglevel error -y -i {video_input} -map 0 -c:v copy -segment_time 5 -f segment -reset_timestamps 1 -q:a 0 {video_output}"
+    # command = f'ffmpeg -loglevel error -y -i {video_input} -map 0 -segment_time 5 -f segment -reset_timestamps 1 -force_key_frames "expr:gte(t,n_forced*5)" -crf 18 -q:a 0 {video_output}'
+    subprocess.run(command, shell=True)
+
+
+def multi_run_wrapper(args):
+    return segment_video(*args)
+
+
+def segment_videos_multiprocessing(input_dir, output_dir, num_workers):
+    print(f"Recursively gathering video paths of {input_dir} ...")
+    gather_paths(input_dir, output_dir)
+
+    print(f"Segmenting videos of {input_dir} ...")
+    with Pool(num_workers) as pool:
+        for _ in tqdm.tqdm(pool.imap_unordered(multi_run_wrapper, paths), total=len(paths)):
+            pass
+
+
+if __name__ == "__main__":
+    input_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/avatars_new/cut"
+    output_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/avatars_new/segmented"
+    num_workers = 50
+
+    segment_videos_multiprocessing(input_dir, output_dir, num_workers)

preprocess/sync_av.py

@@ -0,0 +1,113 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import tqdm
+from eval.syncnet import SyncNetEval
+from eval.syncnet_detect import SyncNetDetector
+from eval.eval_sync_conf import syncnet_eval
+import torch
+import subprocess
+import shutil
+from multiprocessing import Process
+
+paths = []
+
+
+def gather_paths(input_dir, output_dir):
+    # os.makedirs(output_dir, exist_ok=True)
+
+    for video in tqdm.tqdm(sorted(os.listdir(input_dir))):
+        if video.endswith(".mp4"):
+            video_input = os.path.join(input_dir, video)
+            video_output = os.path.join(output_dir, video)
+            if os.path.isfile(video_output):
+                continue
+            paths.append((video_input, video_output))
+        elif os.path.isdir(os.path.join(input_dir, video)):
+            gather_paths(os.path.join(input_dir, video), os.path.join(output_dir, video))
+
+
+def adjust_offset(video_input: str, video_output: str, av_offset: int, fps: int = 25):
+    command = f"ffmpeg -loglevel error -y -i {video_input} -itsoffset {av_offset/fps} -i {video_input} -map 0:v -map 1:a -c copy -q:v 0 -q:a 0 {video_output}"
+    subprocess.run(command, shell=True)
+
+
+def func(sync_conf_threshold, paths, device_id, process_temp_dir):
+    os.makedirs(process_temp_dir, exist_ok=True)
+    device = f"cuda:{device_id}"
+
+    syncnet = SyncNetEval(device=device)
+    syncnet.loadParameters("checkpoints/auxiliary/syncnet_v2.model")
+
+    detect_results_dir = os.path.join(process_temp_dir, "detect_results")
+    syncnet_eval_results_dir = os.path.join(process_temp_dir, "syncnet_eval_results")
+
+    syncnet_detector = SyncNetDetector(device=device, detect_results_dir=detect_results_dir)
+
+    for video_input, video_output in paths:
+        try:
+            av_offset, conf = syncnet_eval(
+                syncnet, syncnet_detector, video_input, syncnet_eval_results_dir, detect_results_dir
+            )
+            if conf >= sync_conf_threshold and abs(av_offset) <= 6:
+                os.makedirs(os.path.dirname(video_output), exist_ok=True)
+                if av_offset == 0:
+                    shutil.copy(video_input, video_output)
+                else:
+                    adjust_offset(video_input, video_output, av_offset)
+        except Exception as e:
+            print(e)
+
+
+def split(a, n):
+    k, m = divmod(len(a), n)
+    return (a[i * k + min(i, m) : (i + 1) * k + min(i + 1, m)] for i in range(n))
+
+
+def sync_av_multi_gpus(input_dir, output_dir, temp_dir, num_workers, sync_conf_threshold):
+    gather_paths(input_dir, output_dir)
+    num_devices = torch.cuda.device_count()
+    if num_devices == 0:
+        raise RuntimeError("No GPUs found")
+    split_paths = list(split(paths, num_workers * num_devices))
+    processes = []
+
+    for i in range(num_devices):
+        for j in range(num_workers):
+            process_index = i * num_workers + j
+            process = Process(
+                target=func,
+                args=(
+                    sync_conf_threshold,
+                    split_paths[process_index],
+                    i,
+                    os.path.join(temp_dir, f"process_{process_index}"),
+                ),
+            )
+            process.start()
+            processes.append(process)
+
+    for process in processes:
+        process.join()
+
+
+if __name__ == "__main__":
+    input_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/ads/affine_transformed"
+    output_dir = "/mnt/bn/maliva-gen-ai-v2/chunyu.li/VoxCeleb2/temp"
+    temp_dir = "temp"
+    num_workers = 20  # How many processes per device
+    sync_conf_threshold = 3
+
+    sync_av_multi_gpus(input_dir, output_dir, temp_dir, num_workers, sync_conf_threshold)

requirements.txt

@@ -0,0 +1,30 @@
+torch==2.2.2
+torchvision==0.17.2
+--extra-index-url https://download.pytorch.org/whl/cu121
+xformers==0.0.26
+triton==2.2.0
+
+diffusers==0.11.1
+transformers==4.38.0
+huggingface-hub==0.25.2
+imageio==2.27.0
+decord==0.6.0
+accelerate==0.26.1
+einops==0.7.0
+omegaconf==2.3.0
+safetensors==0.4.2
+opencv-python==4.9.0.80
+mediapipe==0.10.11
+av==11.0.0
+torch-fidelity==0.3.0
+torchmetrics==1.3.1
+python_speech_features==0.6
+librosa==0.10.1
+scenedetect==0.6.1
+ffmpeg-python==0.2.0
+lpips==0.1.4
+face-alignment==1.4.1
+ninja==1.11.1.1
+pandas==2.0.3
+numpy==1.24.4
+gradio==5.9.1

scripts/inference.py

@@ -0,0 +1,103 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+from omegaconf import OmegaConf
+import torch
+from diffusers import AutoencoderKL, DDIMScheduler
+from latentsync.models.unet import UNet3DConditionModel
+from latentsync.pipelines.lipsync_pipeline import LipsyncPipeline
+from diffusers.utils.import_utils import is_xformers_available
+from accelerate.utils import set_seed
+from latentsync.whisper.audio2feature import Audio2Feature
+
+
+def main(config, args):
+    # Check if the GPU supports float16
+    is_fp16_supported = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] > 7
+    dtype = torch.float16 if is_fp16_supported else torch.float32
+
+    print(f"Input video path: {args.video_path}")
+    print(f"Input audio path: {args.audio_path}")
+    print(f"Loaded checkpoint path: {args.inference_ckpt_path}")
+
+    scheduler = DDIMScheduler.from_pretrained("configs")
+
+    if config.model.cross_attention_dim == 768:
+        whisper_model_path = "checkpoints/whisper/small.pt"
+    elif config.model.cross_attention_dim == 384:
+        whisper_model_path = "checkpoints/whisper/tiny.pt"
+    else:
+        raise NotImplementedError("cross_attention_dim must be 768 or 384")
+
+    audio_encoder = Audio2Feature(model_path=whisper_model_path, device="cuda", num_frames=config.data.num_frames)
+
+    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=dtype)
+    vae.config.scaling_factor = 0.18215
+    vae.config.shift_factor = 0
+
+    unet, _ = UNet3DConditionModel.from_pretrained(
+        OmegaConf.to_container(config.model),
+        args.inference_ckpt_path,  # load checkpoint
+        device="cpu",
+    )
+
+    unet = unet.to(dtype=dtype)
+
+    # set xformers
+    if is_xformers_available():
+        unet.enable_xformers_memory_efficient_attention()
+
+    pipeline = LipsyncPipeline(
+        vae=vae,
+        audio_encoder=audio_encoder,
+        unet=unet,
+        scheduler=scheduler,
+    ).to("cuda")
+
+    if args.seed != -1:
+        set_seed(args.seed)
+    else:
+        torch.seed()
+
+    print(f"Initial seed: {torch.initial_seed()}")
+
+    pipeline(
+        video_path=args.video_path,
+        audio_path=args.audio_path,
+        video_out_path=args.video_out_path,
+        video_mask_path=args.video_out_path.replace(".mp4", "_mask.mp4"),
+        num_frames=config.data.num_frames,
+        num_inference_steps=config.run.inference_steps,
+        guidance_scale=args.guidance_scale,
+        weight_dtype=dtype,
+        width=config.data.resolution,
+        height=config.data.resolution,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--unet_config_path", type=str, default="configs/unet.yaml")
+    parser.add_argument("--inference_ckpt_path", type=str, required=True)
+    parser.add_argument("--video_path", type=str, required=True)
+    parser.add_argument("--audio_path", type=str, required=True)
+    parser.add_argument("--video_out_path", type=str, required=True)
+    parser.add_argument("--guidance_scale", type=float, default=1.0)
+    parser.add_argument("--seed", type=int, default=1247)
+    args = parser.parse_args()
+
+    config = OmegaConf.load(args.unet_config_path)
+
+    main(config, args)

scripts/train_syncnet.py

@@ -0,0 +1,336 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from tqdm.auto import tqdm
+import os, argparse, datetime, math
+import logging
+from omegaconf import OmegaConf
+import shutil
+
+from latentsync.data.syncnet_dataset import SyncNetDataset
+from latentsync.models.syncnet import SyncNet
+from latentsync.models.syncnet_wav2lip import SyncNetWav2Lip
+from latentsync.utils.util import gather_loss, plot_loss_chart
+from accelerate.utils import set_seed
+
+import torch
+from diffusers import AutoencoderKL
+from diffusers.utils.logging import get_logger
+from einops import rearrange
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.utils.data.distributed import DistributedSampler
+from latentsync.utils.util import init_dist, cosine_loss
+
+logger = get_logger(__name__)
+
+
+def main(config):
+    # Initialize distributed training
+    local_rank = init_dist()
+    global_rank = dist.get_rank()
+    num_processes = dist.get_world_size()
+    is_main_process = global_rank == 0
+
+    seed = config.run.seed + global_rank
+    set_seed(seed)
+
+    # Logging folder
+    folder_name = "train" + datetime.datetime.now().strftime(f"-%Y_%m_%d-%H:%M:%S")
+    output_dir = os.path.join(config.data.train_output_dir, folder_name)
+
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+
+    # Handle the output folder creation
+    if is_main_process:
+        os.makedirs(output_dir, exist_ok=True)
+        os.makedirs(f"{output_dir}/checkpoints", exist_ok=True)
+        os.makedirs(f"{output_dir}/loss_charts", exist_ok=True)
+        shutil.copy(config.config_path, output_dir)
+
+    device = torch.device(local_rank)
+
+    if config.data.latent_space:
+        vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)
+        vae.requires_grad_(False)
+        vae.to(device)
+    else:
+        vae = None
+
+    # Dataset and Dataloader setup
+    train_dataset = SyncNetDataset(config.data.train_data_dir, config.data.train_fileslist, config)
+    val_dataset = SyncNetDataset(config.data.val_data_dir, config.data.val_fileslist, config)
+
+    train_distributed_sampler = DistributedSampler(
+        train_dataset,
+        num_replicas=num_processes,
+        rank=global_rank,
+        shuffle=True,
+        seed=config.run.seed,
+    )
+
+    # DataLoaders creation:
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset,
+        batch_size=config.data.batch_size,
+        shuffle=False,
+        sampler=train_distributed_sampler,
+        num_workers=config.data.num_workers,
+        pin_memory=False,
+        drop_last=True,
+        worker_init_fn=train_dataset.worker_init_fn,
+    )
+    
+    num_samples_limit = 640
+
+    val_batch_size = min(
+        num_samples_limit // config.data.num_frames, config.data.batch_size
+    )  # limit batch size to avoid CUDA OOM
+
+    val_dataloader = torch.utils.data.DataLoader(
+        val_dataset,
+        batch_size=val_batch_size,
+        shuffle=False,
+        num_workers=config.data.num_workers,
+        pin_memory=False,
+        drop_last=False,
+        worker_init_fn=val_dataset.worker_init_fn,
+    )
+
+    # Model
+    syncnet = SyncNet(OmegaConf.to_container(config.model)).to(device)
+    # syncnet = SyncNetWav2Lip().to(device)
+
+    optimizer = torch.optim.AdamW(
+        list(filter(lambda p: p.requires_grad, syncnet.parameters())), lr=config.optimizer.lr
+    )
+
+    if config.ckpt.resume_ckpt_path != "":
+        if is_main_process:
+            logger.info(f"Load checkpoint from: {config.ckpt.resume_ckpt_path}")
+        ckpt = torch.load(config.ckpt.resume_ckpt_path, map_location=device)
+
+        syncnet.load_state_dict(ckpt["state_dict"])
+        global_step = ckpt["global_step"]
+        train_step_list = ckpt["train_step_list"]
+        train_loss_list = ckpt["train_loss_list"]
+        val_step_list = ckpt["val_step_list"]
+        val_loss_list = ckpt["val_loss_list"]
+    else:
+        global_step = 0
+        train_step_list = []
+        train_loss_list = []
+        val_step_list = []
+        val_loss_list = []
+
+    # DDP wrapper
+    syncnet = DDP(syncnet, device_ids=[local_rank], output_device=local_rank)
+
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader))
+    num_train_epochs = math.ceil(config.run.max_train_steps / num_update_steps_per_epoch)
+    # validation_steps = int(config.ckpt.save_ckpt_steps // 5)
+    # validation_steps = 100
+
+    if is_main_process:
+        logger.info("***** Running training *****")
+        logger.info(f"  Num examples = {len(train_dataset)}")
+        logger.info(f"  Num Epochs = {num_train_epochs}")
+        logger.info(f"  Instantaneous batch size per device = {config.data.batch_size}")
+        logger.info(f"  Total train batch size (w. parallel & distributed) = {config.data.batch_size * num_processes}")
+        logger.info(f"  Total optimization steps = {config.run.max_train_steps}")
+
+    first_epoch = global_step // num_update_steps_per_epoch
+    num_val_batches = config.data.num_val_samples // (num_processes * config.data.batch_size)
+
+    # Only show the progress bar once on each machine.
+    progress_bar = tqdm(
+        range(0, config.run.max_train_steps), initial=global_step, desc="Steps", disable=not is_main_process
+    )
+
+    # Support mixed-precision training
+    scaler = torch.cuda.amp.GradScaler() if config.run.mixed_precision_training else None
+
+    for epoch in range(first_epoch, num_train_epochs):
+        train_dataloader.sampler.set_epoch(epoch)
+        syncnet.train()
+
+        for step, batch in enumerate(train_dataloader):
+            ### >>>> Training >>>> ###
+
+            frames = batch["frames"].to(device, dtype=torch.float16)
+            audio_samples = batch["audio_samples"].to(device, dtype=torch.float16)
+            y = batch["y"].to(device, dtype=torch.float32)
+
+            if config.data.latent_space:
+                max_batch_size = (
+                    num_samples_limit // config.data.num_frames
+                )  # due to the limited cuda memory, we split the input frames into parts
+                if frames.shape[0] > max_batch_size:
+                    assert (
+                        frames.shape[0] % max_batch_size == 0
+                    ), f"max_batch_size {max_batch_size} should be divisible by batch_size {frames.shape[0]}"
+                    frames_part_results = []
+                    for i in range(0, frames.shape[0], max_batch_size):
+                        frames_part = frames[i : i + max_batch_size]
+                        frames_part = rearrange(frames_part, "b f c h w -> (b f) c h w")
+                        with torch.no_grad():
+                            frames_part = vae.encode(frames_part).latent_dist.sample() * 0.18215
+                        frames_part_results.append(frames_part)
+                    frames = torch.cat(frames_part_results, dim=0)
+                else:
+                    frames = rearrange(frames, "b f c h w -> (b f) c h w")
+                    with torch.no_grad():
+                        frames = vae.encode(frames).latent_dist.sample() * 0.18215
+
+                frames = rearrange(frames, "(b f) c h w -> b (f c) h w", f=config.data.num_frames)
+            else:
+                frames = rearrange(frames, "b f c h w -> b (f c) h w")
+
+            if config.data.lower_half:
+                height = frames.shape[2]
+                frames = frames[:, :, height // 2 :, :]
+
+            # audio_embeds = wav2vec_encoder(audio_samples).last_hidden_state
+
+            # Mixed-precision training
+            with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=config.run.mixed_precision_training):
+                vision_embeds, audio_embeds = syncnet(frames, audio_samples)
+
+            loss = cosine_loss(vision_embeds.float(), audio_embeds.float(), y).mean()
+
+            optimizer.zero_grad()
+
+            # Backpropagate
+            if config.run.mixed_precision_training:
+                scaler.scale(loss).backward()
+                """ >>> gradient clipping >>> """
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(syncnet.parameters(), config.optimizer.max_grad_norm)
+                """ <<< gradient clipping <<< """
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                loss.backward()
+                """ >>> gradient clipping >>> """
+                torch.nn.utils.clip_grad_norm_(syncnet.parameters(), config.optimizer.max_grad_norm)
+                """ <<< gradient clipping <<< """
+                optimizer.step()
+
+            progress_bar.update(1)
+            global_step += 1
+
+            global_average_loss = gather_loss(loss, device)
+            train_step_list.append(global_step)
+            train_loss_list.append(global_average_loss)
+
+            if is_main_process and global_step % config.run.validation_steps == 0:
+                logger.info(f"Validation at step {global_step}")
+                val_loss = validation(
+                    val_dataloader,
+                    device,
+                    syncnet,
+                    cosine_loss,
+                    config.data.latent_space,
+                    config.data.lower_half,
+                    vae,
+                    num_val_batches,
+                )
+                val_step_list.append(global_step)
+                val_loss_list.append(val_loss)
+                logger.info(f"Validation loss at step {global_step} is {val_loss:0.3f}")
+
+            if is_main_process and global_step % config.ckpt.save_ckpt_steps == 0:
+                checkpoint_save_path = os.path.join(output_dir, f"checkpoints/checkpoint-{global_step}.pt")
+                torch.save(
+                    {
+                        "state_dict": syncnet.module.state_dict(),  # to unwrap DDP
+                        "global_step": global_step,
+                        "train_step_list": train_step_list,
+                        "train_loss_list": train_loss_list,
+                        "val_step_list": val_step_list,
+                        "val_loss_list": val_loss_list,
+                    },
+                    checkpoint_save_path,
+                )
+                logger.info(f"Saved checkpoint to {checkpoint_save_path}")
+                plot_loss_chart(
+                    os.path.join(output_dir, f"loss_charts/loss_chart-{global_step}.png"),
+                    ("Train loss", train_step_list, train_loss_list),
+                    ("Val loss", val_step_list, val_loss_list),
+                )
+
+            progress_bar.set_postfix({"step_loss": global_average_loss})
+            if global_step >= config.run.max_train_steps:
+                break
+
+    progress_bar.close()
+    dist.destroy_process_group()
+
+
+@torch.no_grad()
+def validation(val_dataloader, device, syncnet, cosine_loss, latent_space, lower_half, vae, num_val_batches):
+    syncnet.eval()
+
+    losses = []
+    val_step = 0
+    while True:
+        for step, batch in enumerate(val_dataloader):
+            ### >>>> Validation >>>> ###
+
+            frames = batch["frames"].to(device, dtype=torch.float16)
+            audio_samples = batch["audio_samples"].to(device, dtype=torch.float16)
+            y = batch["y"].to(device, dtype=torch.float32)
+
+            if latent_space:
+                num_frames = frames.shape[1]
+                frames = rearrange(frames, "b f c h w -> (b f) c h w")
+                frames = vae.encode(frames).latent_dist.sample() * 0.18215
+                frames = rearrange(frames, "(b f) c h w -> b (f c) h w", f=num_frames)
+            else:
+                frames = rearrange(frames, "b f c h w -> b (f c) h w")
+
+            if lower_half:
+                height = frames.shape[2]
+                frames = frames[:, :, height // 2 :, :]
+
+            with torch.autocast(device_type="cuda", dtype=torch.float16):
+                vision_embeds, audio_embeds = syncnet(frames, audio_samples)
+
+            loss = cosine_loss(vision_embeds.float(), audio_embeds.float(), y).mean()
+
+            losses.append(loss.item())
+
+            val_step += 1
+            if val_step > num_val_batches:
+                syncnet.train()
+                if len(losses) == 0:
+                    raise RuntimeError("No validation data")
+                return sum(losses) / len(losses)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Code to train the expert lip-sync discriminator")
+    parser.add_argument("--config_path", type=str, default="configs/syncnet/syncnet_16_vae.yaml")
+    args = parser.parse_args()
+
+    # Load a configuration file
+    config = OmegaConf.load(args.config_path)
+    config.config_path = args.config_path
+
+    main(config)

scripts/train_unet.py

@@ -0,0 +1,510 @@
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import math
+import argparse
+import shutil
+import datetime
+import logging
+from omegaconf import OmegaConf
+
+from tqdm.auto import tqdm
+from einops import rearrange
+
+import torch
+import torch.nn.functional as F
+import torch.distributed as dist
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+import diffusers
+from diffusers import AutoencoderKL, DDIMScheduler
+from diffusers.utils.logging import get_logger
+from diffusers.optimization import get_scheduler
+from diffusers.utils.import_utils import is_xformers_available
+from accelerate.utils import set_seed
+
+from latentsync.data.unet_dataset import UNetDataset
+from latentsync.models.unet import UNet3DConditionModel
+from latentsync.models.syncnet import SyncNet
+from latentsync.pipelines.lipsync_pipeline import LipsyncPipeline
+from latentsync.utils.util import (
+    init_dist,
+    cosine_loss,
+    reversed_forward,
+)
+from latentsync.utils.util import plot_loss_chart, gather_loss
+from latentsync.whisper.audio2feature import Audio2Feature
+from latentsync.trepa import TREPALoss
+from eval.syncnet import SyncNetEval
+from eval.syncnet_detect import SyncNetDetector
+from eval.eval_sync_conf import syncnet_eval
+import lpips
+
+
+logger = get_logger(__name__)
+
+
+def main(config):
+    # Initialize distributed training
+    local_rank = init_dist()
+    global_rank = dist.get_rank()
+    num_processes = dist.get_world_size()
+    is_main_process = global_rank == 0
+
+    seed = config.run.seed + global_rank
+    set_seed(seed)
+
+    # Logging folder
+    folder_name = "train" + datetime.datetime.now().strftime(f"-%Y_%m_%d-%H:%M:%S")
+    output_dir = os.path.join(config.data.train_output_dir, folder_name)
+
+    # Make one log on every process with the configuration for debugging.
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO,
+    )
+
+    # Handle the output folder creation
+    if is_main_process:
+        diffusers.utils.logging.set_verbosity_info()
+        os.makedirs(output_dir, exist_ok=True)
+        os.makedirs(f"{output_dir}/checkpoints", exist_ok=True)
+        os.makedirs(f"{output_dir}/val_videos", exist_ok=True)
+        os.makedirs(f"{output_dir}/loss_charts", exist_ok=True)
+        shutil.copy(config.unet_config_path, output_dir)
+        shutil.copy(config.data.syncnet_config_path, output_dir)
+
+    device = torch.device(local_rank)
+
+    noise_scheduler = DDIMScheduler.from_pretrained("configs")
+
+    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)
+    vae.config.scaling_factor = 0.18215
+    vae.config.shift_factor = 0
+    vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
+    vae.requires_grad_(False)
+    vae.to(device)
+
+    syncnet_eval_model = SyncNetEval(device=device)
+    syncnet_eval_model.loadParameters("checkpoints/auxiliary/syncnet_v2.model")
+
+    syncnet_detector = SyncNetDetector(device=device, detect_results_dir="detect_results")
+
+    if config.model.cross_attention_dim == 768:
+        whisper_model_path = "checkpoints/whisper/small.pt"
+    elif config.model.cross_attention_dim == 384:
+        whisper_model_path = "checkpoints/whisper/tiny.pt"
+    else:
+        raise NotImplementedError("cross_attention_dim must be 768 or 384")
+
+    audio_encoder = Audio2Feature(
+        model_path=whisper_model_path,
+        device=device,
+        audio_embeds_cache_dir=config.data.audio_embeds_cache_dir,
+        num_frames=config.data.num_frames,
+    )
+
+    unet, resume_global_step = UNet3DConditionModel.from_pretrained(
+        OmegaConf.to_container(config.model),
+        config.ckpt.resume_ckpt_path,  # load checkpoint
+        device=device,
+    )
+
+    if config.model.add_audio_layer and config.run.use_syncnet:
+        syncnet_config = OmegaConf.load(config.data.syncnet_config_path)
+        if syncnet_config.ckpt.inference_ckpt_path == "":
+            raise ValueError("SyncNet path is not provided")
+        syncnet = SyncNet(OmegaConf.to_container(syncnet_config.model)).to(device=device, dtype=torch.float16)
+        syncnet_checkpoint = torch.load(syncnet_config.ckpt.inference_ckpt_path, map_location=device)
+        syncnet.load_state_dict(syncnet_checkpoint["state_dict"])
+        syncnet.requires_grad_(False)
+
+    unet.requires_grad_(True)
+    trainable_params = list(unet.parameters())
+
+    if config.optimizer.scale_lr:
+        config.optimizer.lr = config.optimizer.lr * num_processes
+
+    optimizer = torch.optim.AdamW(trainable_params, lr=config.optimizer.lr)
+
+    if is_main_process:
+        logger.info(f"trainable params number: {len(trainable_params)}")
+        logger.info(f"trainable params scale: {sum(p.numel() for p in trainable_params) / 1e6:.3f} M")
+
+    # Enable xformers
+    if config.run.enable_xformers_memory_efficient_attention:
+        if is_xformers_available():
+            unet.enable_xformers_memory_efficient_attention()
+        else:
+            raise ValueError("xformers is not available. Make sure it is installed correctly")
+
+    # Enable gradient checkpointing
+    if config.run.enable_gradient_checkpointing:
+        unet.enable_gradient_checkpointing()
+
+    # Get the training dataset
+    train_dataset = UNetDataset(config.data.train_data_dir, config)
+    distributed_sampler = DistributedSampler(
+        train_dataset,
+        num_replicas=num_processes,
+        rank=global_rank,
+        shuffle=True,
+        seed=config.run.seed,
+    )
+
+    # DataLoaders creation:
+    train_dataloader = torch.utils.data.DataLoader(
+        train_dataset,
+        batch_size=config.data.batch_size,
+        shuffle=False,
+        sampler=distributed_sampler,
+        num_workers=config.data.num_workers,
+        pin_memory=False,
+        drop_last=True,
+        worker_init_fn=train_dataset.worker_init_fn,
+    )
+
+    # Get the training iteration
+    if config.run.max_train_steps == -1:
+        assert config.run.max_train_epochs != -1
+        config.run.max_train_steps = config.run.max_train_epochs * len(train_dataloader)
+
+    # Scheduler
+    lr_scheduler = get_scheduler(
+        config.optimizer.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=config.optimizer.lr_warmup_steps,
+        num_training_steps=config.run.max_train_steps,
+    )
+
+    if config.run.perceptual_loss_weight != 0 and config.run.pixel_space_supervise:
+        lpips_loss_func = lpips.LPIPS(net="vgg").to(device)
+
+    if config.run.trepa_loss_weight != 0 and config.run.pixel_space_supervise:
+        trepa_loss_func = TREPALoss(device=device)
+
+    # Validation pipeline
+    pipeline = LipsyncPipeline(
+        vae=vae,
+        audio_encoder=audio_encoder,
+        unet=unet,
+        scheduler=noise_scheduler,
+    ).to(device)
+    pipeline.set_progress_bar_config(disable=True)
+
+    # DDP warpper
+    unet = DDP(unet, device_ids=[local_rank], output_device=local_rank)
+
+    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
+    num_update_steps_per_epoch = math.ceil(len(train_dataloader))
+    # Afterwards we recalculate our number of training epochs
+    num_train_epochs = math.ceil(config.run.max_train_steps / num_update_steps_per_epoch)
+
+    # Train!
+    total_batch_size = config.data.batch_size * num_processes
+
+    if is_main_process:
+        logger.info("***** Running training *****")
+        logger.info(f"  Num examples = {len(train_dataset)}")
+        logger.info(f"  Num Epochs = {num_train_epochs}")
+        logger.info(f"  Instantaneous batch size per device = {config.data.batch_size}")
+        logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
+        logger.info(f"  Total optimization steps = {config.run.max_train_steps}")
+    global_step = resume_global_step
+    first_epoch = resume_global_step // num_update_steps_per_epoch
+
+    # Only show the progress bar once on each machine.
+    progress_bar = tqdm(
+        range(0, config.run.max_train_steps),
+        initial=resume_global_step,
+        desc="Steps",
+        disable=not is_main_process,
+    )
+
+    train_step_list = []
+    sync_loss_list = []
+    recon_loss_list = []
+
+    val_step_list = []
+    sync_conf_list = []
+
+    # Support mixed-precision training
+    scaler = torch.cuda.amp.GradScaler() if config.run.mixed_precision_training else None
+
+    for epoch in range(first_epoch, num_train_epochs):
+        train_dataloader.sampler.set_epoch(epoch)
+        unet.train()
+
+        for step, batch in enumerate(train_dataloader):
+            ### >>>> Training >>>> ###
+
+            if config.model.add_audio_layer:
+                if batch["mel"] != []:
+                    mel = batch["mel"].to(device, dtype=torch.float16)
+
+                audio_embeds_list = []
+                try:
+                    for idx in range(len(batch["video_path"])):
+                        video_path = batch["video_path"][idx]
+                        start_idx = batch["start_idx"][idx]
+
+                        with torch.no_grad():
+                            audio_feat = audio_encoder.audio2feat(video_path)
+                        audio_embeds = audio_encoder.crop_overlap_audio_window(audio_feat, start_idx)
+                        audio_embeds_list.append(audio_embeds)
+                except Exception as e:
+                    logger.info(f"{type(e).__name__} - {e} - {video_path}")
+                    continue
+                audio_embeds = torch.stack(audio_embeds_list)  # (B, 16, 50, 384)
+                audio_embeds = audio_embeds.to(device, dtype=torch.float16)
+            else:
+                audio_embeds = None
+
+            # Convert videos to latent space
+            gt_images = batch["gt"].to(device, dtype=torch.float16)
+            gt_masked_images = batch["masked_gt"].to(device, dtype=torch.float16)
+            mask = batch["mask"].to(device, dtype=torch.float16)
+            ref_images = batch["ref"].to(device, dtype=torch.float16)
+
+            gt_images = rearrange(gt_images, "b f c h w -> (b f) c h w")
+            gt_masked_images = rearrange(gt_masked_images, "b f c h w -> (b f) c h w")
+            mask = rearrange(mask, "b f c h w -> (b f) c h w")
+            ref_images = rearrange(ref_images, "b f c h w -> (b f) c h w")
+
+            with torch.no_grad():
+                gt_latents = vae.encode(gt_images).latent_dist.sample()
+                gt_masked_images = vae.encode(gt_masked_images).latent_dist.sample()
+                ref_images = vae.encode(ref_images).latent_dist.sample()
+
+            mask = torch.nn.functional.interpolate(mask, size=config.data.resolution // vae_scale_factor)
+
+            gt_latents = (
+                rearrange(gt_latents, "(b f) c h w -> b c f h w", f=config.data.num_frames) - vae.config.shift_factor
+            ) * vae.config.scaling_factor
+            gt_masked_images = (
+                rearrange(gt_masked_images, "(b f) c h w -> b c f h w", f=config.data.num_frames)
+                - vae.config.shift_factor
+            ) * vae.config.scaling_factor
+            ref_images = (
+                rearrange(ref_images, "(b f) c h w -> b c f h w", f=config.data.num_frames) - vae.config.shift_factor
+            ) * vae.config.scaling_factor
+            mask = rearrange(mask, "(b f) c h w -> b c f h w", f=config.data.num_frames)
+
+            # Sample noise that we'll add to the latents
+            if config.run.use_mixed_noise:
+                # Refer to the paper: https://arxiv.org/abs/2305.10474
+                noise_shared_std_dev = (config.run.mixed_noise_alpha**2 / (1 + config.run.mixed_noise_alpha**2)) ** 0.5
+                noise_shared = torch.randn_like(gt_latents) * noise_shared_std_dev
+                noise_shared = noise_shared[:, :, 0:1].repeat(1, 1, config.data.num_frames, 1, 1)
+
+                noise_ind_std_dev = (1 / (1 + config.run.mixed_noise_alpha**2)) ** 0.5
+                noise_ind = torch.randn_like(gt_latents) * noise_ind_std_dev
+                noise = noise_ind + noise_shared
+            else:
+                noise = torch.randn_like(gt_latents)
+                noise = noise[:, :, 0:1].repeat(
+                    1, 1, config.data.num_frames, 1, 1
+                )  # Using the same noise for all frames, refer to the paper: https://arxiv.org/abs/2308.09716
+
+            bsz = gt_latents.shape[0]
+
+            # Sample a random timestep for each video
+            timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=gt_latents.device)
+            timesteps = timesteps.long()
+
+            # Add noise to the latents according to the noise magnitude at each timestep
+            # (this is the forward diffusion process)
+            noisy_tensor = noise_scheduler.add_noise(gt_latents, noise, timesteps)
+
+            # Get the target for loss depending on the prediction type
+            if noise_scheduler.config.prediction_type == "epsilon":
+                target = noise
+            elif noise_scheduler.config.prediction_type == "v_prediction":
+                raise NotImplementedError
+            else:
+                raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
+
+            unet_input = torch.cat([noisy_tensor, mask, gt_masked_images, ref_images], dim=1)
+
+            # Predict the noise and compute loss
+            # Mixed-precision training
+            with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=config.run.mixed_precision_training):
+                pred_noise = unet(unet_input, timesteps, encoder_hidden_states=audio_embeds).sample
+
+            if config.run.recon_loss_weight != 0:
+                recon_loss = F.mse_loss(pred_noise.float(), target.float(), reduction="mean")
+            else:
+                recon_loss = 0
+
+            pred_latents = reversed_forward(noise_scheduler, pred_noise, timesteps, noisy_tensor)
+
+            if config.run.pixel_space_supervise:
+                pred_images = vae.decode(
+                    rearrange(pred_latents, "b c f h w -> (b f) c h w") / vae.config.scaling_factor
+                    + vae.config.shift_factor
+                ).sample
+
+            if config.run.perceptual_loss_weight != 0 and config.run.pixel_space_supervise:
+                pred_images_perceptual = pred_images[:, :, pred_images.shape[2] // 2 :, :]
+                gt_images_perceptual = gt_images[:, :, gt_images.shape[2] // 2 :, :]
+                lpips_loss = lpips_loss_func(pred_images_perceptual.float(), gt_images_perceptual.float()).mean()
+            else:
+                lpips_loss = 0
+
+            if config.run.trepa_loss_weight != 0 and config.run.pixel_space_supervise:
+                trepa_pred_images = rearrange(pred_images, "(b f) c h w -> b c f h w", f=config.data.num_frames)
+                trepa_gt_images = rearrange(gt_images, "(b f) c h w -> b c f h w", f=config.data.num_frames)
+                trepa_loss = trepa_loss_func(trepa_pred_images, trepa_gt_images)
+            else:
+                trepa_loss = 0
+
+            if config.model.add_audio_layer and config.run.use_syncnet:
+                if config.run.pixel_space_supervise:
+                    syncnet_input = rearrange(pred_images, "(b f) c h w -> b (f c) h w", f=config.data.num_frames)
+                else:
+                    syncnet_input = rearrange(pred_latents, "b c f h w -> b (f c) h w")
+
+                if syncnet_config.data.lower_half:
+                    height = syncnet_input.shape[2]
+                    syncnet_input = syncnet_input[:, :, height // 2 :, :]
+                ones_tensor = torch.ones((config.data.batch_size, 1)).float().to(device=device)
+                vision_embeds, audio_embeds = syncnet(syncnet_input, mel)
+                sync_loss = cosine_loss(vision_embeds.float(), audio_embeds.float(), ones_tensor).mean()
+                sync_loss_list.append(gather_loss(sync_loss, device))
+            else:
+                sync_loss = 0
+
+            loss = (
+                recon_loss * config.run.recon_loss_weight
+                + sync_loss * config.run.sync_loss_weight
+                + lpips_loss * config.run.perceptual_loss_weight
+                + trepa_loss * config.run.trepa_loss_weight
+            )
+
+            train_step_list.append(global_step)
+            if config.run.recon_loss_weight != 0:
+                recon_loss_list.append(gather_loss(recon_loss, device))
+
+            optimizer.zero_grad()
+
+            # Backpropagate
+            if config.run.mixed_precision_training:
+                scaler.scale(loss).backward()
+                """ >>> gradient clipping >>> """
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(unet.parameters(), config.optimizer.max_grad_norm)
+                """ <<< gradient clipping <<< """
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                loss.backward()
+                """ >>> gradient clipping >>> """
+                torch.nn.utils.clip_grad_norm_(unet.parameters(), config.optimizer.max_grad_norm)
+                """ <<< gradient clipping <<< """
+                optimizer.step()
+
+            # Check the grad of attn blocks for debugging
+            # print(unet.module.up_blocks[3].attentions[2].transformer_blocks[0].audio_cross_attn.attn.to_q.weight.grad)
+
+            lr_scheduler.step()
+            progress_bar.update(1)
+            global_step += 1
+
+            ### <<<< Training <<<< ###
+
+            # Save checkpoint and conduct validation
+            if is_main_process and (global_step % config.ckpt.save_ckpt_steps == 0):
+                if config.run.recon_loss_weight != 0:
+                    plot_loss_chart(
+                        os.path.join(output_dir, f"loss_charts/recon_loss_chart-{global_step}.png"),
+                        ("Reconstruction loss", train_step_list, recon_loss_list),
+                    )
+                if config.model.add_audio_layer:
+                    if sync_loss_list != []:
+                        plot_loss_chart(
+                            os.path.join(output_dir, f"loss_charts/sync_loss_chart-{global_step}.png"),
+                            ("Sync loss", train_step_list, sync_loss_list),
+                        )
+                model_save_path = os.path.join(output_dir, f"checkpoints/checkpoint-{global_step}.pt")
+                state_dict = {
+                    "global_step": global_step,
+                    "state_dict": unet.module.state_dict(),  # to unwrap DDP
+                }
+                try:
+                    torch.save(state_dict, model_save_path)
+                    logger.info(f"Saved checkpoint to {model_save_path}")
+                except Exception as e:
+                    logger.error(f"Error saving model: {e}")
+
+                # Validation
+                logger.info("Running validation... ")
+
+                validation_video_out_path = os.path.join(output_dir, f"val_videos/val_video_{global_step}.mp4")
+                validation_video_mask_path = os.path.join(output_dir, f"val_videos/val_video_mask.mp4")
+
+                with torch.autocast(device_type="cuda", dtype=torch.float16):
+                    pipeline(
+                        config.data.val_video_path,
+                        config.data.val_audio_path,
+                        validation_video_out_path,
+                        validation_video_mask_path,
+                        num_frames=config.data.num_frames,
+                        num_inference_steps=config.run.inference_steps,
+                        guidance_scale=config.run.guidance_scale,
+                        weight_dtype=torch.float16,
+                        width=config.data.resolution,
+                        height=config.data.resolution,
+                        mask=config.data.mask,
+                    )
+
+                logger.info(f"Saved validation video output to {validation_video_out_path}")
+
+                val_step_list.append(global_step)
+
+                if config.model.add_audio_layer:
+                    try:
+                        _, conf = syncnet_eval(syncnet_eval_model, syncnet_detector, validation_video_out_path, "temp")
+                    except Exception as e:
+                        logger.info(e)
+                        conf = 0
+                    sync_conf_list.append(conf)
+                    plot_loss_chart(
+                        os.path.join(output_dir, f"loss_charts/sync_conf_chart-{global_step}.png"),
+                        ("Sync confidence", val_step_list, sync_conf_list),
+                    )
+
+            logs = {"step_loss": loss.item(), "lr": lr_scheduler.get_last_lr()[0]}
+            progress_bar.set_postfix(**logs)
+
+            if global_step >= config.run.max_train_steps:
+                break
+
+    progress_bar.close()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    # Config file path
+    parser.add_argument("--unet_config_path", type=str, default="configs/unet.yaml")
+
+    args = parser.parse_args()
+    config = OmegaConf.load(args.unet_config_path)
+    config.unet_config_path = args.unet_config_path
+
+    main(config)