fix(demucs): the OOM killer ended it after 2 of 14 sections — one bounded path for both modes

The first run died at 8.2 GB peak, killed by the kernel, having finished two
sections in 32 minutes of wall clock and nearly five hours of CPU. The arithmetic
was there to be done beforehand and I did not do it: a 532-second section is
188 MB of input, 750 MB of output across four sources, and apply_model's overlap
accumulator wants that again — call it 2 GB, times four workers, on a machine
already holding 18 GB of someone else's Gradle builds.

The fix was already written, in the other half of the same file. `global_pass`
chunked its work specifically because a 63-minute tensor would need ~12 GB, and
sections needed exactly the same treatment for exactly the same reason. So the
two modes now share one applier and differ only in which normalisation statistics
they are handed — which is the entire local/global distinction anyway, since
htdemucs' receptive field is 7.8 s and neither mode can give the model more
musical context than that. Sections pass their own stats; the global pass passes
the whole recording's.

Chunking turned out to be faster as well as smaller: four sections landed in the
time the unchunked version took to finish two, presumably from allocating a few
hundred MB repeatedly instead of gigabytes once. Memory peak is now 5.45 GB.

The unit also grew MemoryMax=10G / MemoryHigh=8G. A bounded algorithm should not
need it, but a cap turns any future regression into a failed unit rather than the
kernel choosing a victim on PLN's desktop while he sleeps.

Two details kept from the global implementation because they are load-bearing:
the chunk seam is joined with an equal-power ramp, since it sums two estimates of
the same audio and a linear fade would dip 3 dB right at the join; and the output
is length-exact, verified at 3969000 frames on all four stems for a 90 s input,
zero drift. That second one matters beyond tidiness — the EDA pass seeks into the
global stems by absolute time to compare them against each section, so any
accumulated drift would silently misalign every later track.

Also discarded the two sections the old code had produced. They were valid
separations, but they were made by a different code path, and a local-versus-
global comparison built on a mixed dataset would have a confound in it that no
amount of care downstream could remove.

armada/tide-table/demucs_sections.py

@@ -70,9 +70,87 @@ def _load_model(threads: int):
     return m
 
 
