feat: add config-driven extraction and audio sync tooling

sync_audio.py

@@ -0,0 +1,217 @@
+#!/usr/bin/env python3
+"""
+Audio cross-correlation for movie_sync — torch GPU auto (ROCm/CUDA) with CPU fallback.
+
+Content-to-content sync (Option A): ignores container start_time/edit-list mux delay.
+If audio is delayed in container, we pad with silence on mux instead of shifting presentation.
+Extracted wavs are pcm_f32le 48k mono (extract.py:41), so correlation is phase-accurate.
+
+Usage:
+  from sync_audio import load_wav, correlate, find_lag
+  a, sr = load_wav("extracted/movie_b_audio.wav")
+  lag_samples, peak, corr = correlate(a, b, sr=48000, method="phat", device="auto")
+"""
+
+from __future__ import annotations
+import os
+import sys
+import math
+from typing import Tuple, Literal
+
+import numpy as np
+
+try:
+    import soundfile as sf  # preferred for f32 wav
+    HAS_SF = True
+except ImportError:
+    HAS_SF = False
+
+try:
+    import torch
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+
+
+def _get_device(pref: str = "auto") -> str:
+    if pref == "cpu":
+        return "cpu"
+    if pref in ("cuda", "auto") and HAS_TORCH and torch.cuda.is_available():
+        return "cuda"
+    return "cpu"
+
+
+def load_wav(path: str, target_sr: int = 48000, mono: bool = True) -> Tuple[np.ndarray, int]:
+    """Load wav as float32 mono. Resample if needed (cheap linear warn)."""
+    if not os.path.exists(path):
+        raise FileNotFoundError(path)
+    if HAS_SF:
+        data, sr = sf.read(path, always_2d=False, dtype="float32")
+    else:
+        # fallback scipy
+        from scipy.io import wavfile
+        sr, data = wavfile.read(path)
+        if data.dtype == np.int16:
+            data = data.astype(np.float32) / 32768.0
+        elif data.dtype == np.int32:
+            data = data.astype(np.float32) / 2147483648.0
+        elif data.dtype != np.float32:
+            data = data.astype(np.float32)
+    if data.ndim > 1:
+        if mono:
+            data = data.mean(axis=1).astype(np.float32)
+    if sr != target_sr:
+        # simple resample via scipy if available, else warn
+        try:
+            from scipy.signal import resample_poly
+            import math as _m
+            g = _m.gcd(sr, target_sr)
+            data = resample_poly(data, target_sr // g, sr // g).astype(np.float32)
+            sr = target_sr
+        except Exception:
+            print(f"WARNING: {path} sr={sr} != {target_sr} and resample failed, using as-is", file=sys.stderr)
+    return data.astype(np.float32), sr
+
+
+def _next_pow2(n: int) -> int:
+    return 1 << (n - 1).bit_length()
+
+
+def _normalize(a: np.ndarray) -> np.ndarray:
+    a = a - a.mean()
+    s = a.std()
+    return a / s if s > 1e-9 else a
+
+
+def correlate(
+    a: np.ndarray,
+    b: np.ndarray,
+    sr: int = 48000,
+    method: Literal["phat", "standard"] = "phat",
+    device: str = "auto",
+    normalize: bool = True,
+) -> Tuple[int, float, np.ndarray]:
+    """
+    Return (lag_samples, peak_value, corr_array).
+    lag >0 means b delayed vs a (b starts later), need to shift b earlier by lag.
+    Convention: corr = ifft(FFT(a) * conj(FFT(b))). peak at lag = argmax.
+    GPU via torch if available, else numpy.
+    """
+    if normalize:
+        a = _normalize(a)
+        b = _normalize(b)
+
+    n = len(a) + len(b) - 1
+    N = _next_pow2(n)
+
+    dev = _get_device(device)
+
+    if dev == "cuda" and HAS_TORCH:
+        # torch path — ROCm and CUDA both via cuda device
+        try:
+            ta = torch.from_numpy(a).to("cuda", dtype=torch.float32)
+            tb = torch.from_numpy(b).to("cuda", dtype=torch.float32)
+            Fa = torch.fft.rfft(ta, n=N)
+            Fb = torch.fft.rfft(tb, n=N)
+            R = Fa * torch.conj(Fb)
+            if method == "phat":
+                R = R / (torch.abs(R) + 1e-10)
+            corr_t = torch.fft.irfft(R, n=N)
+            # fftshift to center 0 lag
+            corr_t = torch.roll(corr_t, shifts=N // 2)
+            corr = corr_t.cpu().numpy()
+            # keep only valid length n centered
+            start = N // 2 - (len(b) - 1)
+            # Actually after roll, we want slice; simpler: use raw without roll and map lag later
+            # Recompute without roll for correct indexing, keep roll for display but fix lag calc
+            # For lag calc we need not shift: peak at N - lag etc. So we use shifted version + offset mapping
+            # To avoid confusion, return shifted corr and map lag accordingly
+            peak_idx = int(np.argmax(np.abs(corr)))
+            # lag = peak_idx - (len(b)-1) - (N//2 - (len(b)-1))? Let's instead use unshifted for lag and shifted for display
+            # Fallback: compute lag from unshifted
+            corr_raw = torch.fft.irfft(R, n=N).cpu().numpy()
+            peak_raw = int(np.argmax(np.abs(corr_raw)))
+            # raw layout: lag 0 at 0, positive lag (b delayed) wraps to end
+            if peak_raw > N // 2:
+                lag = peak_raw - N
+            else:
+                lag = peak_raw
+            peak_val = float(np.abs(corr_raw[peak_raw]))
+            # return shifted display corr + lag
+            return lag, peak_val, corr
+        except Exception as e:
+            print(f"torch cuda failed ({e}), falling back cpu", file=sys.stderr)
+            dev = "cpu"
+
+    # numpy/scipy cpu path
+    Fa = np.fft.rfft(a, n=N)
+    Fb = np.fft.rfft(b, n=N)
+    R = Fa * np.conj(Fb)
+    if method == "phat":
+        R = R / (np.abs(R) + 1e-10)
+    corr_raw = np.fft.irfft(R, n=N)
+    peak_raw = int(np.argmax(np.abs(corr_raw)))
+    if peak_raw > N // 2:
+        lag = peak_raw - N
+    else:
+        lag = peak_raw
+    peak_val = float(np.abs(corr_raw[peak_raw]))
+    # shifted for display centered at 0
+    corr_shifted = np.fft.fftshift(corr_raw)
+    # but we return raw shifted centered? Provide shifted for plotting
+    # For consistency with torch branch we return shifted display array
+    return lag, peak_val, corr_shifted
+
+
+def lag_to_seconds(lag_samples: int, sr: int = 48000) -> float:
+    return lag_samples / float(sr)
+
+
+def apply_lag_with_silence(wav: np.ndarray, lag_samples: int) -> np.ndarray:
+    """
+    Apply lag by padding silence at start (no sound for delay).
+    lag >0: b delayed vs a => pad b at start with lag silence to align to a, or trim a.
+    For mux: if lag>0 (b starts later), we need to delay a by lag or提前 b.
+    This helper pads with zeros (silence) at start; negative lag pads the other direction.
+    Caller decides which file to pad.
+    """
+    if lag_samples == 0:
+        return wav
+    if lag_samples > 0:
+        # b delayed => pad b front with lag silence to align to a
+        return np.concatenate([np.zeros(lag_samples, dtype=np.float32), wav])
+    else:
+        # b earlier => trim or pad a; this pads a-side: when applying to b, negative lag means b earlier, so we trim? Provide padded version for mux: pad with silence then trim alignment
+        # For generic helper: pad with silence at start if lag negative we pad other file, but this is applied to lagging file only
+        return np.concatenate([np.zeros(-lag_samples, dtype=np.float32), wav])
+
+
+# torch device info for Makefile / GUI
+def device_info() -> str:
+    if not HAS_TORCH:
+        return "torch not installed (cpu numpy)"
+    if torch.cuda.is_available():
+        try:
+            name = torch.cuda.get_device_name(0)
+        except Exception:
+            name = "cuda"
+        return f"torch {torch.__version__} + {name} (cuda/ROCm)"
+    return f"torch {torch.__version__} (cpu only)"
+
+
+if __name__ == "__main__":
+    import argparse, pathlib
+    ap = argparse.ArgumentParser()
+    ap.add_argument("a", help="wav A")
+    ap.add_argument("b", help="wav B")
+    ap.add_argument("--phat", action="store_true", help="use GCC-PHAT")
+    ap.add_argument("--device", default="auto", choices=["auto","cpu","cuda"])
+    ap.add_argument("--sr", type=int, default=48000)
+    args = ap.parse_args()
+    a, sa = load_wav(args.a, target_sr=args.sr)
+    b, sb = load_wav(args.b, target_sr=args.sr)
+    method = "phat" if args.phat else "standard"
+    lag, peak, _ = correlate(a, b, sr=args.sr, method=method, device=args.device)
+    print(f"device: {device_info()} method={method}")
+    print(f"lag_samples={lag} lag_sec={lag/args.sr:.6f}s peak={peak:.4f}")
+    print(f"apply: ffmpeg -i video.mkv -i audio.wav -filter_complex \"[1:a]adelay={max(0, lag/args.sr*1000):.0f}|{max(0, lag/args.sr*1000):.0f}[ad]\" -map 0:v -map \"[ad]\" out.mkv  (pad silence for delay)")