install the engines the presenter has been preferring all along

LatentSync and MuseTalk have been selectable by name for a while, and
`engine: "auto"` has ranked both ahead of wav2lip whenever there was a clip to
drive -- wav2lip works at 96x96 and upscales, which is the soft mouth on a 1080p
reel. Nothing ever installed them, so every selection ended at "neither on PATH
nor installed", and the call that would have been made could not have worked
either: LatentSync requires its checkpoint and UNet config as well as the three
paths, and MuseTalk takes a YAML task list and writes into a result directory
rather than taking paths at all.

Coderai orchestrates what it serves, so it installs these. The inference source
is vendored under third_party/ and pinned to a commit, for the reason LongCat's
is: an image someone else pulls has to run without network access to a third
party. The weights come from the Hub on first use, never during a render. Each
engine gets its own venv on its own python, because both pin stacks the server
cannot host.

Upstream MuseTalk takes its BiSeNet face-parsing checkpoint off a Google Drive id
and a pytorch.org URL; both come from a Hub mirror here, so the whole set has one
provenance.

_managed_engine_ready used to answer "an entrypoint and any .pt anywhere", which
reads an interrupted download as a working install. It now asks the module that
knows the file list, and the old heuristic keeps covering an engine an operator
installed by hand. /v1/video/presenter/engines reports what an uninstalled engine
is short of, so a client can tell "not yet" from "never here".

Verified live: LatentSync installed from nothing and lip-synced a 2s 480x480 clip
to real narration in 136s.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014T3rZnLhydX9QqUcfBbqLu

codai/api/compose.py

@@ -1957,9 +1957,22 @@ async def presenter_engines(_auth=Depends(_require_api_auth)):
         matte = matting.engine_name()
     except Exception:
         matte = None
