Initial upload of directory

.gitattributes

@@ -33,3 +33,11 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/demo1_audio.wav filter=lfs diff=lfs merge=lfs -text
+assets/demo1_video.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/demo2_audio.wav filter=lfs diff=lfs merge=lfs -text
+assets/demo2_video.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/demo3_audio.wav filter=lfs diff=lfs merge=lfs -text
+assets/demo3_video.mp4 filter=lfs diff=lfs merge=lfs -text
+assets/framework.png filter=lfs diff=lfs merge=lfs -text
+temp/video.mp4 filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,10 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv

.python-version

@@ -0,0 +1 @@
+3.11

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

@@ -0,0 +1,15 @@
+#OpenLipSync
+
+This is a small repo containing all the required files to run inference for LatentSync1.5. 
+
+TODO:
+add MuseTalk checkpoints
+add LatentSync16 checkpoints
+
+Installation
+- clone the repo
+- On debian based systems run bash debian_setup.sh
+
+Run
+- for inference modify the scritpts/inference.py file add your video and audio path
+- run with uv run python -m scripts.inference

assets/demo1_audio.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f7dd2112dbdc0bece5ee6f26553a4867b65740eb53187ecc8b1a3c1618b2405
+size 307278

assets/demo1_video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed2dd1e2001aa605c3f2d77672a8af4ed55e427a85c55d408adfc3d5076bc872
+size 1240008

assets/demo2_audio.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4916574779fb975367ddcb1f12597205ae15ea8aeaa61ad92d2c1c5d719c3607
+size 634958

assets/demo2_video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c3f10288e0642e587a95c0040e6966f8f6b7e003c3a17b572f72472b896d8ff
+size 1772492

assets/demo3_audio.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d5014567b03d35e0bd813a3725c3129a99722497cd4cf8e036d2c304530ea432
+size 593998

assets/demo3_video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cfa177b2a44f7809f606285c120e270d526caa50d708ec95e0f614d220970e0f
+size 2112370

checkpoints/.gitattributes

@@ -0,0 +1,33 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+diffusion_pytorch_model.safetensors filter=lfs diff=lfs merge=lfs -text

checkpoints/.gitignore

@@ -0,0 +1 @@
+.DS_Store

checkpoints/auxiliary/2DFAN4-cd938726ad.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd938726adb1f15f361263cce2db9cb820c42585fa8796ec72ce19107f369a46
+size 96316515

checkpoints/auxiliary/i3d_torchscript.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bec6519f66ea534e953026b4ae2c65553c17bf105611c746d904657e5860a5e2
+size 51235320

checkpoints/auxiliary/koniq_pretrained.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff9277bcc68ecc10e77d88b6d0a32825ec3c85562095542734ec6212eaaf6d81
+size 109768650

checkpoints/auxiliary/s3fd-619a316812.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:619a31681264d3f7f7fc7a16a42cbbe8b23f31a256f75a366e5a1bcd59b33543
+size 89843225

checkpoints/auxiliary/sfd_face.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d54a87c2b7543b64729c9a25eafd188da15fd3f6e02f0ecec76ae1b30d86c491
+size 89844381

checkpoints/auxiliary/syncnet_v2.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:961e8696f888fce4f3f3a6c3d5b3267cf5b343100b238e79b2659bff2c605442
+size 54573114

checkpoints/auxiliary/vgg16-397923af.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:397923af8e79cdbb6a7127f12361acd7a2f83e06b05044ddf496e83de57a5bf0
+size 553433881

checkpoints/auxiliary/vit_g_hybrid_pt_1200e_ssv2_ft.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a210a92f035dff30c53b46157b612e7a1a5d3c99700e1b2d71da5c399ca7e70
+size 2023804201

checkpoints/latentsync/README.md

@@ -0,0 +1,14 @@
+---
+license: openrail++
+library_name: diffusers
+tags:
+- video-to-video
+---
+
+# The checkpoints of LatentSync
+
+This repo not only stores the pretrained U-Net and SyncNet checkpoints of LatentSync, but also stores the whisper checkpoints, auxiliary checkpoints for detecting face, calculating syncnet confidence score and so on. They have covered all you need for both inference and training of LatentSync
+
+Paper: https://arxiv.org/abs/2412.09262
+
+Code: https://github.com/bytedance/LatentSync

