model-dirs

app.py

@@ -26,11 +26,19 @@ sys.modules["torchvision.transforms.functional_tensor"] = _F
 os.environ["PROCESSED_RESULTS"] = os.path.join(os.getcwd(), "processed_results")
 os.makedirs(os.environ["PROCESSED_RESULTS"], exist_ok=True)
 
-src = "checkpoints"
+src = "/models"
 dst = os.path.expanduser("~/.cache/torch/hub/checkpoints")
 
 os.makedirs(dst, exist_ok=True)
 
+if os.path.exists(src):
+    for item in os.listdir(src):
+        src_path = os.path.join(src, item)
+        dst_path = os.path.join(dst, item)
+        if os.path.isfile(src_path) and not os.path.exists(dst_path):
+            shutil.copy2(src_path, dst_path)
+            print(f"Copied {item} to {dst}")
+
 print("Done copying checkpoints!")
 
 print("Loading LatentSync models...")

config.py

@@ -1,5 +1,10 @@
 """Configuration constants and global settings for OutofLipSync"""
 
+import os
+
+# Models directory - from environment variable or default to /models
+MODELS_DIR = os.getenv("MODELS_DIR", "/models")
+
 # Video settings
 DEFAULT_DURATION = 10
 MIN_DURATION = 5

eval/eval_syncnet_acc.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import argparse
+import os
+import sys
 from tqdm.auto import tqdm
 import torch
 import torch.nn as nn
@@ -23,6 +25,9 @@ from diffusers import AutoencoderKL
 from omegaconf import OmegaConf
 from accelerate.utils import set_seed
 
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import MODELS_DIR
+
 
 def main(config):
     set_seed(config.run.seed)
@@ -31,13 +36,19 @@ def main(config):
 
     if config.data.latent_space:
         vae = AutoencoderKL.from_pretrained(
-            "runwayml/stable-diffusion-inpainting", subfolder="vae", revision="fp16", torch_dtype=torch.float16
+            "runwayml/stable-diffusion-inpainting",
+            subfolder="vae",
+            revision="fp16",
+            torch_dtype=torch.float16,
+            cache_dir=MODELS_DIR,
         )
         vae.requires_grad_(False)
         vae.to(device)
 
     # Dataset and Dataloader setup