+    # An engine coderai can install but has not is neither "installed" nor
+    # permanently absent, and a client that cannot tell the two apart greys the
+    # feature out for good. Say which ones an operator can still bring up, and
+    # what each is short of.
+    installable = {}
+    try:
+        from codai.api import lipsync_engines as _le
+        for _name in _le.MANAGED:
+            if _name not in have:
+                installable[_name] = _le.missing_pieces(_name)
+    except Exception:
+        pass
     return {"engines": have, "available": bool(have),
             "video_engines": [e for e in have if e in _pres.VIDEO_ENGINES],
             "portrait_engines": [e for e in have if e in _pres.PORTRAIT_ENGINES],
+            "installable": installable,
             "layouts": ["full", "pip", "split"],
             # remove_background is a separate model from the lip-sync engine, so
             # a client greys out that toggle on its own.

codai/compose/presenter.py

@@ -71,10 +71,21 @@ def _managed_engine_dir(name: str):
 def _managed_engine_ready(name: str) -> bool:
     """Whether a managed install of this engine is usable WITHOUT downloading.
 
-    An inference entrypoint and at least one checkpoint. Deliberately not
-    "could be downloaded": a compose job is the wrong place to fetch several
-    GB, which is the same rule wav2lip follows above.
+    Deliberately not "could be downloaded": a compose job is the wrong place to
+    fetch several GB, which is the same rule wav2lip follows above.
+
+    :mod:`codai.api.lipsync_engines` owns the real answer, because it knows which
+    files each engine needs rather than guessing from "an entrypoint and at least
+    one checkpoint". The guess counted an interrupted download as a working
+    install. The fallback below keeps an operator's hand-rolled install
+    selectable for an engine this build does not manage.
     """
+    try:
+        from codai.api import lipsync_engines
+        if name in lipsync_engines.MANAGED:
+            return lipsync_engines.is_installed(name)
+    except Exception:
+        pass
     d = _managed_engine_dir(name)
     if not d:
         return False
@@ -180,10 +191,9 @@ def run_engine(engine: str, face_path: str, audio_path: str, out_path: str,
 
     if eng in ("latentsync", "musetalk"):
         # Both are video-driven: they replace the mouth in a clip, so `face`
-        # here is a base clip rather than a still. Both also take the same
-        # three arguments under their own CLI, and both upstream repos drive
-        # them through a scripts/inference.py, so a shim on PATH is tried
-        # first and the managed checkout second.
+        # here is a base clip rather than a still. A shim on PATH is tried
+        # first, because an operator who put one there meant it to be used;
+        # the managed install is the fallback.
         binary = shutil.which(eng)
         if binary:
             cmd = [binary, "--video_path", face_path, "--audio_path", audio_path,
@@ -198,26 +208,21 @@ def run_engine(engine: str, face_path: str, audio_path: str, out_path: str,
             if os.path.isfile(out_path):
                 return out_path
             raise PresenterError(f"{eng} produced no output")
-        d = _managed_engine_dir(eng)
-        script = None
-        for cand in ("scripts/inference.py", "inference.py", "run.py"):
-            if d and (d / cand).exists():
-                script = d / cand
-                break
-        if script is None:
-            raise PresenterError(
-                f"{eng} is neither on PATH nor installed under {d} "
-                f"(expected scripts/inference.py and its checkpoints)")
-        import sys as _sys
-        env_extra = str(d)
-        cmd = [_sys.executable, str(script),
-               "--video_path", os.path.abspath(face_path),
-               "--audio_path", os.path.abspath(audio_path),
-               "--video_out_path", os.path.abspath(out_path)]
-        _run(cmd, cancel=cancel, timeout=timeout, log=log, cwd=str(d),
-             pythonpath=env_extra)
-        if not os.path.isfile(out_path):
-            raise PresenterError(f"{eng} produced no output")
+        # The managed install. Neither engine actually takes those three flags
+        # alone -- LatentSync requires its checkpoint and UNet config too, and
+        # MuseTalk takes a YAML task list rather than paths and writes into a
+        # result directory -- and neither runs on this interpreter. So the module
+        # that installs them also knows how to call them.
+        try:
+            from codai.api import lipsync_engines
+        except Exception as exc:
+            raise PresenterError(f"{eng} support is unavailable: {exc}") from exc
+        try:
+            return lipsync_engines.run(
+                eng, os.path.abspath(face_path), os.path.abspath(audio_path),
+                os.path.abspath(out_path), timeout=timeout, log=log)
+        except Exception as exc:
+            raise PresenterError(str(exc)) from exc
         return out_path
 
     if eng == "sadtalker":

docs/video-compose.md

@@ -329,6 +329,24 @@ identity-locked video generation are here too.
   instead of replacing a mouth. Any other engine installed
   under its own name can be named instead — no API change for a new one. Ask
   `GET /v1/video/presenter/engines` what this install has.
+* **Installing latentsync / musetalk.** Coderai manages both: the inference source
+  is vendored in the repo under `third_party/` and pinned to a commit, the weights
+  come from the Hub on first use, and each engine gets its own venv on its own
+  python (both pin stacks the server cannot host — diffusers 0.32/0.30 against the
+  server's 0.42, and MuseTalk needs mmpose). A render never installs them, because
+  a compose job is the wrong place to fetch several GB, so they are reported as
+  unavailable until installed:
+
+  ```python
+  from codai.api import lipsync_engines
+  lipsync_engines.ensure_installed("latentsync", log=print)   # code + weights + venv
+  lipsync_engines.is_available("latentsync")
+  lipsync_engines.missing_pieces("musetalk")                  # what an install is short of
+  ```
+
+  `CODERAI_LATENTSYNC_DIR` / `CODERAI_MUSETALK_DIR` point at an existing install,
+  and `CODERAI_LATENTSYNC_VENV` / `CODERAI_MUSETALK_VENV` at an existing venv.
+  An engine put on PATH by hand still wins over the managed install.
 * **`layout`**: `full` replaces the scene's visuals (a scene may then have none);
   `pip` composites a bubble over them inside the caption safe area — and captions
   at the bottom move above it automatically; `split` stacks presenter and B-roll,

requirements-latentsync.txt

+# Isolated venv for LatentSync (ByteDance), the video-driven lip-sync engine.
+#
+# It cannot share the server's environment: it pins diffusers 0.32 and numpy 1.x
+# against the server's diffusers 0.42 and numpy 2.x, and its UNet3DConditionModel
+# is written against that diffusers. So it runs in its own venv on its own python
+# and coderai drives scripts.inference as a subprocess
+# (codai/api/lipsync_engines.py). Do NOT merge these into requirements.txt.
+#
+# python 3.11: decord publishes nothing for 3.13 and the pinned torch is a cp311
+# build. 3.11 is the newest interpreter the whole set agrees on.
+#
+# Trimmed against what the inference path actually imports. Upstream's
+# requirements.txt also lists gradio (its demo app), mediapipe, face-alignment,
+# lpips, python_speech_features and scenedetect — none of which any module
+# reached from scripts.inference imports.
+
+--extra-index-url https://download.pytorch.org/whl/cu124
+torch==2.5.1+cu124
+torchvision==0.20.1+cu124
+
+diffusers==0.32.2
+transformers==4.48.0
+accelerate==0.26.1
+huggingface-hub==0.30.2
+
+# The lip-sync pipeline and the vendored Whisper audio encoder.
+einops==0.7.0
+omegaconf==2.3.0
+numpy==1.26.4
+opencv-python==4.9.0.80
+librosa==0.10.1
+soundfile>=0.12.1
+decord==0.6.0
+kornia==0.8.0
+tqdm>=4.66.0
+
+# Face detection for the mouth region. onnxruntime is not in any import line in
+# the tree -- insightface loads it lazily from inside model_zoo, so dropping it
+# as "unused" leaves an install that passes every file-level check and then
+# raises ImportError on the first frame.
+insightface==0.7.3
+onnxruntime-gpu==1.21.0
+
+# The sampling speed-up scripts.inference enables with --enable_deepcache.
+DeepCache==0.1.1
+
+ffmpeg-python==0.2.0
+imageio==2.31.1
+imageio-ffmpeg==0.5.1

requirements-musetalk.txt

+# Isolated venv for MuseTalk (Tencent Music Entertainment Lyra Lab).
+#
+# Separate for the same reason as LatentSync: it pins diffusers 0.30 and numpy
+# 1.x, and it needs mmpose, whose mmcv extension is built against one exact torch
+# build. coderai drives scripts.inference as a subprocess
+# (codai/api/lipsync_engines.py). Do NOT merge these into requirements.txt.
+#
+# python 3.11 — mmcv publishes no wheel for 3.13, and numpy 1.23 none either.
+#
+# Two notes on what is NOT here. Upstream's requirements.txt lists
+# tensorflow==2.12.0 and tensorboard: nothing reachable from scripts.inference
+# imports either, and they are ~600 MB. It also lists gdown, only used by
+# download_weights.sh to pull the BiSeNet face-parsing checkpoint off a Google
+# Drive id; coderai takes that file from a Hub mirror instead.
+#
+# mmcv is the fragile one. openmmlab publishes prebuilt wheels for a handful of
+# torch/CUDA pairs and torch 2.4.1+cu121 is one of them, which is why the torch
+# pin here is older than LatentSync's. If this host has no matching wheel, mmcv
+# builds from source and needs nvcc; ensure_installed() reports that rather than
+# leaving a half-built venv behind.
+
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch==2.4.1+cu121
+torchvision==0.19.1+cu121
+
+diffusers==0.30.2
+transformers==4.39.2
+accelerate==0.28.0
+huggingface-hub==0.30.2
+
+numpy==1.23.5
+opencv-python==4.9.0.80
+soundfile==0.12.1
+librosa==0.11.0
+einops==0.8.1
+omegaconf
+tqdm>=4.66.0
+imageio[ffmpeg]
+ffmpeg-python
+requests
+
+# DWPose landmarks. preprocessing.py calls mmpose.apis at import time, so these
+# are not optional.
+-f https://download.openmmlab.com/mmcv/dist/cu121/torch2.4/index.html
+mmcv==2.2.0
+mmengine>=0.10.3
+mmpose==1.3.2

tests/test_ltx2_audio_and_engines.py

@@ -134,19 +134,27 @@ def test_the_new_engines_are_video_driven_and_preferred():
 
 
 def test_an_engine_is_only_available_when_its_assets_are_here(monkeypatch, tmp_path):
+    """The directory heuristic, on an engine coderai does not manage.
+
+    It used to answer for latentsync and musetalk too, and "an entrypoint plus
+    any .pt anywhere" read an interrupted download as a working install. Those
+    two now answer from codai.api.lipsync_engines, which knows the file list —
+    see tests/test_lipsync_engines.py. The heuristic still covers an engine an
+    operator installed by hand under its own name, which is what is checked here.
+    """
     from codai.compose import presenter as P
     monkeypatch.setattr(P.shutil, "which", lambda n: None)
     monkeypatch.setattr(P, "_managed_engine_dir", lambda n: tmp_path / n)
     # Nothing installed.
-    assert P._managed_engine_ready("latentsync") is False
+    assert P._managed_engine_ready("hallo") is False
     # An entrypoint but no weights is NOT ready: a compose job is the wrong
     # place to discover a 5 GB download is missing.
-    d = tmp_path / "latentsync"
+    d = tmp_path / "hallo"
     (d / "scripts").mkdir(parents=True)
     (d / "scripts" / "inference.py").write_text("x")
-    assert P._managed_engine_ready("latentsync") is False
+    assert P._managed_engine_ready("hallo") is False
     (d / "unet.pt").write_text("w")
-    assert P._managed_engine_ready("latentsync") is True
+    assert P._managed_engine_ready("hallo") is True
 
 
 def test_available_engines_is_reported_in_preference_order(monkeypatch):

third_party/LATENTSYNC-PINNED

+LatentSync (ByteDance) — vendored inference source.
+
+Upstream:  https://github.com/bytedance/LatentSync
+Commit:    a229c3948406bc2cf6eaf4873e662e70c6a04746  (2025-06-20)
+Licence:   Apache-2.0 — see LATENTSYNC-LICENSE
+
+What is here and why: `latentsync/` (the UNet, the lip-sync pipeline, the vendored
+OpenAI Whisper audio encoder and their utilities), `scripts/inference.py` and the
+`configs/` the entrypoint reads at runtime. Training scripts, the SyncNet configs,
+the Gradio app, the demo assets and the evaluation corpus are not here.
+
+Not cloned at install time, for the same reason the LongCat source is vendored: an
+image someone else pulls has to run without network access to a third party. The
+WEIGHTS are a separate matter and are fetched from the Hub on first use.

third_party/MUSETALK-PINNED

+MuseTalk (Tencent Music Entertainment Lyra Lab) — vendored inference source.
+
+Upstream:  https://github.com/TMElyralab/MuseTalk
+Commit:    0a89dec45a0192b824e3cf4daf96c239440c5ed8  (2025-09-26)
+Licence:   MIT — see MUSETALK-LICENSE
+
+What is here and why: `musetalk/` (the UNet wrapper, the VAE, face parsing, the
+DWPose preprocessing and the vendored OpenAI Whisper audio encoder),
+`scripts/inference.py` and `configs/` (the task-list templates). Training, the
+realtime demo, the Gradio app and the sample media are not here.
+
+Not cloned at install time — see LATENTSYNC-PINNED for the reasoning. The WEIGHTS
+are fetched from the Hub on first use; upstream's download_weights.sh takes the
+BiSeNet face-parsing checkpoint off a Google Drive id, which coderai does not use:
+it comes from a Hub mirror so the whole set has one provenance.

third_party/latentsync/configs/audio.yaml

+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

third_party/latentsync/configs/scheduler_config.json

+{
+  "_class_name": "DDIMScheduler",
+  "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
+}

third_party/latentsync/configs/unet/stage1.yaml

+data:
+  syncnet_config_path: configs/syncnet/syncnet_16_pixel_attn.yaml
+  train_output_dir: debug/unet
+  train_fileslist: /mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/data_v10_core.txt
+  train_data_dir: ""
+  audio_embeds_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/embeds
+  audio_mel_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/mel
+
+  val_video_path: assets/demo1_video.mp4
+  val_audio_path: assets/demo1_audio.wav
+  batch_size: 1 # 24
+  num_workers: 12 # 12
+  num_frames: 16
+  resolution: 256
+  mask_image_path: latentsync/utils/mask.png
+  audio_sample_rate: 16000
+  video_fps: 25
+  audio_feat_length: [2, 2]
+
+ckpt:
+  resume_ckpt_path: checkpoints/latentsync_unet.pt
+  save_ckpt_steps: 10000
+
+run:
+  pixel_space_supervise: false
+  use_syncnet: false
+  sync_loss_weight: 0.05
+  perceptual_loss_weight: 0.1 # 0.1
+  recon_loss_weight: 1 # 1
+  guidance_scale: 1.5 # [1.0 - 3.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: 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
+  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]
+
+  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: 24
+    temporal_attention_dim_div: 1
+    zero_initialize: true

third_party/latentsync/configs/unet/stage1_512.yaml

+data:
+  syncnet_config_path: configs/syncnet/syncnet_16_pixel_attn.yaml
+  train_output_dir: debug/unet
+  train_fileslist: /mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/data_v10_core.txt
+  train_data_dir: ""
+  audio_embeds_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/embeds
+  audio_mel_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/mel
+
+  val_video_path: assets/demo1_video.mp4
+  val_audio_path: assets/demo1_audio.wav
+  batch_size: 1 # 8
+  num_workers: 12 # 12
+  num_frames: 16
+  resolution: 512
+  mask_image_path: latentsync/utils/mask.png
+  audio_sample_rate: 16000
+  video_fps: 25
+  audio_feat_length: [2, 2]
+
+ckpt:
+  resume_ckpt_path: checkpoints/latentsync_unet.pt
+  save_ckpt_steps: 10000
+
+run:
+  pixel_space_supervise: false
+  use_syncnet: false
+  sync_loss_weight: 0.05
+  perceptual_loss_weight: 0.1 # 0.1
+  recon_loss_weight: 1 # 1
+  guidance_scale: 1.5 # [1.0 - 3.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: 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
+  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]
+
+  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: 24
+    temporal_attention_dim_div: 1
+    zero_initialize: true

third_party/latentsync/configs/unet/stage2.yaml

+data:
+  syncnet_config_path: configs/syncnet/syncnet_16_pixel_attn.yaml
+  train_output_dir: debug/unet
+  train_fileslist: /mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/data_v10_core.txt
+  train_data_dir: ""
+  audio_embeds_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/embeds
+  audio_mel_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/mel
+
+  val_video_path: assets/demo1_video.mp4
+  val_audio_path: assets/demo1_audio.wav
+  batch_size: 1 # 4
+  num_workers: 12 # 12
+  num_frames: 16
+  resolution: 256
+  mask_image_path: latentsync/utils/mask.png
+  audio_sample_rate: 16000
+  video_fps: 25
+  audio_feat_length: [2, 2]
+
+ckpt:
+  resume_ckpt_path: checkpoints/latentsync_unet.pt
+  save_ckpt_steps: 10000
+
+run:
+  pixel_space_supervise: true
+  use_syncnet: true
+  sync_loss_weight: 0.05
+  perceptual_loss_weight: 0.1 # 0.1
+  recon_loss_weight: 1 # 1
+  guidance_scale: 1.5 # [1.0 - 3.0]
+  trepa_loss_weight: 10
+  inference_steps: 20
+  trainable_modules:
+    - motion_modules.
+    - attentions.
+  seed: 1247
+  use_mixed_noise: true
+  mixed_noise_alpha: 1 # 1
+  mixed_precision_training: true
+  enable_gradient_checkpointing: 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
+  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]
+
+  use_motion_module: true
+  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: 24
+    temporal_attention_dim_div: 1
+    zero_initialize: true

third_party/latentsync/configs/unet/stage2_512.yaml

+data:
+  syncnet_config_path: configs/syncnet/syncnet_16_pixel_attn.yaml
+  train_output_dir: debug/unet
+  train_fileslist: /mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/data_v10_core.txt
+  train_data_dir: ""
+  audio_embeds_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/embeds
+  audio_mel_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/mel
+
+  val_video_path: assets/demo1_video.mp4
+  val_audio_path: assets/demo1_audio.wav
+  batch_size: 1 # 4
+  num_workers: 12 # 12
+  num_frames: 16
+  resolution: 512
+  mask_image_path: latentsync/utils/mask.png
+  audio_sample_rate: 16000
+  video_fps: 25
+  audio_feat_length: [2, 2]
+
+ckpt:
+  resume_ckpt_path: checkpoints/latentsync_unet.pt
+  save_ckpt_steps: 10000
+
+run:
+  pixel_space_supervise: true
+  use_syncnet: true
+  sync_loss_weight: 0.05
+  perceptual_loss_weight: 0.1 # 0.1
+  recon_loss_weight: 1 # 1
+  guidance_scale: 1.5 # [1.0 - 3.0]
+  trepa_loss_weight: 10
+  inference_steps: 20
+  trainable_modules:
+    - motion_modules.
+    - attentions.
+  seed: 1247
+  use_mixed_noise: true
+  mixed_noise_alpha: 1 # 1
+  mixed_precision_training: true
+  enable_gradient_checkpointing: 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
+  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]
+
+  use_motion_module: true
+  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: 24
+    temporal_attention_dim_div: 1
+    zero_initialize: true

third_party/latentsync/configs/unet/stage2_efficient.yaml

+data:
+  syncnet_config_path: configs/syncnet/syncnet_16_pixel_attn.yaml
+  train_output_dir: debug/unet
+  train_fileslist: /mnt/bn/maliva-gen-ai-v2/chunyu.li/fileslist/data_v10_core.txt
+  train_data_dir: ""
+  audio_embeds_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/embeds
+  audio_mel_cache_dir: /mnt/bn/maliva-gen-ai-v2/chunyu.li/audio_cache/mel
+
+  val_video_path: assets/demo1_video.mp4
+  val_audio_path: assets/demo1_audio.wav
+  batch_size: 1 # 4
+  num_workers: 12 # 12
+  num_frames: 16
+  resolution: 256
+  mask_image_path: latentsync/utils/mask.png
+  audio_sample_rate: 16000
+  video_fps: 25
+  audio_feat_length: [2, 2]
+
+ckpt:
+  resume_ckpt_path: checkpoints/latentsync_unet.pt
+  save_ckpt_steps: 10000
+
+run:
+  pixel_space_supervise: true
+  use_syncnet: true
+  sync_loss_weight: 0.05
+  perceptual_loss_weight: 0.1 # 0.1
+  recon_loss_weight: 1 # 1
+  guidance_scale: 1.5 # [1.0 - 3.0]
+  trepa_loss_weight: 0
+  inference_steps: 20
+  trainable_modules:
+    - motion_modules.
+    - attn2.
+  seed: 1247
+  use_mixed_noise: true
+  mixed_noise_alpha: 1 # 1
+  mixed_precision_training: true
+  enable_gradient_checkpointing: 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
+  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]
+
+  use_motion_module: true
+  motion_module_resolutions: [1, 2, 4, 8]
+  motion_module_mid_block: false
+  motion_module_decoder_only: true
+  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: 24
+    temporal_attention_dim_div: 1
+    zero_initialize: true