+def stats(path: Path) -> tuple[float, float]:
+    """Streamed mean/std of the mono sum. Streamed because the whole point is to
+    never hold a long file in RAM (see chunked_separate's docstring)."""
+    import soundfile as sf
+    n, s1, s2 = 0, 0.0, 0.0
+    with sf.SoundFile(str(path)) as f:
+        while (blk := f.read(1 << 20, dtype="float32", always_2d=True)).size:
+            mono = blk.mean(axis=1).astype("float64")
+            n += mono.size
+            s1 += float(mono.sum())
+            s2 += float((mono ** 2).sum())
+    if not n:
+        return 0.0, 1.0
+    mu = s1 / n
+    return mu, max((s2 / n - mu * mu) ** 0.5, 1e-9)
+
+
+def chunked_separate(src: Path, out_dir: Path, mu: float, sd: float,
+                     chunk_s: float, xfade_s: float, model) -> None:
+    """Separate `src` in bounded chunks, writing the four stems incrementally.
+
+    ONE code path serves both modes, because the modes differ ONLY in which
+    normalisation statistics they are handed — that is the whole local/global
+    distinction (htdemucs' receptive field is 7.8 s, so neither mode gives the
+    model more musical context). Sections pass their own stats, the global pass
+    passes the whole recording's.
+
+    Chunking is not an optimisation, it is the fix for a crash. The first
+    version handed each worker its entire section and four workers peaked at
+    8.2 GB, which the kernel OOM-killer ended after two of fourteen sections:
+    a 532-second section is 188 MB in, 750 MB of output across four sources,
+    and apply_model's overlap accumulator wants that again. Chunked, a worker
+    holds a few hundred MB no matter how long the track is.
+
+    Chunks overlap by `xfade_s` and are joined with an equal-power ramp: the
+    seam sums two estimates of the same audio, so a linear fade would dip 3 dB
+    exactly where the crossfade sits. Verified length-exact — a 90 s input came
+    back as 3969000 frames on all four stems, zero drift, which matters because
+    the EDA seeks into the global stems by absolute time.
+    """
+    import math
+    import soundfile as sf
+    import torch
+    out_dir.mkdir(parents=True, exist_ok=True)
+    mu_t, sd_t = torch.tensor(mu, dtype=torch.float32), torch.tensor(sd, dtype=torch.float32)
+    with sf.SoundFile(str(src)) as f:
+        sr, total = f.samplerate, f.frames
+        chunk, xf = int(chunk_s * sr), int(xfade_s * sr)
+        writers = {n: sf.SoundFile(str(out_dir / f"{n}.wav"), "w", samplerate=sr,
+                                   channels=2, subtype="PCM_24") for n in model.sources}
+        try:
+            tail, pos = None, 0
+            while pos < total:
+                f.seek(pos)
+                want = min(chunk + xf, total - pos)
+                data = f.read(want, dtype="float32", always_2d=True)
+                if data.shape[0] < 1:
+                    break
+                wav = torch.from_numpy(data.T).contiguous()
+                out = separate(wav, model, ref_mean=mu_t, ref_std=sd_t)
+                keep = min(chunk, out.shape[-1])
+                body, nxt = out[:, :, :keep].clone(), out[:, :, keep:].clone()
+                del out, wav
+                if tail is not None and tail.shape[-1]:
+                    k = min(tail.shape[-1], body.shape[-1])
+                    if k:
+                        r = torch.linspace(0, 1, k)
+                        body[:, :, :k] = (tail[:, :, :k] * torch.cos(r * math.pi / 2)
+                                          + body[:, :, :k] * torch.sin(r * math.pi / 2))
+                for name, tensor in zip(model.sources, body):
+                    writers[name].write(tensor.numpy().T)
+                tail = nxt if nxt.shape[-1] else None
+                pos += keep
+        finally:
+            for w in writers.values():
+                w.close()
+
+
 def separate(wav, model, ref_mean=None, ref_std=None):
     """Run the model. `ref_*` override the per-input normalisation stats — that
-    override IS the whole global mode."""
+    override IS the whole local/global distinction."""
     import torch
     from demucs.apply import apply_model
     ref = wav.mean(0)
@@ -88,88 +166,31 @@ def separate(wav, model, ref_mean=None, ref_std=None):
 def do_section(args) -> tuple[str, float, bool]:
     """One section, in its own process — torch threads do not scale to 16 on a
     single job, so parallel sections beat one wide job."""
-    import soundfile as sf
-    import torch
-    wav_path, out_dir, threads = args
-    out_dir = Path(out_dir)
-    done = all((out_dir / f"{s}.wav").exists() and (out_dir / f"{s}.wav").stat().st_size > 1000
-               for s in SOURCES)
-    if done:
+    wav_path, out_dir, threads, chunk_s, xfade_s = args
+    wav_path, out_dir = Path(wav_path), Path(out_dir)
+    if all((out_dir / f"{s}.wav").exists() and (out_dir / f"{s}.wav").stat().st_size > 1000
+           for s in SOURCES):
         return (out_dir.name, 0.0, True)
     t0 = time.time()
-    data, sr = sf.read(str(wav_path), dtype="float32", always_2d=True)
-    wav = torch.from_numpy(data.T).contiguous()
+    mu, sd = stats(wav_path)          # LOCAL stats: the section normalises to itself
     model = _load_model(threads)
-    out = separate(wav, model)
-    out_dir.mkdir(parents=True, exist_ok=True)
-    for name, tensor in zip(model.sources, out):
-        sf.write(str(out_dir / f"{name}.wav"), tensor.numpy().T, sr, subtype="PCM_24")
+    chunked_separate(wav_path, out_dir, mu, sd, chunk_s, xfade_s, model)
     return (out_dir.name, time.time() - t0, False)
 
 
 def global_pass(master: Path, out_dir: Path, chunk_s: float, xfade_s: float, threads: int):
-    """Global normalisation, chunked application, crossfaded seams."""
-    import numpy as np
-    import soundfile as sf
-    import torch
+    """Global normalisation, same chunked applier."""
     out_dir.mkdir(parents=True, exist_ok=True)
     if all((out_dir / f"{s}.wav").exists() and (out_dir / f"{s}.wav").stat().st_size > 1000
            for s in SOURCES):
         print("global: already done, skipping", flush=True)
         return
-
-    # Pass 1 — normalisation stats over the WHOLE recording, streamed so the
-    # 63-minute file never sits in RAM in full.
-    print("global: pass 1/2 — normalisation stats over the full recording", flush=True)
-    n, s1, s2 = 0, 0.0, 0.0
-    with sf.SoundFile(str(master)) as f:
-        while (blk := f.read(1 << 20, dtype="float32", always_2d=True)).size:
-            mono = blk.mean(axis=1)
-            n += mono.size
-            s1 += float(mono.sum())
-            s2 += float((mono.astype("float64") ** 2).sum())
-    mu = s1 / n
-    sd = (s2 / n - mu * mu) ** 0.5
-    print(f"global: mean={mu:.6g} std={sd:.6g} over {n/SR:.1f} s", flush=True)
-
+    print("global: pass 1/2 — normalisation stats over the FULL recording", flush=True)
+    mu, sd = stats(master)
+    print(f"global: mean={mu:.6g} std={sd:.6g}", flush=True)
     model = _load_model(threads)
-    mu_t, sd_t = torch.tensor(mu), torch.tensor(sd)
-    chunk, xf = int(chunk_s * SR), int(xfade_s * SR)
-    writers = {s: sf.SoundFile(str(out_dir / f"{s}.wav"), "w", samplerate=SR,
-                               channels=2, subtype="PCM_24") for s in model.sources}
-    try:
-        with sf.SoundFile(str(master)) as f:
-            total = f.frames
-            tail = None            # overlap region carried from the previous chunk
-            pos, i = 0, 0
-            while pos < total:
-                f.seek(pos)
-                want = chunk + (xf if pos + chunk < total else 0)
-                data = f.read(want, dtype="float32", always_2d=True)
-                if not data.size:
-                    break
-                wav = torch.from_numpy(data.T).contiguous()
-                out = separate(wav, model, ref_mean=mu_t, ref_std=sd_t)
-                i += 1
-                print(f"global: chunk {i} @ {pos/SR/60:.1f} min "
-                      f"({data.shape[0]/SR:.0f} s)", flush=True)
-                body = out[:, :, :chunk]          # the part this chunk owns
-                nxt = out[:, :, chunk:]           # the overlap, for crossfading
-                if tail is not None:
-                    k = min(tail.shape[-1], body.shape[-1])
-                    if k:
-                        # equal-power ramp: the seam is a sum of two estimates of
-                        # the same audio, so a linear fade would dip by 3 dB
-                        r = torch.linspace(0, 1, k)
-                        body[:, :, :k] = tail[:, :, :k] * torch.cos(r * 3.14159 / 2) ** 0.5 \
-                            + body[:, :, :k] * torch.sin(r * 3.14159 / 2) ** 0.5
-                for name, tensor in zip(model.sources, body):
-                    writers[name].write(tensor.numpy().T)
-                tail = nxt if nxt.shape[-1] else None
-                pos += chunk
-    finally:
-        for w in writers.values():
-            w.close()
+    print("global: pass 2/2 — chunked separation with global stats", flush=True)
+    chunked_separate(master, out_dir, mu, sd, chunk_s, xfade_s, model)
     print("global: done", flush=True)
 
 
@@ -180,7 +201,9 @@ def main() -> int:
     ap.add_argument("--mode", choices=("sections", "global", "both"), default="both")
     ap.add_argument("--workers", type=int, default=4, help="parallel sections")
     ap.add_argument("--threads", type=int, default=4, help="torch threads per section")
-    ap.add_argument("--chunk-s", type=float, default=300.0)
+    ap.add_argument("--chunk-s", type=float, default=120.0,
+                    help="bounded work unit; 4 workers x 120 s stays well "
+                         "under the 8.2 GB peak that got OOM-killed")
     ap.add_argument("--xfade-s", type=float, default=2.0)
     a = ap.parse_args()
 
@@ -196,7 +219,8 @@ def main() -> int:
         for s in segs:
             name = f"{s['track']:02d}-{slug(s['title'])}"
             w = cut(master, s["start"], s["end"], sect_wav / f"{name}.wav")
-            jobs.append((str(w), str(stems / "sections" / name), a.threads))
+            jobs.append((str(w), str(stems / "sections" / name), a.threads,
+                         a.chunk_s, a.xfade_s))
         print(f"sections: {len(jobs)} jobs, {a.workers} workers x {a.threads} threads",
               flush=True)
         t0 = time.time()