checkpoints/latentsync/config.json

@@ -0,0 +1,3 @@
+{
+    "Name": "LatentSync"
+}

checkpoints/latentsync/latentsync_syncnet.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:38fa63bad3ed2332f647c40a5dc616cb0e233db8579f698f62af4c41965c4da5
+size 1488019828

checkpoints/latentsync/latentsync_unet.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:63197c73d21ad55ddf2b6e5cc38d0a19a1e494317aefe2707c6b6c6fc952f3c7
+size 3400080614

checkpoints/sd-vae-ft-mse/README.md

@@ -0,0 +1,83 @@
+---
+license: mit
+tags:
+- stable-diffusion
+- stable-diffusion-diffusers
+inference: false
+---
+# Improved Autoencoders
+
+## Utilizing
+These weights are intended to be used with the [🧨 diffusers library](https://github.com/huggingface/diffusers). If you are looking for the model to use with the original [CompVis Stable Diffusion codebase](https://github.com/CompVis/stable-diffusion), [come here](https://huggingface.co/stabilityai/sd-vae-ft-mse-original).
+
+#### How to use with 🧨 diffusers
+You can integrate this fine-tuned VAE decoder to your existing `diffusers` workflows, by including a `vae` argument to the `StableDiffusionPipeline`
+```py
+from diffusers.models import AutoencoderKL
+from diffusers import StableDiffusionPipeline
+
+model = "CompVis/stable-diffusion-v1-4"
+vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
+pipe = StableDiffusionPipeline.from_pretrained(model, vae=vae)
+```
+
+## Decoder Finetuning
+We publish two kl-f8 autoencoder versions, finetuned from the original [kl-f8 autoencoder](https://github.com/CompVis/latent-diffusion#pretrained-autoencoding-models) on a 1:1 ratio of [LAION-Aesthetics](https://laion.ai/blog/laion-aesthetics/) and LAION-Humans, an unreleased subset containing only SFW images of humans. The intent was to fine-tune on the Stable Diffusion training set (the autoencoder was originally trained on OpenImages) but also enrich the dataset with images of humans to improve the reconstruction of faces.
+The first, _ft-EMA_, was resumed from the original checkpoint, trained for 313198 steps and uses EMA weights. It uses the same loss configuration as the original checkpoint (L1 + LPIPS).
+The second, _ft-MSE_, was resumed from _ft-EMA_ and uses EMA weights and was trained for another 280k steps using a different loss, with more emphasis 
+on MSE reconstruction (MSE + 0.1 * LPIPS). It produces somewhat ``smoother'' outputs. The batch size for both versions was 192 (16 A100s, batch size 12 per GPU).
+To keep compatibility with existing models, only the decoder part was finetuned; the checkpoints can be used as a drop-in replacement for the existing autoencoder.
+
+_Original kl-f8 VAE vs f8-ft-EMA vs f8-ft-MSE_
+
+## Evaluation 
+### COCO 2017 (256x256, val, 5000 images)
+| Model    | train steps | rFID | PSNR         | SSIM          | PSIM          | Link                                                                              | Comments                                                                                        
+|----------|---------|------|--------------|---------------|---------------|-----------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------|
+|          |         |      |              |               |               |                                                                                   |                                                                                                 |
+| original | 246803        | 4.99 | 23.4 +/- 3.8 | 0.69 +/- 0.14 | 1.01 +/- 0.28 | https://ommer-lab.com/files/latent-diffusion/kl-f8.zip                            | as used in SD                                                                                   |
+| ft-EMA   | 560001        | 4.42 | 23.8 +/- 3.9 | 0.69 +/- 0.13 | 0.96 +/- 0.27 | https://huggingface.co/stabilityai/sd-vae-ft-ema-original/resolve/main/vae-ft-ema-560000-ema-pruned.ckpt | slightly better overall, with EMA                                                               |
+| ft-MSE   | 840001        | 4.70 | 24.5 +/- 3.7 | 0.71 +/- 0.13 | 0.92 +/- 0.27 | https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt | resumed with EMA from ft-EMA, emphasis on MSE (rec. loss = MSE + 0.1 * LPIPS), smoother outputs |
+
+
+### LAION-Aesthetics 5+ (256x256, subset, 10000 images)
+| Model    | train steps | rFID | PSNR         | SSIM          | PSIM          | Link                                                                              | Comments                                                                                        
+|----------|-----------|------|--------------|---------------|---------------|-----------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------|
+|          |           |      |              |               |               |                                                                                   |                                                                                                 |
+| original | 246803         | 2.61 | 26.0 +/- 4.4 | 0.81 +/- 0.12 | 0.75 +/- 0.36 | https://ommer-lab.com/files/latent-diffusion/kl-f8.zip                            | as used in SD                                                                                   |
+| ft-EMA   | 560001          | 1.77 | 26.7 +/- 4.8 | 0.82 +/- 0.12 | 0.67 +/- 0.34 | https://huggingface.co/stabilityai/sd-vae-ft-ema-original/resolve/main/vae-ft-ema-560000-ema-pruned.ckpt | slightly better overall, with EMA                                                               |
+| ft-MSE   | 840001          | 1.88 | 27.3 +/- 4.7 | 0.83 +/- 0.11 | 0.65 +/- 0.34 | https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt | resumed with EMA from ft-EMA, emphasis on MSE (rec. loss = MSE + 0.1 * LPIPS), smoother outputs |
+
+
+### Visual
+_Visualization of reconstructions on  256x256 images from the COCO2017 validation dataset._ 
+
+<p align="center">
+  <br>
+  <b>
+256x256: ft-EMA (left), ft-MSE (middle), original (right)</b>
+</p>
+
+<p align="center">
+<img src=https://huggingface.co/stabilityai/stable-diffusion-decoder-finetune/resolve/main/eval/ae-decoder-tuning-reconstructions/merged/00025_merged.png />
+</p>
+
+<p align="center">
+<img src=https://huggingface.co/stabilityai/stable-diffusion-decoder-finetune/resolve/main/eval/ae-decoder-tuning-reconstructions/merged/00011_merged.png />
+</p>
+
+<p align="center">
+<img src=https://huggingface.co/stabilityai/stable-diffusion-decoder-finetune/resolve/main/eval/ae-decoder-tuning-reconstructions/merged/00037_merged.png />
+</p>
+
+<p align="center">
+<img src=https://huggingface.co/stabilityai/stable-diffusion-decoder-finetune/resolve/main/eval/ae-decoder-tuning-reconstructions/merged/00043_merged.png />
+</p>
+
+<p align="center">
+<img src=https://huggingface.co/stabilityai/stable-diffusion-decoder-finetune/resolve/main/eval/ae-decoder-tuning-reconstructions/merged/00053_merged.png />
+</p>
+
+<p align="center">
+<img src=https://huggingface.co/stabilityai/stable-diffusion-decoder-finetune/resolve/main/eval/ae-decoder-tuning-reconstructions/merged/00029_merged.png />
+</p>

checkpoints/sd-vae-ft-mse/config.json

@@ -0,0 +1,29 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.4.2",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "in_channels": 3,
+  "latent_channels": 4,
+  "layers_per_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 256,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ]
+}

checkpoints/sd-vae-ft-mse/diffusion_pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b4889b6b1d4ce7ae320a02dedaeff1780ad77d415ea0d744b476155c6377ddc
+size 334707217

checkpoints/sd-vae-ft-mse/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815
+size 334643276

checkpoints/whisper/tiny.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9
+size 75572083

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

debian_deps.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+sudo apt -y install libgl1
+sudo apt -y install ffmpeg
+sudo apt -y install curl
+curl -LsSf https://astral.sh/uv/install.sh | sh
+# OpenCV dependencies
+uv init --python=3.11
+uv add pip
+uv run pip install -r requirements.txt

inference.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+uv run python -m scripts.inference

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")))
+
+    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).cuda()
+            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)