-    dataset = SyncNetDataset(config.data.val_data_dir, config.data.val_fileslist, config)
+    dataset = SyncNetDataset(
+        config.data.val_data_dir, config.data.val_fileslist, config
+    )
 
     test_dataloader = torch.utils.data.DataLoader(
         dataset,
@@ -52,7 +63,9 @@ def main(config):
     syncnet = StableSyncNet(OmegaConf.to_container(config.model)).to(device)
 
     print(f"Load checkpoint from: {config.ckpt.inference_ckpt_path}")
-    checkpoint = torch.load(config.ckpt.inference_ckpt_path, map_location=device, weights_only=True)
+    checkpoint = torch.load(
+        config.ckpt.inference_ckpt_path, map_location=device, weights_only=True
+    )
 
     syncnet.load_state_dict(checkpoint["state_dict"])
     syncnet.to(dtype=torch.float16)
@@ -80,7 +93,9 @@ def main(config):
                 with torch.no_grad():
                     frames = vae.encode(frames).latent_dist.sample() * 0.18215
 
-                frames = rearrange(frames, "(b f) c h w -> b (f c) h w", f=config.data.num_frames)
+                frames = rearrange(
+                    frames, "(b f) c h w -> b (f c) h w", f=config.data.num_frames
+                )
             else:
                 frames = rearrange(frames, "b f c h w -> b (f c) h w")
 
@@ -102,14 +117,18 @@ def main(config):
 
             if global_step >= num_val_batches:
                 progress_bar.close()
-                print(f"SyncNet Accuracy: {num_correct_preds / num_total_preds*100:.2f}%")
+                print(
+                    f"SyncNet Accuracy: {num_correct_preds / num_total_preds * 100:.2f}%"
+                )
                 return
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Code to test the accuracy of SyncNet")
 
-    parser.add_argument("--config_path", type=str, default="configs/syncnet/syncnet_16_latent.yaml")
+    parser.add_argument(
+        "--config_path", type=str, default="configs/syncnet/syncnet_16_latent.yaml"
+    )
     args = parser.parse_args()
 
     # Load a configuration file

latentsync/utils/util.py

@@ -49,9 +49,7 @@ def read_video(video_path: str, change_fps=True, use_decord=True):
         if os.path.exists(temp_dir):
             shutil.rmtree(temp_dir)
         os.makedirs(temp_dir, exist_ok=True)
-        command = (
-            f"ffmpeg -loglevel error -y -nostdin -i {video_path} -r 25 -crf 18 {os.path.join(temp_dir, 'video.mp4')}"
-        )
+        command = f"ffmpeg -loglevel error -y -nostdin -i {video_path} -r 25 -crf 18 {os.path.join(temp_dir, 'video.mp4')}"
         subprocess.run(command, shell=True)
         target_video_path = os.path.join(temp_dir, "video.mp4")
     else:
@@ -127,7 +125,9 @@ def write_video(video_output_path: str, video_frames: np.ndarray, fps: int):
 
 def write_video_cv2(video_output_path: str, video_frames: np.ndarray, fps: int):
     height, width = video_frames[0].shape[:2]
-    out = cv2.VideoWriter(video_output_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
+    out = cv2.VideoWriter(
+        video_output_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)
+    )
     # out = cv2.VideoWriter(video_output_path, cv2.VideoWriter_fourcc(*"vp09"), fps, (width, height))
     for frame in video_frames:
         frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
@@ -162,7 +162,9 @@ def check_video_fps(video_path: str):
     cam = cv2.VideoCapture(video_path)
     fps = cam.get(cv2.CAP_PROP_FPS)
     if fps != 25:
-        raise ValueError(f"Video FPS is not 25, it is {fps}. Please convert the video to 25 FPS.")
+        raise ValueError(
+            f"Video FPS is not 25, it is {fps}. Please convert the video to 25 FPS."
+        )
 
 
 def one_step_sampling(ddim_scheduler, pred_noise, timesteps, x_t):
@@ -175,7 +177,9 @@ def one_step_sampling(ddim_scheduler, pred_noise, timesteps, x_t):
     if ddim_scheduler.config.prediction_type == "epsilon":
         beta_prod_t = beta_prod_t[:, None, None, None, None]
         alpha_prod_t = alpha_prod_t[:, None, None, None, None]
-        pred_original_sample = (x_t - beta_prod_t ** (0.5) * pred_noise) / alpha_prod_t ** (0.5)
+        pred_original_sample = (
+            x_t - beta_prod_t ** (0.5) * pred_noise
+        ) / alpha_prod_t ** (0.5)
     else:
         raise NotImplementedError("This prediction type is not implemented yet")
 
@@ -269,16 +273,25 @@ def count_video_time(video_path):
 
 def check_ffmpeg_installed():
     # Run the ffmpeg command with the -version argument to check if it's installed
-    result = subprocess.run("ffmpeg -version", stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
+    result = subprocess.run(
+        "ffmpeg -version", stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True
+    )
     if not result.returncode == 0:
-        raise FileNotFoundError("ffmpeg not found, please install it by:\n    $ conda install -c conda-forge ffmpeg")
+        raise FileNotFoundError(
+            "ffmpeg not found, please install it by:\n    $ conda install -c conda-forge ffmpeg"
+        )
 
 
-def check_model_and_download(ckpt_path: str, huggingface_model_id: str = "ByteDance/LatentSync-1.5"):
+def check_model_and_download(
+    ckpt_path: str, huggingface_model_id: str = "ByteDance/LatentSync-1.5"
+):
     if not os.path.exists(ckpt_path):
         ckpt_path_obj = Path(ckpt_path)
         download_cmd = f"huggingface-cli download {huggingface_model_id} {Path(*ckpt_path_obj.parts[1:])} --local-dir {Path(ckpt_path_obj.parts[0])}"
         subprocess.run(download_cmd, shell=True)
+        print(f"Downloaded model to {ckpt_path}")
+    else:
+        print(f"Model already exists: {ckpt_path}")
 
 
 class dummy_context:

latentsync/whisper/audio2feature.py

@@ -15,8 +15,9 @@ class Audio2Feature:
         audio_embeds_cache_dir=None,
         num_frames=16,
         audio_feat_length=[2, 2],
+        download_root=None,
     ):
-        self.model = load_model(model_path, device)
+        self.model = load_model(model_path, device, download_root=download_root)
         self.audio_embeds_cache_dir = audio_embeds_cache_dir
         if audio_embeds_cache_dir is not None and audio_embeds_cache_dir != "":
             Path(audio_embeds_cache_dir).mkdir(parents=True, exist_ok=True)

lipsync.py

@@ -8,11 +8,12 @@ import torch
 from DeepCache import DeepCacheSDHelper
 from latentsync.pipelines.lipsync_pipeline import LipsyncPipeline
 from shared.model_manager import ModelManager
+from config import MODELS_DIR
 
 torch.backends.cudnn.benchmark = True
 torch.backends.cudnn.deterministic = False
 
-os.makedirs("checkpoints", exist_ok=True)
+os.makedirs(MODELS_DIR, exist_ok=True)
 
 
 def get_quality_params(level: str) -> tuple:

preprocess/data_processing_pipeline.py

@@ -14,6 +14,7 @@
 
 import argparse
 import os
+import sys
 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
@@ -25,14 +26,22 @@ from preprocess.filter_visual_quality import filter_visual_quality_multi_gpus
 from preprocess.remove_incorrect_affined import remove_incorrect_affined_multiprocessing
 from latentsync.utils.util import check_model_and_download
 
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import MODELS_DIR
+
 
 def data_processing_pipeline(
-    total_num_workers, per_gpu_num_workers, resolution, sync_conf_threshold, temp_dir, input_dir
+    total_num_workers,
+    per_gpu_num_workers,
+    resolution,
+    sync_conf_threshold,
+    temp_dir,
+    input_dir,
 ):
     print("Checking models are downloaded...")
-    check_model_and_download("checkpoints/auxiliary/syncnet_v2.model")
-    check_model_and_download("checkpoints/auxiliary/sfd_face.pth")
-    check_model_and_download("checkpoints/auxiliary/koniq_pretrained.pkl")
+    check_model_and_download(f"{MODELS_DIR}/auxiliary/syncnet_v2.model")
+    check_model_and_download(f"{MODELS_DIR}/auxiliary/sfd_face.pth")
+    check_model_and_download(f"{MODELS_DIR}/auxiliary/koniq_pretrained.pkl")
 
     print("Removing broken videos...")
     remove_broken_videos_multiprocessing(input_dir, total_num_workers)
@@ -55,19 +64,39 @@ def data_processing_pipeline(
     # 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(segmented_dir, affine_transformed_dir, temp_dir, resolution, per_gpu_num_workers // 2)
+    affine_transformed_dir = os.path.join(
+        os.path.dirname(input_dir), "affine_transformed"
+    )
+    affine_transform_multi_gpus(
+        segmented_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)
+    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)
+    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__":

preprocess/filter_visual_quality.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import sys
 import tqdm
 import torch
 import torchvision
@@ -23,6 +24,9 @@ from decord import VideoReader
 from einops import rearrange
 from eval.hyper_iqa import HyperNet, TargetNet
 
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+from config import MODELS_DIR
+
 
 paths = []
 
@@ -38,7 +42,9 @@ def gather_paths(input_dir, output_dir):
                 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))
