retrain_lipsync.py

#!/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  ->  <features-dir>/feat_<i>.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()