#!/usr/bin/env python3 """ Retrain the LIPSYNC core's alignment model (avh-align). The lipsync core (avh-align_core) detects audio/visual *misalignment*. Its ``FusionModel`` is trained **self-supervised on REAL data only**: for every video frame it must place its highest synchronization score on the temporally-aligned audio window (offset 0). Fakes break that alignment at inference time. So retraining needs only real videos and no labels. Pipeline (mirrors the core's own train pipeline, reusing the core's code): Phase A preprocess (dlib mouth ROI + wav) and extract AVHuBERT per-frame features -> /feat_.npz with keys 'visual','audio' + an internal metadata CSV (path,num_frames) that ``avh-align_core/dataset.py:FeatureDataset`` understands. Phase B train ``avh-align_core/model.py:FusionModel`` with the alignment objective from ``avh-align_core/train.py``: sync = LogSoftmax(outputs)[:, tau]; loss = -sum(sync) Metadata -------- CSV with a column ``file_path`` (or ``path``) listing REAL video files. No label. Output ------ ``retrained_models/lipsync/fusion_retrained.pt`` — a checkpoint dict ``{"state_dict": ...}`` (what ``avh_align.py:load_models`` loads). Use it: python detect.py INPUT --lipsync-model retrained_models/lipsync/fusion_retrained.pt Never writes inside the core directory. Auto-relaunches in the `avh` conda env. """ import os import sys # --------------------------------------------------------------------------- # # CONFIG # --------------------------------------------------------------------------- # PROJECT_DIR = os.path.dirname(os.path.abspath(__file__)) CORE_DIR = os.path.join(PROJECT_DIR, "avh-align_core") LIPSYNC_PY = "/opt/conda/envs/avh/bin/python" DEFAULT_OUT = os.path.join(PROJECT_DIR, "retrained_models", "lipsync", "fusion_retrained.pt") DEFAULT_FEATURES = os.path.join(PROJECT_DIR, "retrained_models", "lipsync", "features") _RELAUNCH_FLAG = "_DF_LIPSYNC_RETRAIN_RELAUNCHED" def _ensure_env(): if os.environ.get(_RELAUNCH_FLAG) == "1": return if os.path.exists(LIPSYNC_PY) and os.path.abspath(sys.executable) != os.path.abspath(LIPSYNC_PY): env = dict(os.environ) env[_RELAUNCH_FLAG] = "1" os.execve(LIPSYNC_PY, [LIPSYNC_PY] + sys.argv, env) os.environ[_RELAUNCH_FLAG] = "1" def _load_core_namespace(): """Exec avh_align.py truncated before its self-test; this also applies the required np.float monkeypatch and loads AVHuBERT + dlib models.""" os.chdir(CORE_DIR) sys.path.insert(0, CORE_DIR) with open(os.path.join(CORE_DIR, "avh_align.py")) as fh: src = fh.read().split("#### testing")[0] # drop trailing self-test block ns = {"__name__": "__lipsync_core__", "__file__": os.path.join(CORE_DIR, "avh_align.py")} exec(compile(src, "avh_align.py", "exec"), ns) ns["load_models"]("model") return ns def extract_phase(ns, df, path_col, features_dir, skip_existing): """Phase A: build per-video .npz features + metadata rows.""" import numpy as np preproc_dir = os.path.join(features_dir, "_preproc") os.makedirs(features_dir, exist_ok=True) os.makedirs(preproc_dir, exist_ok=True) preprocess_video = ns["preprocess_video"] extract_feats = ns["extract_feats"] model, task = ns["model"], ns["task"] rows = [] for i, (_, row) in enumerate(df.iterrows()): fpath = str(row[path_col]) if not os.path.isabs(fpath): fpath = os.path.join(os.path.dirname(df.attrs["meta_path"]), fpath) stem = "feat_%05d" % i npz_path = os.path.join(features_dir, stem + ".npz") if skip_existing and os.path.exists(npz_path): try: nf = int(np.load(npz_path)["visual"].shape[0]) rows.append({"path": stem + ".mp4", "num_frames": nf}) print(" [cached] %s (%d frames)" % (os.path.basename(fpath), nf)) continue except Exception: pass if not os.path.exists(fpath): print(" [skip] missing file: %s" % fpath) continue try: out_dir = os.path.join(preproc_dir, stem) os.makedirs(out_dir, exist_ok=True) pre = preprocess_video(fpath, out_dir) if not pre: print(" [skip] preprocess failed: %s" % fpath) continue roi_path, audio_fn = pre f_audio, f_video, _ = extract_feats(model, task, roi_path, audio_fn) n = min(len(f_video), len(f_audio)) if n < 1: print(" [skip] no frames: %s" % fpath) continue visual = np.asarray(f_video[:n]) audio = np.asarray(f_audio[:n]) np.savez(npz_path, visual=visual, audio=audio) rows.append({"path": stem + ".mp4", "num_frames": int(n)}) print(" [ok] %s -> %s (%d frames)" % (os.path.basename(fpath), stem, n)) except Exception as exc: print(" [err] %s: %s" % (fpath, exc)) return rows def train_phase(ns, train_csv, features_dir, args, device): """Phase B: alignment training (replicates avh-align_core/train.py).""" import importlib import torch import torch.nn as nn from torch.utils.data import DataLoader FeatureDataset = importlib.import_module("dataset").FeatureDataset FusionModel = ns["FusionModel"] dataset = FeatureDataset(train_csv, features_dir, tau=args.tau) num_videos = dataset.num_videos num_workers = max(0, min(args.num_workers, num_videos)) loader = DataLoader(dataset, batch_size=args.batch_size, num_workers=num_workers, pin_memory=True) model = FusionModel().to(device) model.device = device optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) logsoftmax = nn.LogSoftmax(dim=1) print("\nTraining FusionModel | videos=%d total_frames=%d tau=%d" % (num_videos, len(dataset), args.tau)) for epoch in range(args.epochs): model.train() total_loss, total = 0.0, 0 for step, batch in enumerate(loader): visual_frame, audio_window, _, _ = batch bs = visual_frame.size(0) visual_frame = visual_frame.to(device) audio_window = audio_window.to(device) # repeat the central visual frame across the (2*tau+1) audio window visual_central = visual_frame.unsqueeze(1).repeat(1, 2 * args.tau + 1, 1) outputs = model(visual_central, audio_window).squeeze(-1) # [B, 2*tau+1] sync = logsoftmax(outputs)[:, args.tau] loss = -torch.sum(sync) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() total += bs if args.log_interval and step % args.log_interval == 0 and step > 0: print(" epoch %d step %d avg_loss=%.6f" % (epoch + 1, step, total_loss / max(total, 1))) print("epoch %3d/%d loss=%.6f" % (epoch + 1, args.epochs, total_loss / max(total, 1))) return model def main(): import argparse parser = argparse.ArgumentParser(description="Retrain the lipsync alignment FusionModel on real videos.") parser.add_argument("--metadata", required=True, help="CSV with column file_path (or path); REAL videos only") parser.add_argument("--out", default=DEFAULT_OUT, help="output checkpoint .pt path") parser.add_argument("--features-dir", default=DEFAULT_FEATURES, help="where to cache extracted .npz features") parser.add_argument("--skip-existing", action="store_true", help="reuse cached .npz features when present") parser.add_argument("--extract-only", action="store_true", help="run Phase A only (no training)") parser.add_argument("--tau", type=int, default=15, help="audio window half-width (matches core default)") parser.add_argument("--epochs", type=int, default=10) parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--batch-size", type=int, default=256) parser.add_argument("--num-workers", type=int, default=4) parser.add_argument("--log-interval", type=int, default=200) args = parser.parse_args() import pandas as pd import torch metadata_path = os.path.abspath(args.metadata) out_path = os.path.abspath(args.out) features_dir = os.path.abspath(args.features_dir) df = pd.read_csv(metadata_path) df.attrs["meta_path"] = metadata_path path_col = "file_path" if "file_path" in df.columns else ("path" if "path" in df.columns else None) if path_col is None: sys.exit("metadata must have a column: file_path (or path)") print("Loading lipsync core (AVHuBERT + dlib) ...") ns = _load_core_namespace() print("\n[Phase A] feature extraction ...") rows = extract_phase(ns, df, path_col, features_dir, args.skip_existing) if not rows: sys.exit("No features extracted; check metadata paths.") train_csv = os.path.join(features_dir, "train_metadata.csv") pd.DataFrame(rows, columns=["path", "num_frames"]).to_csv(train_csv, index=False) print("Wrote metadata: %s (%d videos)" % (train_csv, len(rows))) if args.extract_only: print("\n--extract-only set; skipping training.") return device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print("\n[Phase B] alignment training on %s ..." % device) model = train_phase(ns, train_csv, features_dir, args, device) os.makedirs(os.path.dirname(out_path), exist_ok=True) torch.save({"state_dict": model.state_dict(), "tau": args.tau, "epochs": args.epochs}, out_path) print("\n=== lipsync retraining done ===") print("saved : %s" % out_path) print("\nUse it: python detect.py INPUT --lipsync-model %s" % out_path) if __name__ == "__main__": _ensure_env() main()