+            gather_paths(
+                os.path.join(input_dir, video), os.path.join(output_dir, video)
+            )
 
 
 def read_video(video_path: str):
@@ -61,13 +67,17 @@ def func(paths, device_id):
 
     # load the pre-trained model on the koniq-10k dataset
     model_hyper.load_state_dict(
-        (torch.load("checkpoints/auxiliary/koniq_pretrained.pkl", map_location=device, weights_only=True))
+        (torch.load(f"{MODELS_DIR}/auxiliary/koniq_pretrained.pkl", map_location=device, weights_only=True))
+    )
+        )
     )
 
     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)),
+            torchvision.transforms.Normalize(
+                mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)
+            ),
         ]
     )
 
@@ -76,7 +86,9 @@ def func(paths, device_id):
             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
+            paras = model_hyper(
+                video_frames
+            )  # 'paras' contains the network weights conveyed to target network
 
             # Building target network
             model_target = TargetNet(paras).to(device)
@@ -84,11 +96,15 @@ def func(paths, device_id):
                 param.requires_grad = False
 
             # Quality prediction
-            pred = model_target(paras["target_in_vec"])  # 'paras['target_in_vec']' is the input to target net
+            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}")
+            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)

preprocess/sync_av.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import sys
 import tqdm
 from eval.syncnet import SyncNetEval
 from eval.syncnet_detect import SyncNetDetector
