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

@@ -9,4 +9,4 @@ app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

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)