third_party/latentsync/latentsync/data/syncnet_dataset.py

+# 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 pathlib import Path
+
+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.image_processor = ImageProcessor(resolution=config.data.resolution)
+        self.audio_mel_cache_dir = config.data.audio_mel_cache_dir
+        Path(self.audio_mel_cache_dir).mkdir(parents=True, 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)
+            if wrong_start_idx == start_idx:
+                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):
+        self.worker_id = worker_id
+
+    def __getitem__(self, idx):
+        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, weights_only=True)
+                    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)
+
+                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

third_party/latentsync/latentsync/data/unet_dataset.py

+# 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 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
+import torch.nn.functional as F
+from pathlib import Path
+
+
+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
+
+        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.image_processor = ImageProcessor(
+            self.resolution, mask_image=load_fixed_mask(self.resolution, config.data.mask_image_path)
+        )
+        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
+        Path(self.audio_mel_cache_dir).mkdir(parents=True, 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)
+        gt_frames_index = np.arange(start_idx, start_idx + self.num_frames, dtype=int)
+
+        while True:
+            ref_start_idx = random.randint(0, total_num_frames - self.num_frames)
+            if ref_start_idx > start_idx - self.num_frames and ref_start_idx < start_idx + self.num_frames:
+                continue
+            ref_frames_index = np.arange(ref_start_idx, ref_start_idx + self.num_frames, dtype=int)
+            break
+
+        gt_frames = video_reader.get_batch(gt_frames_index).asnumpy()
+        ref_frames = video_reader.get_batch(ref_frames_index).asnumpy()
+
+        return gt_frames, ref_frames, start_idx
+
+    def worker_init_fn(self, worker_id):
+        self.worker_id = worker_id
+
+    def __getitem__(self, idx):
+        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
+
+                gt_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, weights_only=True)
+                        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_pixel_values, masked_pixel_values, masks = self.image_processor.prepare_masks_and_masked_images(
+                    gt_frames, affine_transform=False
+                )  # (f, c, h, w)
+                ref_pixel_values = self.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_pixel_values=gt_pixel_values,
+            masked_pixel_values=masked_pixel_values,
+            ref_pixel_values=ref_pixel_values,
+            mel=mel,
+            masks=masks,
+            video_path=video_path,
+            start_idx=start_idx,
+        )
+
+        return sample