@@ -22,6 +23,10 @@ import subprocess
 import shutil
 from multiprocessing import Process
 
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import MODELS_DIR
+
+
 paths = []
 
 
@@ -36,11 +41,13 @@ def gather_paths(input_dir, output_dir):
                 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))
+            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}"
+    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)
 
 
@@ -49,17 +56,23 @@ def func(sync_conf_threshold, paths, device_id, process_temp_dir):
     device = f"cuda:{device_id}"
 
     syncnet = SyncNetEval(device=device)
-    syncnet.loadParameters("checkpoints/auxiliary/syncnet_v2.model")
+    syncnet.loadParameters(f"{MODELS_DIR}/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)
+    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
+                syncnet,
+                syncnet_detector,
+                video_input,
+                syncnet_eval_results_dir,
+                detect_results_dir,
             )
 
             if conf >= sync_conf_threshold and abs(av_offset) <= 6:
@@ -77,7 +90,9 @@ def split(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):
+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:
@@ -111,4 +126,6 @@ if __name__ == "__main__":
     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)
+    sync_av_multi_gpus(
+        input_dir, output_dir, temp_dir, num_workers, sync_conf_threshold
+    )

scripts/inference.py

@@ -14,10 +14,14 @@
 
 import argparse
 import os
+import sys
 from omegaconf import OmegaConf
 import torch
 from diffusers import AutoencoderKL, DDIMScheduler
 from latentsync.models.unet import UNet3DConditionModel
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import MODELS_DIR
 from latentsync.pipelines.lipsync_pipeline import LipsyncPipeline
 from accelerate.utils import set_seed
 from latentsync.whisper.audio2feature import Audio2Feature
@@ -56,7 +60,9 @@ def main(config, args):
         audio_feat_length=config.data.audio_feat_length,
     )
 
-    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=dtype)
+    vae = AutoencoderKL.from_pretrained(
+        "stabilityai/sd-vae-ft-mse", torch_dtype=dtype, cache_dir=MODELS_DIR
+    )
     vae.config.scaling_factor = 0.18215
     vae.config.shift_factor = 0
 

scripts/train_syncnet.py

@@ -13,11 +13,18 @@
 # limitations under the License.
 
 from tqdm.auto import tqdm
-import os, argparse, datetime, math
+import os
+import sys
+import argparse
+import datetime
+import math
 import logging
 from omegaconf import OmegaConf
 import shutil
 
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import MODELS_DIR
+
 from latentsync.data.syncnet_dataset import SyncNetDataset
 from latentsync.models.stable_syncnet import StableSyncNet
 from latentsync.models.wav2lip_syncnet import Wav2LipSyncNet
@@ -67,15 +74,21 @@ def main(config):
     device = torch.device(local_rank)
 
     if config.data.latent_space:
-        vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)
+        vae = AutoencoderKL.from_pretrained(
+            "stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16, cache_dir=MODELS_DIR
+        )
         vae.requires_grad_(False)
         vae.to(device)
     else:
         vae = None
 
     # Dataset and Dataloader setup
