scripts/dataset/wetdry.py

"""Wet/dry mixer for a2a output — a strength dial that needs no regeneration.

Once the LoRA has applied a space you cannot subtract it, but you can mix the
original dry source back underneath. Because a2a preserves timing (measured
10-30ms lag, 0.87-0.95 envelope correlation) the two line up well enough to sum.

Sub-sample offsets would comb-filter, so the wet track is first aligned to the
dry by envelope cross-correlation before mixing.

  python wetdry.py dry.wav wet.wav outdir 25 50 75

Writes outdir/<wetname>_wet25.wav etc. 100 = untouched wet, 0 = dry.
"""

import subprocess
import sys
import wave
from pathlib import Path

import numpy as np

SR = 48000


def load(p: Path) -> np.ndarray:
    tmp = Path("/tmp/_wd_" + p.stem + ".wav")
    subprocess.run(
        ["ffmpeg", "-y", "-v", "error", "-i", str(p), "-ac", "2", "-ar", str(SR), str(tmp)],
        capture_output=True,
    )
    with wave.open(str(tmp)) as w:
        x = np.frombuffer(w.readframes(w.getnframes()), dtype=np.int16).astype(np.float64) / 32768.0
    return x.reshape(-1, 2)


def save(x: np.ndarray, p: Path) -> None:
    x = np.clip(x, -1.0, 1.0)
    with wave.open(str(p), "wb") as w:
        w.setnchannels(2)
        w.setsampwidth(2)
        w.setframerate(SR)
        w.writeframes((x * 32767).astype(np.int16).tobytes())


def best_lag(dry: np.ndarray, wet: np.ndarray, hop: int = 240) -> int:
    """Envelope cross-correlation lag, in samples."""
    def env(x):
        m = x.mean(axis=1)
        n = len(m) - len(m) % hop
        e = np.sqrt((m[:n].reshape(-1, hop) ** 2).mean(axis=1)) + 1e-9
        return (e - e.mean()) / (e.std() + 1e-9)

    a, b = env(dry), env(wet)
    n = min(len(a), len(b))
    c = np.correlate(a[:n], b[:n], "full")
    return int((np.argmax(c) - (n - 1)) * hop)


def main() -> int:
    if len(sys.argv) < 5:
        print(__doc__)
        return 2
    dry_p, wet_p, outdir = Path(sys.argv[1]), Path(sys.argv[2]), Path(sys.argv[3])
    mixes = [float(a) for a in sys.argv[4:]]
    outdir.mkdir(parents=True, exist_ok=True)

    dry, wet = load(dry_p), load(wet_p)
    lag = best_lag(dry, wet)
    if lag > 0:
        wet = np.vstack([np.zeros((lag, 2)), wet])
    elif lag < 0:
        wet = wet[-lag:]
    n = min(len(dry), len(wet))
    dry, wet = dry[:n], wet[:n]
    print(f"  aligned by {lag} samples ({lag / SR * 1000:+.0f} ms)")

    for m in mixes:
        w = m / 100.0
        out = dry * (1 - w) + wet * w
        # keep peak comparable to the wet source so mixes are level-matched
        peak = np.abs(out).max()
        ref = np.abs(wet).max()
        if peak > 0:
            out = out * (ref / peak) if peak > ref else out
        dest = outdir / f"{wet_p.stem}_wet{int(m):03d}.wav"
        save(out, dest)
        print(f"  {dest.name}")
    return 0


if __name__ == "__main__":
    sys.exit(main())