third_party/latentsync/latentsync/models/resnet.py

+# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/resnet.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from einops import rearrange
+
+
+class InflatedConv3d(nn.Conv2d):
+    def forward(self, x):
+        video_length = x.shape[2]
+
+        x = rearrange(x, "b c f h w -> (b f) c h w")
+        x = super().forward(x)
+        x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length)
+
+        return x
+
+
+class InflatedGroupNorm(nn.GroupNorm):
+    def forward(self, x):
+        video_length = x.shape[2]
+
+        x = rearrange(x, "b c f h w -> (b f) c h w")
+        x = super().forward(x)
+        x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length)
+
+        return x
+
+
+class Upsample3D(nn.Module):
+    def __init__(self, channels, use_conv=False, use_conv_transpose=False, out_channels=None, name="conv"):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_conv_transpose = use_conv_transpose
+        self.name = name
+
+        conv = None
+        if use_conv_transpose:
+            raise NotImplementedError
+        elif use_conv:
+            self.conv = InflatedConv3d(self.channels, self.out_channels, 3, padding=1)
+
+    def forward(self, hidden_states, output_size=None):
+        assert hidden_states.shape[1] == self.channels
+
+        if self.use_conv_transpose:
+            raise NotImplementedError
+
+        # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
+        dtype = hidden_states.dtype
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(torch.float32)
+
+        # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
+        if hidden_states.shape[0] >= 64:
+            hidden_states = hidden_states.contiguous()
+
+        # if `output_size` is passed we force the interpolation output
+        # size and do not make use of `scale_factor=2`
+        if output_size is None:
+            hidden_states = F.interpolate(hidden_states, scale_factor=[1.0, 2.0, 2.0], mode="nearest")
+        else:
+            hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")
+
+        # If the input is bfloat16, we cast back to bfloat16
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(dtype)
+
+        hidden_states = self.conv(hidden_states)
+
+        return hidden_states
+
+
+class Downsample3D(nn.Module):
+    def __init__(self, channels, use_conv=False, out_channels=None, padding=1, name="conv"):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.padding = padding
+        stride = 2
+        self.name = name
+
+        if use_conv:
+            self.conv = InflatedConv3d(self.channels, self.out_channels, 3, stride=stride, padding=padding)
+        else:
+            raise NotImplementedError
+
+    def forward(self, hidden_states):
+        assert hidden_states.shape[1] == self.channels
+        if self.use_conv and self.padding == 0:
+            raise NotImplementedError
+
+        assert hidden_states.shape[1] == self.channels
+        hidden_states = self.conv(hidden_states)
+
+        return hidden_states
+
+
+class ResnetBlock3D(nn.Module):
+    def __init__(
+        self,
+        *,
+        in_channels,
+        out_channels=None,
+        conv_shortcut=False,
+        dropout=0.0,
+        temb_channels=512,
+        groups=32,
+        groups_out=None,
+        pre_norm=True,
+        eps=1e-6,
+        non_linearity="swish",
+        time_embedding_norm="default",
+        output_scale_factor=1.0,
+        use_in_shortcut=None,
+        use_inflated_groupnorm=False,
+    ):
+        super().__init__()
+        self.pre_norm = pre_norm
+        self.pre_norm = True
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+        self.time_embedding_norm = time_embedding_norm
+        self.output_scale_factor = output_scale_factor
+
+        if groups_out is None:
+            groups_out = groups
+
+        assert use_inflated_groupnorm != None
+        if use_inflated_groupnorm:
+            self.norm1 = InflatedGroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+        else:
+            self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+
+        self.conv1 = InflatedConv3d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        if temb_channels is not None:
+            if self.time_embedding_norm == "default":
+                time_emb_proj_out_channels = out_channels
+            elif self.time_embedding_norm == "scale_shift":
+                time_emb_proj_out_channels = out_channels * 2
+            else:
+                raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ")
+
+            self.time_emb_proj = torch.nn.Linear(temb_channels, time_emb_proj_out_channels)
+        else:
+            self.time_emb_proj = None
+
+        if self.time_embedding_norm == "scale_shift":
+            self.double_len_linear = torch.nn.Linear(time_emb_proj_out_channels, 2 * time_emb_proj_out_channels)
+        else:
+            self.double_len_linear = None
+
+        if use_inflated_groupnorm:
+            self.norm2 = InflatedGroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
+        else:
+            self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
+
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = InflatedConv3d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        if non_linearity == "swish":
+            self.nonlinearity = lambda x: F.silu(x)
+        elif non_linearity == "mish":
+            self.nonlinearity = Mish()
+        elif non_linearity == "silu":
+            self.nonlinearity = nn.SiLU()
+
+        self.use_in_shortcut = self.in_channels != self.out_channels if use_in_shortcut is None else use_in_shortcut
+
+        self.conv_shortcut = None
+        if self.use_in_shortcut:
+            self.conv_shortcut = InflatedConv3d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, input_tensor, temb):
+        hidden_states = input_tensor
+
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.nonlinearity(hidden_states)
+
+        hidden_states = self.conv1(hidden_states)
+
+        if temb is not None:
+            if temb.dim() == 2:
+                # input (1, 1280)
+                temb = self.time_emb_proj(self.nonlinearity(temb))
+                temb = temb[:, :, None, None, None]  # unsqueeze
+            else:
+                # input (1, 1280, 16)
+                temb = temb.permute(0, 2, 1)
+                temb = self.time_emb_proj(self.nonlinearity(temb))
+                if self.double_len_linear is not None:
+                    temb = self.double_len_linear(self.nonlinearity(temb))
+                temb = temb.permute(0, 2, 1)
+                temb = temb[:, :, :, None, None]
+
+        if temb is not None and self.time_embedding_norm == "default":
+            hidden_states = hidden_states + temb
+
+        hidden_states = self.norm2(hidden_states)
+
+        if temb is not None and self.time_embedding_norm == "scale_shift":
+            scale, shift = torch.chunk(temb, 2, dim=1)
+            hidden_states = hidden_states * (1 + scale) + shift
+
+        hidden_states = self.nonlinearity(hidden_states)
+
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        if self.conv_shortcut is not None:
+            input_tensor = self.conv_shortcut(input_tensor)
+
+        output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
+
+        return output_tensor
+
+
+class Mish(torch.nn.Module):
+    def forward(self, hidden_states):
+        return hidden_states * torch.tanh(torch.nn.functional.softplus(hidden_states))