-    train_dataset = SyncNetDataset(config.data.train_data_dir, config.data.train_fileslist, config)
-    val_dataset = SyncNetDataset(config.data.val_data_dir, config.data.val_fileslist, config)
+    train_dataset = SyncNetDataset(
+        config.data.train_data_dir, config.data.train_fileslist, config
+    )
+    val_dataset = SyncNetDataset(
+        config.data.val_data_dir, config.data.val_fileslist, config
+    )
 
     train_distributed_sampler = DistributedSampler(
         train_dataset,
@@ -118,7 +131,8 @@ def main(config):
     # syncnet = Wav2LipSyncNet().to(device)
 
     optimizer = torch.optim.AdamW(
-        list(filter(lambda p: p.requires_grad, syncnet.parameters())), lr=config.optimizer.lr
+        list(filter(lambda p: p.requires_grad, syncnet.parameters())),
+        lr=config.optimizer.lr,
     )
 
     global_step = 0
@@ -130,7 +144,9 @@ def main(config):
     if config.ckpt.resume_ckpt_path != "":
         if is_main_process:
             logger.info(f"Load checkpoint from: {config.ckpt.resume_ckpt_path}")
-        ckpt = torch.load(config.ckpt.resume_ckpt_path, map_location=device, weights_only=True)
+        ckpt = torch.load(
+            config.ckpt.resume_ckpt_path, map_location=device, weights_only=True
+        )
 
         syncnet.load_state_dict(ckpt["state_dict"])
 
@@ -145,7 +161,9 @@ def main(config):
     syncnet = DDP(syncnet, device_ids=[local_rank], output_device=local_rank)
 
     num_update_steps_per_epoch = math.ceil(len(train_dataloader))
-    num_train_epochs = math.ceil(config.run.max_train_steps / num_update_steps_per_epoch)
+    num_train_epochs = math.ceil(
+        config.run.max_train_steps / num_update_steps_per_epoch
+    )
 
     if is_main_process:
         logger.info("***** Running training *****")
@@ -158,15 +176,22 @@ def main(config):
         logger.info(f"  Total optimization steps = {config.run.max_train_steps}")
 
     first_epoch = global_step // num_update_steps_per_epoch
-    num_val_batches = config.data.num_val_samples // (num_processes * config.data.batch_size)
+    num_val_batches = config.data.num_val_samples // (
+        num_processes * config.data.batch_size
+    )
 
     # Only show the progress bar once on each machine.
     progress_bar = tqdm(
-        range(0, config.run.max_train_steps), initial=global_step, desc="Steps", disable=not is_main_process
+        range(0, config.run.max_train_steps),
+        initial=global_step,
+        desc="Steps",
+        disable=not is_main_process,
     )
 
     # Support mixed-precision training
-    scaler = torch.amp.GradScaler("cuda") if config.run.mixed_precision_training else None
+    scaler = (
+        torch.amp.GradScaler("cuda") if config.run.mixed_precision_training else None
+    )
 
     for epoch in range(first_epoch, num_train_epochs):
         train_dataloader.sampler.set_epoch(epoch)
@@ -186,15 +211,17 @@ def main(config):
                     num_samples_limit // config.data.num_frames
                 )  # due to the limited cuda memory, we split the input frames into parts
                 if frames.shape[0] > max_batch_size:
-                    assert (
-                        frames.shape[0] % max_batch_size == 0
-                    ), f"max_batch_size {max_batch_size} should be divisible by batch_size {frames.shape[0]}"
+                    assert frames.shape[0] % max_batch_size == 0, (
+                        f"max_batch_size {max_batch_size} should be divisible by batch_size {frames.shape[0]}"
+                    )
                     frames_part_results = []
                     for i in range(0, frames.shape[0], max_batch_size):
                         frames_part = frames[i : i + max_batch_size]
                         frames_part = rearrange(frames_part, "b f c h w -> (b f) c h w")
                         with torch.no_grad():
-                            frames_part = vae.encode(frames_part).latent_dist.sample() * 0.18215
+                            frames_part = (
+                                vae.encode(frames_part).latent_dist.sample() * 0.18215
+                            )
                         frames_part_results.append(frames_part)
                     frames = torch.cat(frames_part_results, dim=0)
                 else:
@@ -202,7 +229,9 @@ def main(config):
                     with torch.no_grad():
                         frames = vae.encode(frames).latent_dist.sample() * 0.18215
 
