[Model] Opt-in AudioSeal watermark for generated speech

benchmark/audio_watermark_check.py

@@ -0,0 +1,247 @@
+#!/usr/bin/env python3
+"""Checks behind the opt-in AudioSeal watermark (``mstar/model/audio_watermark.py``).
+
+Needs ``pip install -e '.[audioseal,audio]'``; ``quality`` also needs ``pesq``.
+
+``exactness``: on a synthetic 3 s, 24 kHz speech-like signal, the streamed
+watermark (aligned, uneven and batched chunkings) against AudioSeal's own
+whole-utterance ``get_watermark``, plus AudioSeal 0.2's ``streaming()`` for
+contrast (it does not carry history, upstream issue #105)::
+
+    python -m benchmark.audio_watermark_check exactness --device cuda
+
+``cost``: one ``apply_pcm16`` call per codec chunk length (1 to 25 Qwen3-TTS
+frames of 1920 samples), batch 1 and 8, streams already under way::
+
+    python -m benchmark.audio_watermark_check cost --device cuda
+
+``quality``: from the per-request WAVs ``benchmark.tts_speech_bench
+--save-audio-dir`` writes, one directory served with the watermark off and
+one with it on. The clean files are marked here chunk by chunk along the
+codec's schedule (served greedy speech does not repeat across server
+restarts, so served on/off pairs are not the same speech) for SNR, PESQ and
+streamed-vs-whole; the detector then scores every served file, whole and one
+codec chunk at a time, and whole copies resampled to 16 and 48 kHz::
+
+    python -m benchmark.audio_watermark_check quality --clean off_wav --marked on_wav --out quality.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from mstar.model.audio_watermark import AudioSealDetector, AudioSealWatermarker, AudioWatermarkConfig, WatermarkStream
+
+MESSAGE = 0xBEEF
+FRAME = 1920  # Qwen3-TTS samples per codec frame (24 kHz)
+SCHEDULE = (1, 3, 8, 16)  # the codec's first windows, in frames
+CHUNK = 25  # frames per window after the ramp
+RESAMPLED_RATES = (16000, 48000)  # copies the detector must still find the mark in
+
+
+def _stream(wmk: AudioSealWatermarker, x: torch.Tensor, sizes: list[int]) -> torch.Tensor:
+    stream, out, pos = WatermarkStream(), [], 0
+    for i, n in enumerate(sizes):
+        out.append(wmk.apply([x[pos : pos + n]], [stream], [i == len(sizes) - 1])[0])
+        pos += n
+    return torch.cat(out)
+
+
+def exactness(args) -> dict:
+    from audioseal.libs.moshi.utils.compile import no_compile
+
+    torch.manual_seed(0)
+    cfg = AudioWatermarkConfig(enabled=True, message=MESSAGE)
+    wmk = AudioSealWatermarker.build(cfg, args.device)
+    sr = 24000
+    t = torch.arange(3 * sr) / sr
+    phase = 2 * torch.pi * torch.cumsum(140 + 30 * torch.sin(2 * torch.pi * 0.7 * t), 0) / sr
+    x = sum(torch.sin(k * phase) / k for k in range(1, 12)) * (0.5 + 0.5 * torch.sin(2 * torch.pi * 3 * t)) * 0.15
+    x = (x + 0.005 * torch.randn_like(x)).to(args.device)
+    msg = torch.tensor([cfg.message_bits()], device=args.device)
+
+    def library(audio):
+        # AudioSeal's own call in full fp32; on a GPU the watermarker's TF32 cuDNN convs land a few int16 steps off
+        matmul, cudnn = torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32
+        torch.backends.cuda.matmul.allow_tf32 = torch.backends.cudnn.allow_tf32 = False
+        try:
+            with torch.no_grad(), no_compile():
+                return audio + wmk.generator.get_watermark(audio.view(1, 1, -1), message=msg).view(-1)
+        finally:
+            torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32 = matmul, cudnn
+
+    ref = library(x)
+    n = x.shape[0]
+    aligned = [1920, 5760, 15360, 48000]
+    aligned.append(n - sum(aligned))
+    uneven = [2048] * 10 + [n - 20480]
+    report = {
+        "device": str(args.device),
+        "samples": n,
+        "offline_max_abs": (wmk.offline(x) - ref).abs().max().item(),
+        "aligned_max_abs": (_stream(wmk, x, aligned) - ref).abs().max().item(),
+        "uneven_max_abs": (_stream(wmk, x, uneven) - ref).abs().max().item(),
+    }
+    xs = [x, torch.roll(x, 777), 0.5 * torch.flip(x, [0])]
+    streams = [WatermarkStream() for _ in xs]
+    outs: list[list[torch.Tensor]] = [[] for _ in xs]
+    pos = 0
+    for i, size in enumerate(aligned):
+        marked = wmk.apply([xi[pos : pos + size] for xi in xs], streams, [i == len(aligned) - 1] * len(xs))
+        for k, m in enumerate(marked):
+            outs[k].append(m)
+        pos += size
+    errors = [(torch.cat(o) - library(xi)).abs().max().item() for o, xi in zip(outs, xs, strict=True)]
+    report["batched_max_abs"] = max(errors)
+    with torch.no_grad(), no_compile(), wmk.generator.streaming(1):
+        parts, pos = [], 0
+        for size in aligned:
+            parts.append(wmk.generator(x[pos : pos + size].view(1, 1, -1), message=msg).view(-1))
+            pos += size
+    report["library_streaming_max_abs"] = (torch.cat(parts) - ref).abs().max().item()
+    return report
+
+
+def cost(args) -> dict:
+    device = torch.device(args.device)
+    sync = torch.cuda.synchronize if device.type == "cuda" else (lambda: None)
+    wmk = AudioSealWatermarker.build(AudioWatermarkConfig(enabled=True, message=MESSAGE), device)
+    rows = []
+    for batch in (1, 8):
+        for frames in (1, 3, 8, 16, 25):
+            streams = [WatermarkStream() for _ in range(batch)]
+            chunks = [(torch.randn(FRAME * frames, device=device) * 3000).to(torch.int16) for _ in range(batch)]
+            for _ in range(5):
+                wmk.apply_pcm16(chunks, streams)
+            times = []
+            for _ in range(args.iters):
+                sync()
+                t0 = time.perf_counter()
+                wmk.apply_pcm16(chunks, streams)
+                sync()
+                times.append(1000 * (time.perf_counter() - t0))
+            times.sort()
+            p50, p90 = statistics.median(times), times[int(0.9 * len(times))]
+            rows.append({"batch": batch, "frames": frames, "audio_ms": FRAME * frames / 24,
+                         "ms_p50": round(p50, 2), "ms_p90": round(p90, 2)})
+            print(json.dumps(rows[-1]), flush=True)
+    return {"device": torch.cuda.get_device_name(device) if device.type == "cuda" else "cpu", "rows": rows}
+
+
+def _chunk_edges(n: int) -> list[int]:
+    edges, pos, sizes = [0], 0, [s * FRAME for s in SCHEDULE]
+    while pos < n:
+        pos = min(n, pos + (sizes.pop(0) if sizes else CHUNK * FRAME))
+        edges.append(pos)
+    return edges
+
+
+def _wavs(directory: str) -> dict[str, np.ndarray]:
+    import soundfile as sf
+
+    return {p.stem: sf.read(p, dtype="int16")[0] for p in sorted(Path(directory).glob("*.wav"))}
+
+
+def _rates(marked: list[tuple[float, int]], clean: list[tuple[float, int]], threshold: float) -> dict:
+    return {
+        "marked": len(marked),
+        "marked_detected": sum(s > threshold for s, _ in marked),
+        "marked_message_exact": sum(m == MESSAGE for _, m in marked),
+        "clean": len(clean),
+        "clean_flagged": sum(s > threshold for s, _ in clean),
+        "clean_max_score": round(max((s for s, _ in clean), default=0.0), 3),
+    }
+
+
+def _detect(det: AudioSealDetector, files: dict[str, np.ndarray]) -> tuple[list, dict[int, list], dict[int, list]]:
+    import torchaudio
+
+    whole, per_chunk, resampled = [], {}, {}
+    for x in files.values():
+        if len(x) < FRAME:
+            continue
+        audio = torch.from_numpy(x.astype(np.float32) / 32767)
+        whole.append(det.detect(audio))
+        for rate in RESAMPLED_RATES:
+            copy = torchaudio.functional.resample(audio, 24000, rate)
+            resampled.setdefault(rate, []).append(det.detect(copy, rate))
+        edges = _chunk_edges(len(x))
+        for a, b in zip(edges, edges[1:], strict=False):
+            if b - a == FRAME * ((b - a) // FRAME):  # whole codec windows only
+                score = det.detect(torch.from_numpy(x[a:b].astype(np.float32) / 32767))
+                per_chunk.setdefault((b - a) // FRAME, []).append(score)
+    return whole, per_chunk, resampled
+
+
+def quality(args) -> dict:
+    import torchaudio
+    from pesq import pesq
+
+    def to_16k(audio: np.ndarray) -> np.ndarray:
+        return torchaudio.functional.resample(torch.from_numpy(audio), 24000, 16000).numpy()
+
+    wmk = AudioSealWatermarker.build(AudioWatermarkConfig(enabled=True, message=MESSAGE), args.device)
+    det = AudioSealDetector.build(args.device)
+    clean = _wavs(args.clean)
+    snrs, pesqs, lsb = [], [], []
+    for x in clean.values():
+        pcm = torch.from_numpy(x.copy())
+        stream, edges = WatermarkStream(), _chunk_edges(len(x))
+        spans = zip(edges, edges[1:], strict=False)
+        parts = [wmk.apply_pcm16([pcm[a:b]], [stream], [b == edges[-1]])[0] for a, b in spans]
+        streamed = torch.cat(parts).cpu().numpy()
+        whole = wmk.apply_pcm16([pcm], [WatermarkStream()], [True])[0].cpu().numpy()
+        lsb.append(int(np.abs(streamed.astype(np.int32) - whole.astype(np.int32)).max()))
+        c, m = x.astype(np.float64) / 32767, streamed.astype(np.float64) / 32767
+        snrs.append(10 * np.log10((c**2).sum() / max(((m - c) ** 2).sum(), 1e-20)))
+        pesqs.append(pesq(16000, to_16k(c), to_16k(m), "wb"))
+    report = {
+        "files": len(snrs),
+        "snr_db_mean_min": [round(statistics.mean(snrs), 1), round(min(snrs), 1)],
+        "pesq_wb_mean_min": [round(statistics.mean(pesqs), 2), round(min(pesqs), 2)],
+        "streamed_vs_whole_max_lsb": max(lsb),
+    }
+    if args.marked:
+        on_whole, on_chunks, on_resampled = _detect(det, _wavs(args.marked))
+        off_whole, off_chunks, off_resampled = _detect(det, clean)
+        report["served_whole_files"] = _rates(on_whole, off_whole, args.threshold)
+        for rate in RESAMPLED_RATES:
+            report[f"served_whole_files_at_{rate // 1000}k"] = _rates(
+                on_resampled[rate], off_resampled[rate], args.threshold
+            )
+        report["served_per_chunk"] = {
+            f"{f} frames ({f * FRAME // 24} ms)": _rates(on_chunks.get(f, []), off_chunks.get(f, []), args.threshold)
+            for f in sorted(set(on_chunks) | set(off_chunks))
+        }
+    return report
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    sub = ap.add_subparsers(dest="command", required=True)
+    for name in ("exactness", "cost", "quality"):
+        p = sub.add_parser(name)
+        p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+        p.add_argument("--out", default=None, help="also write the report as JSON")
+        if name == "cost":
+            p.add_argument("--iters", type=int, default=40)
+        if name == "quality":
+            p.add_argument("--clean", required=True, help="WAVs served with the watermark off")
+            p.add_argument("--marked", default=None, help="WAVs served with the watermark on")
+            p.add_argument("--threshold", type=float, default=0.5)
+    args = ap.parse_args()
+    report = {"exactness": exactness, "cost": cost, "quality": quality}[args.command](args)
+    print(json.dumps(report, indent=1))
+    if args.out:
+        Path(args.out).write_text(json.dumps(report, indent=1))
+
+
+if __name__ == "__main__":
+    main()

docs/installation.rst

@@ -112,6 +112,12 @@ Model families and some output formats need extra packages, exposed as pip *extr
    * - ``.[audio]``
      - ``soundfile`` — only needed to return **non-WAV** audio containers (mp3/flac/…)
        from the OpenAI/SDK audio surfaces. WAV/PCM output works without it.
+   * - ``.[audioseal]``
+     - ``audioseal`` (``>=0.2,<0.3``, MIT code and weights) for the opt-in AudioSeal
+       watermark on generated speech (Qwen3-TTS ``audio_watermark``, see the
+       Qwen3-TTS notes in :doc:`models`). Not part of ``.[all]``; a deployment that
+       enables the watermark without it fails at startup with this install hint.
+       The ``python -m mstar.model.audio_watermark detect`` CLI also needs ``.[audio]``.
    * - ``.[dev]``
      - ``ruff`` + ``pytest`` for linting and the test suite.
    * - ``.[all]``

docs/models.rst

@@ -297,6 +297,26 @@ Qwen3-TTS notes
   decoder bit for bit. ``model_kwargs: {codec_dtype: bfloat16}`` in the
   deployment YAML halves its GPU time when throughput matters more than
   bit-exactness (the Talker is bf16 either way).
+- Opt-in AudioSeal watermark on the codec's output (off by default; install
+  ``pip install -e '.[audioseal]'``). Turn it on in the deployment YAML::
+
+     model_kwargs:
+       audio_watermark:
+         enabled: true      # default false
+         message: 0xBEEF    # 16-bit payload: int 0..65535 or a 16-character 0/1 string (default 0)
+         alpha: 1.0         # strength, in (0, 4] (default 1.0)
+         generator: audioseal_wm_streaming   # AudioSeal card name or checkpoint path (default)
+
+  Unknown keys or out-of-range values fail startup, and so does ``enabled: true``
+  without the package. Each request's PCM is marked chunk by chunk with its own
+  streaming state, so a streamed utterance carries the same mark as a whole one.
+  Check files with ``python -m mstar.model.audio_watermark detect --message 0xBEEF out.wav``
+  (``--message``, the deployment's payload: without it the verdict is the
+  score alone, and an unmarked file re-encoded as MP3 can score above 0.5;
+  ``--threshold``, default 0.5; ``--device``; needs ``.[audio]`` for
+  ``soundfile``); a file at another rate than 24 kHz is resampled to 24 kHz
+  first, since the mark is found only at the rate it was added at. Detect on
+  whole utterances: a clip under a few seconds can miss the mark. The helper lives in ``mstar/model/audio_watermark.py``.
 - Weight loading checks coverage in both directions: a parameter the checkpoint
   does not fill, or a checkpoint tensor the port does not load, fails startup.
 - The 12 Hz codec does not require the system SoX executable. M* imports only

mstar/engine/engine.py

@@ -1657,6 +1657,28 @@ def _postprocess_batch(
                     rid, batch.node_name, batch.step_context.graph_walk,
                 )
                 batch.register_failure(rid, error)
+        if type(submodule).batch_postprocess is not NodeSubmodule.batch_postprocess:
+            live = [
+                (rid, node_inputs)
+                for rid, node_inputs in zip(batch.request_ids, batch.inputs, strict=True)
+                if rid not in batch.failed_requests
+            ]
+            if not live:
+                return
+            rids = [rid for rid, _ in live]
+            try:
+                submodule.batch_postprocess(
+                    outputs,
+                    {rid: batch.per_request_info_wrapped[rid] for rid in rids},
+                    [node_inputs for _, node_inputs in live],
+                )
+            except Exception as error:
+                logger.exception(
+                    "batch_postprocess failed for requests %s (node=%s, walk=%s)",
+                    rids, batch.node_name, batch.step_context.graph_walk,
+                )
+                for rid in rids:
+                    batch.register_failure(rid, error)
 
     def exec_and_postprocess(
         self, batch: ExecutingBatch
@@ -1676,7 +1698,6 @@ def exec_and_postprocess(
                 self.postprocess_batch(batch, outputs)
             else:
                 submodule = self._submodules[batch.node_name].submodule
-                # no submodule overrides it as of now
                 if type(submodule).batch_postprocess is not NodeSubmodule.batch_postprocess:
                     submodule.batch_postprocess(
                         outputs, batch.per_request_info_wrapped, batch.inputs,