third_party/latentsync/latentsync/models/stable_syncnet.py

+# 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
+from torch import nn
+from einops import rearrange
+from torch.nn import functional as F
+from .attention import Attention
+
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.models.attention import FeedForward
+from einops import rearrange
+
+
+class StableSyncNet(nn.Module):
+    def __init__(self, config, gradient_checkpointing=False):
+        super().__init__()
+        self.audio_encoder = DownEncoder2D(
+            in_channels=config["audio_encoder"]["in_channels"],
+            block_out_channels=config["audio_encoder"]["block_out_channels"],
+            downsample_factors=config["audio_encoder"]["downsample_factors"],
+            dropout=config["audio_encoder"]["dropout"],
+            attn_blocks=config["audio_encoder"]["attn_blocks"],
+            gradient_checkpointing=gradient_checkpointing,
+        )
+
+        self.visual_encoder = DownEncoder2D(
+            in_channels=config["visual_encoder"]["in_channels"],
+            block_out_channels=config["visual_encoder"]["block_out_channels"],
+            downsample_factors=config["visual_encoder"]["downsample_factors"],
+            dropout=config["visual_encoder"]["dropout"],
+            attn_blocks=config["visual_encoder"]["attn_blocks"],
+            gradient_checkpointing=gradient_checkpointing,
+        )
+
+        self.eval()
+
+    def forward(self, image_sequences, audio_sequences):
+        vision_embeds = self.visual_encoder(image_sequences)  # (b, c, 1, 1)
+        audio_embeds = self.audio_encoder(audio_sequences)  # (b, c, 1, 1)
+
+        vision_embeds = vision_embeds.reshape(vision_embeds.shape[0], -1)  # (b, c)
+        audio_embeds = audio_embeds.reshape(audio_embeds.shape[0], -1)  # (b, c)
+
+        # Make them unit vectors
+        vision_embeds = F.normalize(vision_embeds, p=2, dim=1)
+        audio_embeds = F.normalize(audio_embeds, p=2, dim=1)
+
+        return vision_embeds, audio_embeds
+
+
+class ResnetBlock2D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout: float = 0.0,
+        norm_num_groups: int = 32,
+        eps: float = 1e-6,
+        act_fn: str = "silu",
+        downsample_factor=2,
+    ):
+        super().__init__()
+
+        self.norm1 = nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=eps, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        self.norm2 = nn.GroupNorm(num_groups=norm_num_groups, num_channels=out_channels, eps=eps, affine=True)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+        if act_fn == "relu":
+            self.act_fn = nn.ReLU()
+        elif act_fn == "silu":
+            self.act_fn = nn.SiLU()
+
+        if in_channels != out_channels:
+            self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+        else:
+            self.conv_shortcut = None
+
+        if isinstance(downsample_factor, list):
+            downsample_factor = tuple(downsample_factor)
+
+        if downsample_factor == 1:
+            self.downsample_conv = None
+        else:
+            self.downsample_conv = nn.Conv2d(
+                out_channels, out_channels, kernel_size=3, stride=downsample_factor, padding=0
+            )
+            self.pad = (0, 1, 0, 1)
+            if isinstance(downsample_factor, tuple):
+                if downsample_factor[0] == 1:
+                    self.pad = (0, 1, 1, 1)  # The padding order is from back to front
+                elif downsample_factor[1] == 1:
+                    self.pad = (1, 1, 0, 1)
+
+    def forward(self, input_tensor):
+        hidden_states = input_tensor
+
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.act_fn(hidden_states)
+
+        hidden_states = self.conv1(hidden_states)
+        hidden_states = self.norm2(hidden_states)
+        hidden_states = self.act_fn(hidden_states)
+
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+
+        if self.conv_shortcut is not None:
+            input_tensor = self.conv_shortcut(input_tensor)
+
+        hidden_states += input_tensor
+
+        if self.downsample_conv is not None:
+            hidden_states = F.pad(hidden_states, self.pad, mode="constant", value=0)
+            hidden_states = self.downsample_conv(hidden_states)
+
+        return hidden_states
+
+
+class AttentionBlock2D(nn.Module):
+    def __init__(self, query_dim, norm_num_groups=32, dropout=0.0):
+        super().__init__()
+        self.norm1 = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=query_dim, eps=1e-6, affine=True)
+        self.norm2 = nn.LayerNorm(query_dim)
+        self.norm3 = nn.LayerNorm(query_dim)
+
+        self.ff = FeedForward(query_dim, dropout=dropout, activation_fn="geglu")
+
+        self.conv_in = nn.Conv2d(query_dim, query_dim, kernel_size=1, stride=1, padding=0)
+        self.conv_out = nn.Conv2d(query_dim, query_dim, kernel_size=1, stride=1, padding=0)
+
+        self.attn = Attention(query_dim=query_dim, heads=8, dim_head=query_dim // 8, dropout=dropout, bias=True)
+
+    def forward(self, hidden_states):
+        assert hidden_states.dim() == 4, f"Expected hidden_states to have ndim=4, but got ndim={hidden_states.dim()}."
+
+        batch, channel, height, width = hidden_states.shape
+        residual = hidden_states
+
+        hidden_states = self.norm1(hidden_states)
+        hidden_states = self.conv_in(hidden_states)
+        hidden_states = rearrange(hidden_states, "b c h w -> b (h w) c")
+
+        norm_hidden_states = self.norm2(hidden_states)
+
+        hidden_states = self.attn(norm_hidden_states, attention_mask=None) + hidden_states
+        hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
+
+        hidden_states = rearrange(hidden_states, "b (h w) c -> b c h w", h=height, w=width).contiguous()
+        hidden_states = self.conv_out(hidden_states)
+
+        hidden_states = hidden_states + residual
+        return hidden_states
+
+
+class DownEncoder2D(nn.Module):
+    def __init__(
+        self,
+        in_channels=4 * 16,
+        block_out_channels=[64, 128, 256, 256],
+        downsample_factors=[2, 2, 2, 2],
+        layers_per_block=2,
+        norm_num_groups=32,
+        attn_blocks=[1, 1, 1, 1],
+        dropout: float = 0.0,
+        act_fn="silu",
+        gradient_checkpointing=False,
+    ):
+        super().__init__()
+        self.layers_per_block = layers_per_block
+        self.gradient_checkpointing = gradient_checkpointing
+
+        # in
+        self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1)
+
+        # down
+        self.down_blocks = nn.ModuleList([])
+
+        output_channels = block_out_channels[0]
+        for i, block_out_channel in enumerate(block_out_channels):
+            input_channels = output_channels
+            output_channels = block_out_channel
+
+            down_block = ResnetBlock2D(
+                in_channels=input_channels,
+                out_channels=output_channels,
+                downsample_factor=downsample_factors[i],
+                norm_num_groups=norm_num_groups,
+                dropout=dropout,
+                act_fn=act_fn,
+            )
+
+            self.down_blocks.append(down_block)
+
+            if attn_blocks[i] == 1:
+                attention_block = AttentionBlock2D(query_dim=output_channels, dropout=dropout)
+                self.down_blocks.append(attention_block)
+
+        # out
+        self.norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
+        self.act_fn_out = nn.ReLU()
+
+    def forward(self, hidden_states):
+        hidden_states = self.conv_in(hidden_states)
+
+        # down
+        for down_block in self.down_blocks:
+            if self.gradient_checkpointing:
+                hidden_states = torch.utils.checkpoint.checkpoint(down_block, hidden_states, use_reentrant=False)
+            else:
+                hidden_states = down_block(hidden_states)
+
+        # post-process
+        hidden_states = self.norm_out(hidden_states)
+        hidden_states = self.act_fn_out(hidden_states)
+
+        return hidden_states