-                frames = rearrange(frames, "(b f) c h w -> b (f c) h w", f=config.data.num_frames)
+                frames = rearrange(
+                    frames, "(b f) c h w -> b (f c) h w", f=config.data.num_frames
+                )
             else:
                 frames = rearrange(frames, "b f c h w -> b (f c) h w")
 
@@ -211,14 +240,22 @@ def main(config):
                 frames = frames[:, :, height // 2 :, :]
 
             # Disable gradient sync for the first N-1 steps, enable sync on the final step
-            with syncnet.no_sync() if (index + 1) % config.data.gradient_accumulation_steps != 0 else dummy_context():
+            with (
+                syncnet.no_sync()
+                if (index + 1) % config.data.gradient_accumulation_steps != 0
+                else dummy_context()
+            ):
                 # Mixed-precision training
                 with torch.autocast(
-                    device_type="cuda", dtype=torch.float16, enabled=config.run.mixed_precision_training
+                    device_type="cuda",
+                    dtype=torch.float16,
+                    enabled=config.run.mixed_precision_training,
                 ):
                     vision_embeds, audio_embeds = syncnet(frames, audio_samples)
 
-                loss = cosine_loss(vision_embeds.float(), audio_embeds.float(), y).mean()
+                loss = cosine_loss(
+                    vision_embeds.float(), audio_embeds.float(), y
+                ).mean()
                 loss = loss / config.data.gradient_accumulation_steps
 
                 # Backpropagate
@@ -230,7 +267,9 @@ def main(config):
             if (index + 1) % config.data.gradient_accumulation_steps == 0:
                 """>>> gradient clipping >>>"""
                 scaler.unscale_(optimizer)
-                torch.nn.utils.clip_grad_norm_(syncnet.parameters(), config.optimizer.max_grad_norm)
+                torch.nn.utils.clip_grad_norm_(
+                    syncnet.parameters(), config.optimizer.max_grad_norm
+                )
                 """ <<< gradient clipping <<< """
                 scaler.step(optimizer)
                 scaler.update()
@@ -255,15 +294,21 @@ def main(config):
                     )
                     val_step_list.append(global_step)
                     val_loss_list.append(val_loss)
-                    logger.info(f"Validation loss at step {global_step} is {val_loss:0.3f}")
+                    logger.info(
+                        f"Validation loss at step {global_step} is {val_loss:0.3f}"
+                    )
                     plot_loss_chart(
-                        os.path.join(output_dir, f"loss_charts/loss_chart-{global_step}.png"),
+                        os.path.join(
+                            output_dir, f"loss_charts/loss_chart-{global_step}.png"
+                        ),
                         ("Train loss", train_step_list, train_loss_list),
                         ("Val loss", val_step_list, val_loss_list),
                     )
 
                 if is_main_process and global_step % config.ckpt.save_ckpt_steps == 0:
-                    checkpoint_save_path = os.path.join(output_dir, f"checkpoints/checkpoint-{global_step}.pt")
+                    checkpoint_save_path = os.path.join(
+                        output_dir, f"checkpoints/checkpoint-{global_step}.pt"
+                    )
                     torch.save(
                         {
                             "state_dict": syncnet.module.state_dict(),  # to unwrap DDP
@@ -288,7 +333,9 @@ def main(config):
 
 
 @torch.no_grad()
-def validation(val_dataloader, device, syncnet, latent_space, lower_half, vae, num_val_batches):
+def validation(
+    val_dataloader, device, syncnet, latent_space, lower_half, vae, num_val_batches
+):
     syncnet.eval()
 
     losses = []
@@ -330,7 +377,9 @@ def validation(val_dataloader, device, syncnet, latent_space, lower_half, vae, n
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Code to train the SyncNet")
-    parser.add_argument("--config_path", type=str, default="configs/syncnet/syncnet_16_pixel.yaml")
+    parser.add_argument(
+        "--config_path", type=str, default="configs/syncnet/syncnet_16_pixel.yaml"
+    )
     args = parser.parse_args()
 
     # Load a configuration file

scripts/train_unet.py

@@ -13,11 +13,15 @@
 # limitations under the License.
 
 import os
+import sys
 import math
 import argparse
 import shutil
 import datetime
 import logging
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from config import MODELS_DIR
 from omegaconf import OmegaConf
 
 from tqdm.auto import tqdm
@@ -93,7 +97,7 @@ def main(config):
     noise_scheduler = DDIMScheduler.from_pretrained("configs")
 
     vae = AutoencoderKL.from_pretrained(
-        "stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16
+        "stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16, cache_dir=MODELS_DIR
     )
     vae.config.scaling_factor = 0.18215
     vae.config.shift_factor = 0

shared/face_detection/detector.py

@@ -1,6 +1,11 @@
+import os
+import sys
 import numpy as np
 import torch
 
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+from config import MODELS_DIR
+
 INSIGHTFACE_DETECT_SIZE = 512
 
 
@@ -56,7 +61,7 @@ class FaceDetector:
 
         self.app = FaceAnalysis(
             allowed_modules=["detection", "landmark_2d_106"],
-            root="checkpoints/auxiliary",
+            root=f"{MODELS_DIR}/auxiliary",
             providers=["CUDAExecutionProvider"],
         )
         self.app.prepare(

shared/model_manager.py

@@ -39,10 +39,14 @@ class ModelManager:
         """Load Whisper audio encoder (lazy loaded)"""
         if self._whisper_encoder is None:
             from latentsync.whisper.audio2feature import Audio2Feature
+            from config import MODELS_DIR
 
             logger.info(f"Loading Whisper encoder from {model_path}...")
             self._whisper_encoder = Audio2Feature(
-                model_path=model_path, device=device, num_frames=num_frames
+                model_path=model_path,
+                device=device,
+                num_frames=num_frames,
+                download_root=f"{MODELS_DIR}/whisper",
             )
             logger.info("Whisper encoder loaded")
         return self._whisper_encoder
@@ -53,8 +57,12 @@ class ModelManager:
             from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
 
             logger.info("Loading VAE...")
+            from config import MODELS_DIR
+
             vae = AutoencoderKL.from_pretrained(
-                "stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16
+                "stabilityai/sd-vae-ft-mse",
+                torch_dtype=torch.float16,
+                cache_dir=MODELS_DIR,
             )
             vae.config.scaling_factor = 0.18215
             vae.config.shift_factor = 0
@@ -74,20 +82,23 @@ class ModelManager:
         """Load LatentSync UNet (lazy loaded)"""
         if self._latentsync_unet is None:
             from latentsync.models.unet import UNet3DConditionModel
+            from config import MODELS_DIR
 
-            logger.info("Downloading LatentSync-1.6 models...")
-            os.makedirs("checkpoints", exist_ok=True)
-            snapshot_download(
-                repo_id="ByteDance/LatentSync-1.6", local_dir="./checkpoints"
-            )
+            unet_path = f"{MODELS_DIR}/latentsync_unet.pt"
+
+            if not os.path.exists(unet_path):
+                logger.info("Downloading LatentSync-1.6 models...")
+                os.makedirs(MODELS_DIR, exist_ok=True)
+                snapshot_download(
+                    repo_id="ByteDance/LatentSync-1.6", local_dir=MODELS_DIR
+                )
 
             logger.info("Loading LatentSync UNet...")
             config = self.get_latentsync_config()
 
-            inference_ckpt_path = "checkpoints/latentsync_unet.pt"
             unet, _ = UNet3DConditionModel.from_pretrained(
                 OmegaConf.to_container(config.model),
-                inference_ckpt_path,
+                unet_path,
                 device="cpu",
             )
             unet = unet.to(dtype=torch.float16).to(device)

shared/vae/loader.py

@@ -1,8 +1,13 @@
 """VAE loader (placeholder - actual loading handled by ModelManager)"""
 
+import os
+import sys
 import torch
 from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
 
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+from config import MODELS_DIR
+
 
 def load_vae(device: str = "cuda"):
     """Load VAE from HuggingFace
@@ -14,7 +19,7 @@ def load_vae(device: str = "cuda"):
         VAE model
     """
     vae = AutoencoderKL.from_pretrained(
-        "stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16
+        "stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16, cache_dir=MODELS_DIR
     )
     vae.config.scaling_factor = 0.18215
     vae.config.shift_factor = 0