third_party/latentsync/latentsync/models/utils.py

+# 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.
+
+def zero_module(module):
+    # Zero out the parameters of a module and return it.
+    for p in module.parameters():
+        p.detach().zero_()
+    return module

third_party/latentsync/latentsync/models/wav2lip_syncnet.py

+# Adapted from https://github.com/primepake/wav2lip_288x288/blob/master/models/syncnetv2.py
+# The code here is for ablation study.
+
+from torch import nn
+from torch.nn import functional as F
+
+
+class Wav2LipSyncNet(nn.Module):
+    def __init__(self, act_fn="leaky"):
+        super().__init__()
+
+        # input image sequences: (15, 128, 256)
+        self.visual_encoder = nn.Sequential(
+            Conv2d(15, 32, kernel_size=(7, 7), stride=1, padding=3, act_fn=act_fn), # (128, 256)
+            Conv2d(32, 64, kernel_size=5, stride=(1, 2), padding=1, act_fn=act_fn), # (126, 127)
+            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(64, 128, kernel_size=3, stride=2, padding=1, act_fn=act_fn), # (63, 64)
+            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(128, 256, kernel_size=3, stride=3, padding=1, act_fn=act_fn), # (21, 22)
+            Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(256, 512, kernel_size=3, stride=2, padding=1, act_fn=act_fn), # (11, 11)
+            Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(512, 1024, kernel_size=3, stride=2, padding=1, act_fn=act_fn), # (6, 6)
+            Conv2d(1024, 1024, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(1024, 1024, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(1024, 1024, kernel_size=3, stride=2, padding=1, act_fn="relu"), # (3, 3)
+            Conv2d(1024, 1024, kernel_size=3, stride=1, padding=0, act_fn="relu"), # (1, 1)
+            Conv2d(1024, 1024, kernel_size=1, stride=1, padding=0, act_fn="relu"),
+        )
+
+        # input audio sequences: (1, 80, 16)
+        self.audio_encoder = nn.Sequential(
+            Conv2d(1, 32, kernel_size=3, stride=1, padding=1, act_fn=act_fn),
+            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1, act_fn=act_fn), # (27, 16)
+            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(64, 128, kernel_size=3, stride=3, padding=1, act_fn=act_fn), # (9, 6)
+            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1, act_fn=act_fn), # (3, 3)
+            Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(256, 512, kernel_size=3, stride=1, padding=1, act_fn=act_fn),
+            Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True, act_fn=act_fn),
+            Conv2d(512, 1024, kernel_size=3, stride=1, padding=0, act_fn="relu"), # (1, 1)
+            Conv2d(1024, 1024, kernel_size=1, stride=1, padding=0, act_fn="relu"),
+        )
+
+    def forward(self, image_sequences, audio_sequences):
+        vision_embeds = self.visual_encoder(image_sequences)  # (b, c, 1, 1)
+        audio_embeds = self.audio_encoder(audio_sequences)  # (b, c, 1, 1)
+
+        vision_embeds = vision_embeds.reshape(vision_embeds.shape[0], -1)  # (b, c)
+        audio_embeds = audio_embeds.reshape(audio_embeds.shape[0], -1)  # (b, c)
+
+        # Make them unit vectors
+        vision_embeds = F.normalize(vision_embeds, p=2, dim=1)
+        audio_embeds = F.normalize(audio_embeds, p=2, dim=1)
+
+        return vision_embeds, audio_embeds
+
+
+class Conv2d(nn.Module):
+    def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, act_fn="relu", *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.conv_block = nn.Sequential(nn.Conv2d(cin, cout, kernel_size, stride, padding), nn.BatchNorm2d(cout))
+        if act_fn == "relu":
+            self.act_fn = nn.ReLU()
+        elif act_fn == "tanh":
+            self.act_fn = nn.Tanh()
+        elif act_fn == "silu":
+            self.act_fn = nn.SiLU()
+        elif act_fn == "leaky":
+            self.act_fn = nn.LeakyReLU(0.2, inplace=True)
+
+        self.residual = residual
+
+    def forward(self, x):
+        out = self.conv_block(x)
+        if self.residual:
+            out += x
+        return self.act_fn(out)

third_party/latentsync/latentsync/trepa/loss.py

+# 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 torch.nn.functional as F
+from einops import rearrange
+from .third_party.VideoMAEv2.utils import load_videomae_model
+from ..utils.util import check_model_and_download
+
+
+class TREPALoss:
+    def __init__(
+        self,
+        device="cuda",
+        ckpt_path="checkpoints/auxiliary/vit_g_hybrid_pt_1200e_ssv2_ft.pth",
+        with_cp=False,
+    ):
+        check_model_and_download(ckpt_path)
+        self.model = load_videomae_model(device, ckpt_path, with_cp).eval().to(dtype=torch.float16)
+        self.model.requires_grad_(False)
+
+    def __call__(self, videos_fake, videos_real):
+        batch_size = videos_fake.shape[0]
+        num_frames = videos_fake.shape[2]
+        videos_fake = rearrange(videos_fake.clone(), "b c f h w -> (b f) c h w")
+        videos_real = rearrange(videos_real.clone(), "b c f h w -> (b f) c h w")
+
+        videos_fake = F.interpolate(videos_fake, size=(224, 224), mode="bicubic")
+        videos_real = F.interpolate(videos_real, size=(224, 224), mode="bicubic")
+
+        videos_fake = rearrange(videos_fake, "(b f) c h w -> b c f h w", f=num_frames)
+        videos_real = rearrange(videos_real, "(b f) c h w -> b c f h w", f=num_frames)
+
+        # Because input pixel range is [-1, 1], and model expects pixel range to be [0, 1]
+        videos_fake = (videos_fake / 2 + 0.5).clamp(0, 1)
+        videos_real = (videos_real / 2 + 0.5).clamp(0, 1)
+
+        feats_fake = self.model.forward_features(videos_fake)
+        feats_real = self.model.forward_features(videos_real)
+
+        feats_fake = F.normalize(feats_fake, p=2, dim=1)
+        feats_real = F.normalize(feats_real, p=2, dim=1)
+
+        return F.mse_loss(feats_fake, feats_real)
+
+
+if __name__ == "__main__":
+    torch.manual_seed(42)
+
+    # input shape: (b, c, f, h, w)
+    videos_fake = torch.randn(2, 3, 16, 256, 256, requires_grad=True).to(device="cuda", dtype=torch.float16)
+    videos_real = torch.randn(2, 3, 16, 256, 256, requires_grad=True).to(device="cuda", dtype=torch.float16)
+
+    trepa_loss = TREPALoss(device="cuda", with_cp=True)
+    loss = trepa_loss(videos_fake, videos_real)
+    print(loss)

third_party/latentsync/latentsync/trepa/third_party/VideoMAEv2/utils.py

+import os
+import torch
+import requests
+from tqdm import tqdm
+from torchvision import transforms
+from .videomaev2_finetune import vit_giant_patch14_224
+
+
+def to_normalized_float_tensor(vid):
+    return vid.permute(3, 0, 1, 2).to(torch.float32) / 255
+
+
+# NOTE: for those functions, which generally expect mini-batches, we keep them
+# as non-minibatch so that they are applied as if they were 4d (thus image).
+# this way, we only apply the transformation in the spatial domain
+def resize(vid, size, interpolation="bilinear"):
+    # NOTE: using bilinear interpolation because we don't work on minibatches
+    # at this level
+    scale = None
+    if isinstance(size, int):
+        scale = float(size) / min(vid.shape[-2:])
+        size = None
+    return torch.nn.functional.interpolate(vid, size=size, scale_factor=scale, mode=interpolation, align_corners=False)
+
+
+class ToFloatTensorInZeroOne(object):
+    def __call__(self, vid):
+        return to_normalized_float_tensor(vid)
+
+
+class Resize(object):
+    def __init__(self, size):
+        self.size = size
+
+    def __call__(self, vid):
+        return resize(vid, self.size)
+
+
+def preprocess_videomae(videos):
+    transform = transforms.Compose([ToFloatTensorInZeroOne(), Resize((224, 224))])
+    return torch.stack([transform(f) for f in torch.from_numpy(videos)])
+
+
+def load_videomae_model(device, ckpt_path=None, with_cp=False):
+    if ckpt_path is None:
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        ckpt_path = os.path.join(current_dir, "vit_g_hybrid_pt_1200e_ssv2_ft.pth")
+
+    if not os.path.exists(ckpt_path):
+        # download the ckpt to the path
+        ckpt_url = "https://pjlab-gvm-data.oss-cn-shanghai.aliyuncs.com/internvideo/videomaev2/vit_g_hybrid_pt_1200e_ssv2_ft.pth"
+        response = requests.get(ckpt_url, stream=True, allow_redirects=True)
+        total_size = int(response.headers.get("content-length", 0))
+        block_size = 1024
+
+        with tqdm(total=total_size, unit="B", unit_scale=True) as progress_bar:
+            with open(ckpt_path, "wb") as fw:
+                for data in response.iter_content(block_size):
+                    progress_bar.update(len(data))
+                    fw.write(data)
+
+    model = vit_giant_patch14_224(
+        img_size=224,
+        pretrained=False,
+        num_classes=174,
+        all_frames=16,
+        tubelet_size=2,
+        drop_path_rate=0.3,
+        use_mean_pooling=True,
+        with_cp=with_cp,
+    )
+
+    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+    for model_key in ["model", "module"]:
+        if model_key in ckpt:
+            ckpt = ckpt[model_key]
+            break
+    model.load_state_dict(ckpt)
+
+    del ckpt
+    torch.cuda.empty_cache()
+    return model.to(device)

third_party/latentsync/latentsync/whisper/whisper/__main__.py

+from .transcribe import cli
+
+
+cli()

third_party/latentsync/latentsync/whisper/whisper/assets/gpt2/special_tokens_map.json

+{"bos_token": "<|endoftext|>", "eos_token": "<|endoftext|>", "unk_token": "<|endoftext|>"}
\ No newline at end of file

third_party/latentsync/latentsync/whisper/whisper/assets/gpt2/tokenizer_config.json

+{"unk_token": "<|endoftext|>", "bos_token": "<|endoftext|>", "eos_token": "<|endoftext|>", "add_prefix_space": false, "model_max_length": 1024, "special_tokens_map_file": null, "name_or_path": "gpt2", "tokenizer_class": "GPT2Tokenizer"}
\ No newline at end of file

third_party/latentsync/latentsync/whisper/whisper/assets/multilingual/added_tokens.json

+{"<|endoftext|>": 50257}

third_party/latentsync/latentsync/whisper/whisper/assets/multilingual/special_tokens_map.json

+{"bos_token": "<|endoftext|>", "eos_token": "<|endoftext|>", "unk_token": "<|endoftext|>"}
\ No newline at end of file

third_party/latentsync/latentsync/whisper/whisper/assets/multilingual/tokenizer_config.json

+{"unk_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "bos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "eos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "add_prefix_space": false, "model_max_length": 1024, "special_tokens_map_file": null, "name_or_path": "multilingual", "errors": "replace", "tokenizer_class": "GPT2Tokenizer"}
\ No newline at end of file

third_party/latentsync/latentsync/whisper/whisper/normalizers/__init__.py

+from .basic import BasicTextNormalizer
+from .english import EnglishTextNormalizer

third_party/musetalk/musetalk/utils/__init__.py

+import sys
+from os.path import abspath, dirname
+current_dir = dirname(abspath(__file__))
+parent_dir = dirname(current_dir)
+sys.path.append(parent_dir+'/utils')

third_party/musetalk/musetalk/utils/face_detection/README.md

+The code for Face Detection in this folder has been taken from the wonderful [face_alignment](https://github.com/1adrianb/face-alignment) repository. This has been modified to take batches of faces at a time. 
\ No newline at end of file