From 020aa4db4fe5c390a9563252ece256110cc91ee9 Mon Sep 17 00:00:00 2001 From: Yash Hayaran Date: Wed, 2 Sep 2026 16:11:30 +0530 Subject: [PATCH 1/4] feat(tts): vendor Perth watermarking in EasyMagpie codec stage --- .../serving/test_codec_model.py | 3 +- .../serving/test_watermark.py | 162 +++++++++++++++++ tools/easymagpie_vllm_omni/README.md | 12 ++ .../easymagpie_vllm_omni/codec/model.py | 8 + .../easymagpie_vllm_omni/watermark.py | 169 ++++++++++++++++++ tools/easymagpie_vllm_omni/requirements.txt | 11 ++ .../scripts/benchmark_incremental_server.py | 2 + .../scripts/benchmark_server.py | 2 + 8 files changed, 368 insertions(+), 1 deletion(-) create mode 100644 tests/collections/tts/easymagpie_vllm_omni/serving/test_watermark.py create mode 100644 tools/easymagpie_vllm_omni/easymagpie_vllm_omni/watermark.py diff --git a/tests/collections/tts/easymagpie_vllm_omni/serving/test_codec_model.py b/tests/collections/tts/easymagpie_vllm_omni/serving/test_codec_model.py index d455c24f5c51..fb71cf8abcc7 100644 --- a/tests/collections/tts/easymagpie_vllm_omni/serving/test_codec_model.py +++ b/tests/collections/tts/easymagpie_vllm_omni/serving/test_codec_model.py @@ -57,7 +57,8 @@ def forward(self, codes: torch.Tensor) -> torch.Tensor: ([1025, 3, 1025, 103], 6), ], ) -def test_forward_trims_terminal_control_subframes(terminal_row, expected_samples): +def test_forward_trims_terminal_control_subframes(terminal_row, expected_samples, monkeypatch): + monkeypatch.setenv("NEMOTRON_TTS_PERTH_WATERMARK", "0") model = EasyMagpieCodecForConditionalGeneration.__new__(EasyMagpieCodecForConditionalGeneration) torch.nn.Module.__init__(model) model.config = SimpleNamespace( diff --git a/tests/collections/tts/easymagpie_vllm_omni/serving/test_watermark.py b/tests/collections/tts/easymagpie_vllm_omni/serving/test_watermark.py new file mode 100644 index 000000000000..681e3aadb662 --- /dev/null +++ b/tests/collections/tts/easymagpie_vllm_omni/serving/test_watermark.py @@ -0,0 +1,162 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import types +from unittest import mock + +import pytest +import torch +from easymagpie_vllm_omni import watermark + + +class _FakeAudioProcessor: + def signal_to_magphase(self, signal: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return signal, torch.zeros_like(signal) + + def magphase_to_signal(self, magnitude: torch.Tensor, _phase: torch.Tensor) -> torch.Tensor: + return magnitude + + +class _FakeEncoder: + def __init__(self) -> None: + self.batch_sizes: list[int] = [] + + def __call__(self, magnitude: torch.Tensor) -> tuple[torch.Tensor, None]: + self.batch_sizes.append(int(magnitude.shape[0])) + return magnitude + 0.125, None + + +class _FakePerthNet: + def __init__(self) -> None: + self.hp = types.SimpleNamespace(sample_rate=22_050, n_fft=8) + self.device = torch.device("cpu") + self.ap = _FakeAudioProcessor() + self.encoder = _FakeEncoder() + + +def setup_function() -> None: + watermark._WATERMARKER = None + watermark._INITIALIZED_DEVICE = None + watermark._SHORT_AUDIO_WARNING_EMITTED = False + + +def test_unindexed_cuda_device_is_normalized() -> None: + with ( + mock.patch.object(torch.cuda, "is_available", return_value=True), + mock.patch.object(torch.cuda, "current_device", return_value=2), + ): + assert watermark._normalized_device("cuda") == "cuda:2" + + +def test_equal_length_waveforms_are_batched_and_length_is_preserved() -> None: + perth_net = _FakePerthNet() + watermarker = types.SimpleNamespace(perth_net=perth_net) + waveforms = [torch.zeros(16), torch.full((16,), 0.5), torch.zeros(12)] + + with mock.patch.object(watermark, "initialize_watermarker", return_value=watermarker): + result = watermark.watermark_waveforms( + waveforms, + sample_rate=22_050, + device="cpu", + ) + + assert perth_net.encoder.batch_sizes == [2, 1] + assert [tensor.numel() for tensor in result] == [16, 16, 12] + torch.testing.assert_close(result[0], torch.full((16,), 0.125)) + torch.testing.assert_close(result[1], torch.full((16,), 0.625)) + + +def test_short_waveform_is_left_unchanged() -> None: + perth_net = _FakePerthNet() + watermarker = types.SimpleNamespace(perth_net=perth_net) + original = torch.tensor([0.1, -0.2, 0.3, -0.4]) + + with mock.patch.object(watermark, "initialize_watermarker", return_value=watermarker): + result = watermark.watermark_waveforms( + [original.clone()], + sample_rate=22_050, + device="cpu", + ) + + assert perth_net.encoder.batch_sizes == [] + torch.testing.assert_close(result[0], original) + + +def test_explicit_disable_bypasses_initialization() -> None: + waveforms = [torch.zeros(16)] + with ( + mock.patch.dict("os.environ", {"NEMOTRON_TTS_PERTH_WATERMARK": "off"}), + mock.patch.object(watermark, "initialize_watermarker") as initialize, + ): + result = watermark.watermark_waveforms( + waveforms, + sample_rate=22_050, + device="cpu", + ) + + initialize.assert_not_called() + assert result is waveforms + + +def test_enabled_initialization_failure_is_fatal() -> None: + broken_perth = types.SimpleNamespace( + PerthImplicitWatermarker=mock.Mock(side_effect=ValueError("bad checkpoint")) + ) + with ( + mock.patch.dict("os.environ", {"NEMOTRON_TTS_PERTH_WATERMARK": "1"}), + mock.patch.dict("sys.modules", {"perth": broken_perth}), + pytest.raises(RuntimeError, match="could not be initialized"), + ): + watermark.initialize_watermarker("cpu") + + +def test_encoder_is_compiled_on_cuda() -> None: + encoder = object() + perth_net = types.SimpleNamespace( + hp=types.SimpleNamespace(sample_rate=32_000, n_fft=8), + encoder=encoder, + device="cuda:0", + ap=_FakeAudioProcessor(), + ) + watermarker = types.SimpleNamespace(perth_net=perth_net) + compiled = mock.Mock(name="compiled_encoder", side_effect=lambda mag: (mag, None)) + fake_perth = types.SimpleNamespace( + PerthImplicitWatermarker=mock.Mock(return_value=watermarker) + ) + + with ( + mock.patch.dict("os.environ", {"NEMOTRON_TTS_PERTH_WATERMARK": "1"}), + mock.patch.dict("sys.modules", {"perth": fake_perth}), + mock.patch.object(watermark, "_normalized_device", return_value="cuda:0"), + mock.patch.object(torch, "compile", return_value=compiled) as compile_fn, + mock.patch.object(torch.cuda, "is_available", return_value=False), + ): + result = watermark.initialize_watermarker("cuda:0") + + compile_fn.assert_called_once_with(encoder, mode="default", dynamic=True) + assert result.perth_net.encoder is compiled + + +def test_encoder_is_not_compiled_on_cpu() -> None: + encoder = object() + perth_net = types.SimpleNamespace( + hp=types.SimpleNamespace(sample_rate=32_000, n_fft=8), + encoder=encoder, + device="cpu", + ) + watermarker = types.SimpleNamespace(perth_net=perth_net) + fake_perth = types.SimpleNamespace( + PerthImplicitWatermarker=mock.Mock(return_value=watermarker) + ) + + with ( + mock.patch.dict("os.environ", {"NEMOTRON_TTS_PERTH_WATERMARK": "1"}), + mock.patch.dict("sys.modules", {"perth": fake_perth}), + mock.patch.object(torch, "compile") as compile_fn, + ): + result = watermark.initialize_watermarker("cpu") + + compile_fn.assert_not_called() + assert result.perth_net.encoder is encoder diff --git a/tools/easymagpie_vllm_omni/README.md b/tools/easymagpie_vllm_omni/README.md index d27b116ead98..3e5fb504fea4 100644 --- a/tools/easymagpie_vllm_omni/README.md +++ b/tools/easymagpie_vllm_omni/README.md @@ -54,6 +54,18 @@ Mamba's selective-state-update kernel requires shape- and GPU-specific tuning, s suboptimal performance. Reuse the same Triton/vLLM cache directories across launches so repeated runs accumulate better kernels; for an explicit sweep, run `python scripts/tune_mamba_ssu.py --model converted_model` and restart. +### Perth watermarking + +EasyMagpie audio is watermarked after native codec decoding and before float +PCM leaves the vLLM-Omni codec stage. Equal-length outputs are processed as a +GPU batch, including outputs produced by incremental streaming requests. The +watermarker resamples codec audio to Perth's model rate and back while +preserving each output's original sample count. + +Watermarking is enabled by default, and engine startup fails if Perth or its +bundled checkpoint cannot be loaded. For controlled quality comparisons only, +set `NEMOTRON_TTS_PERTH_WATERMARK=0` to disable it. + ### Quick start — offline synthesis See the [`offline_demo.ipynb`](../../tutorials/tts/easymagpie_vllm_omni/offline_demo.ipynb) tutorial to check how diff --git a/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/codec/model.py b/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/codec/model.py index 5cbb73fa0b38..31c489bd40ab 100644 --- a/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/codec/model.py +++ b/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/codec/model.py @@ -22,6 +22,7 @@ import torch.nn as nn from easymagpie_vllm_omni.codec.config import EasyMagpieCodecConfig from easymagpie_vllm_omni.codec.packed import PackedEasyMagpieCodec +from easymagpie_vllm_omni.watermark import initialize_watermarker, watermark_waveforms from vllm.config import VllmConfig from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFuncCalculator from vllm.model_executor.models.utils import AutoWeightsLoader @@ -73,6 +74,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.has_preprocess = False self.has_postprocess = False self.requires_raw_input_tokens = True + initialize_watermarker(vllm_config.device_config.device) def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: return torch.zeros((input_ids.shape[0], 1), dtype=torch.float32, device=input_ids.device) @@ -183,6 +185,12 @@ def forward( outputs.append(packed_audio[offset : offset + valid_samples].float()) frame_offset += frames offset += samples + if codec_codes is not None or runtime_additional_information: + outputs = watermark_waveforms( + outputs, + sample_rate=self.config.output_sample_rate, + device=packed_audio.device, + ) sample_rate = torch.tensor(self.config.output_sample_rate, dtype=torch.int32) return OmniOutput( text_hidden_states=None, diff --git a/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/watermark.py b/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/watermark.py new file mode 100644 index 000000000000..2134d55476e1 --- /dev/null +++ b/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/watermark.py @@ -0,0 +1,169 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GPU-batched Perth watermarking for EasyMagpie codec waveforms.""" + +from __future__ import annotations + +import logging +import os +from typing import Any + +import torch +import torch.nn.functional as F + +logger = logging.getLogger(__name__) + +_DISABLED_VALUES = {"0", "false", "no", "off"} +_WATERMARKER: Any | None = None +_INITIALIZED_DEVICE: str | None = None +_SHORT_AUDIO_WARNING_EMITTED = False + + +def _normalized_device(device: str | torch.device) -> str: + parsed = torch.device(device) + if parsed.type == "cuda" and parsed.index is None and torch.cuda.is_available(): + parsed = torch.device("cuda", torch.cuda.current_device()) + return str(parsed) + + +def watermarking_enabled() -> bool: + """Return whether production Perth watermarking is enabled.""" + value = os.environ.get("NEMOTRON_TTS_PERTH_WATERMARK", "1") + return value.strip().lower() not in _DISABLED_VALUES + + +def initialize_watermarker(device: str | torch.device) -> Any | None: + """Load Perth on ``device`` or fail startup when watermarking is enabled.""" + global _INITIALIZED_DEVICE, _WATERMARKER + + if not watermarking_enabled(): + logger.warning("Perth watermarking is explicitly disabled by NEMOTRON_TTS_PERTH_WATERMARK") + return None + + requested_device = _normalized_device(device) + if _WATERMARKER is not None: + if requested_device != _INITIALIZED_DEVICE: + raise RuntimeError( + "Perth was initialized on " + f"{_INITIALIZED_DEVICE}, not requested device {requested_device}" + ) + return _WATERMARKER + + try: + import perth + + watermarker_type = perth.PerthImplicitWatermarker + if watermarker_type is None: + raise ImportError("perth.PerthImplicitWatermarker is unavailable") + _WATERMARKER = watermarker_type(device=requested_device) + perth_net = _WATERMARKER.perth_net + if torch.device(requested_device).type == "cuda": + perth_net.encoder = torch.compile( + perth_net.encoder, mode="default", dynamic=True + ) + if torch.cuda.is_available(): + _warmup_compiled_encoder(_WATERMARKER) + except Exception as error: + _WATERMARKER = None + raise RuntimeError( + "Perth watermarking is enabled but its model could not be initialized" + ) from error + + _INITIALIZED_DEVICE = requested_device + logger.info( + "Perth watermarking enabled on device=%s (model sample rate=%d)", + requested_device, + int(_WATERMARKER.perth_net.hp.sample_rate), + ) + return _WATERMARKER + + +def _warmup_compiled_encoder(watermarker: Any) -> None: + """Trace the compiled encoder on startup and 8-frame-scale spectrograms.""" + perth_net = watermarker.perth_net + window = max(int(perth_net.hp.n_fft), 1) + dummy_lengths = (window * 4, window * 8) + with torch.inference_mode(): + for length in dummy_lengths: + dummy = torch.zeros(length, device=perth_net.device) + magnitudes, _ = perth_net.ap.signal_to_magphase(dummy.unsqueeze(0)) + perth_net.encoder(magnitudes) + + +def _restore_length(watermarked: torch.Tensor, original: torch.Tensor) -> torch.Tensor: + target_samples = int(original.shape[-1]) + current_samples = int(watermarked.shape[-1]) + if current_samples > target_samples: + return watermarked[..., :target_samples] + if current_samples < target_samples: + return F.pad(watermarked, (0, target_samples - current_samples)) + return watermarked + + +def watermark_waveforms( + waveforms: list[torch.Tensor], + *, + sample_rate: int, + device: str | torch.device, +) -> list[torch.Tensor]: + """Watermark non-empty waveforms, batching tensors with equal lengths.""" + global _SHORT_AUDIO_WARNING_EMITTED + + if not waveforms or not watermarking_enabled(): + return waveforms + + watermarker = initialize_watermarker(device) + if watermarker is None: + return waveforms + + perth_net = watermarker.perth_net + perth_device = perth_net.device + perth_rate = int(perth_net.hp.sample_rate) + minimum_samples = int(perth_net.hp.n_fft // 2 + 1) + + resample = None + if sample_rate != perth_rate: + from torchaudio.functional import resample + + length_groups: dict[int, list[int]] = {} + for index, waveform in enumerate(waveforms): + samples = int(waveform.numel()) + if samples: + length_groups.setdefault(samples, []).append(index) + + with torch.inference_mode(): + for indices in length_groups.values(): + originals = torch.stack( + [ + waveforms[index].detach().float().reshape(-1).to(perth_device) + for index in indices + ] + ) + signals = originals + if sample_rate != perth_rate: + assert resample is not None + signals = resample(signals, sample_rate, perth_rate) + + if signals.shape[-1] < minimum_samples: + if not _SHORT_AUDIO_WARNING_EMITTED: + logger.warning( + "Skipping Perth for audio shorter than %d samples at %d Hz", + minimum_samples, + perth_rate, + ) + _SHORT_AUDIO_WARNING_EMITTED = True + continue + + magnitudes, phases = perth_net.ap.signal_to_magphase(signals) + marked_magnitudes, _ = perth_net.encoder(magnitudes) + marked = perth_net.ap.magphase_to_signal(marked_magnitudes, phases) + if sample_rate != perth_rate: + assert resample is not None + marked = resample(marked, perth_rate, sample_rate) + marked = _restore_length(marked, originals).clamp_(-1.0, 1.0) + + for batch_index, waveform_index in enumerate(indices): + waveforms[waveform_index] = marked[batch_index].reshape(-1) + + return waveforms diff --git a/tools/easymagpie_vllm_omni/requirements.txt b/tools/easymagpie_vllm_omni/requirements.txt index 43d93f14547b..58f929fad049 100644 --- a/tools/easymagpie_vllm_omni/requirements.txt +++ b/tools/easymagpie_vllm_omni/requirements.txt @@ -1,7 +1,18 @@ # Runtime dependencies used by EasyMagpie or the speech-serving endpoint. +audioread==3.0.1 +decorator==5.2.1 +joblib==1.5.1 +lazy-loader==0.4 +librosa==0.11.0 numpy>=1.26 +pooch==1.8.2 PyYAML>=6.0 +resemble-perth==1.0.1 safetensors>=0.8.0 +scikit-learn==1.7.1 +soxr==0.5.0.post1 +threadpoolctl==3.6.0 +torchaudio soundfile>=0.13.1 tokenizers transformers>=5.5.3 diff --git a/tools/easymagpie_vllm_omni/scripts/benchmark_incremental_server.py b/tools/easymagpie_vllm_omni/scripts/benchmark_incremental_server.py index ec2da8512320..b4ec76b87daa 100644 --- a/tools/easymagpie_vllm_omni/scripts/benchmark_incremental_server.py +++ b/tools/easymagpie_vllm_omni/scripts/benchmark_incremental_server.py @@ -315,6 +315,7 @@ def parse_args(): parser.add_argument("--timeout", type=float, default=300.0, help="Per-request timeout in seconds") parser.add_argument("--no-warmup", action="store_true", help="Skip one warmup request per worker") parser.add_argument("--output-dir", default=None, help="If set, save each generated waveform") + parser.add_argument("--seed", type=int, default=1234, help="Prompt sampling seed (default: %(default)s)") return parser.parse_args() @@ -329,6 +330,7 @@ def main() -> None: if args.send_delay_ms < 0: raise ValueError("--send-delay-ms cannot be negative") + random.seed(args.seed) text_items = base._load_items(args.text_file) if not text_items: raise ValueError(f"No usable lines found in {args.text_file}") diff --git a/tools/easymagpie_vllm_omni/scripts/benchmark_server.py b/tools/easymagpie_vllm_omni/scripts/benchmark_server.py index 837a03928fd3..b19d01a71072 100644 --- a/tools/easymagpie_vllm_omni/scripts/benchmark_server.py +++ b/tools/easymagpie_vllm_omni/scripts/benchmark_server.py @@ -341,8 +341,10 @@ def main() -> None: parser.add_argument("--timeout", type=float, default=300, help="Per-request timeout, s (default: 300)") parser.add_argument("--no-warmup", action="store_true", help="Skip warmup phase (concurrency requests)") parser.add_argument("--output-dir", default=None, help="If set, write each waveform to /.wav") + parser.add_argument("--seed", type=int, default=1234, help="Prompt sampling seed (default: %(default)s)") args = parser.parse_args() + random.seed(args.seed) items = _load_items(args.text_file) if not items: print(f"ERROR: no usable lines found in {args.text_file}") From be63c74b1d859493d41a7cfadd3bac0129845191 Mon Sep 17 00:00:00 2001 From: Yash Hayaran Date: Fri, 4 Sep 2026 10:08:44 +0530 Subject: [PATCH 2/4] feat(perth): add optional dependencies and update README for watermarking support --- tools/easymagpie_vllm_omni/README.md | 13 +++++++++---- tools/easymagpie_vllm_omni/pyproject.toml | 7 +++++++ tools/easymagpie_vllm_omni/requirements-perth.txt | 6 ++++++ tools/easymagpie_vllm_omni/requirements.txt | 13 +------------ 4 files changed, 23 insertions(+), 16 deletions(-) create mode 100644 tools/easymagpie_vllm_omni/requirements-perth.txt diff --git a/tools/easymagpie_vllm_omni/README.md b/tools/easymagpie_vllm_omni/README.md index 3e5fb504fea4..2b002665de7a 100644 --- a/tools/easymagpie_vllm_omni/README.md +++ b/tools/easymagpie_vllm_omni/README.md @@ -42,7 +42,7 @@ cd tools/easymagpie_vllm_omni conda create -n easymagpie-vllm python=3.12 -y conda activate easymagpie-vllm pip install -r requirements.txt -pip install -e . +pip install -e ".[perth]" # optionally for notebook pip install ipykernel python -m ipykernel install --user \ @@ -50,6 +50,9 @@ python -m ipykernel install --user \ --display-name "Python (easymagpie-vllm)" ``` +The `perth` extra is required for default watermarking. `pip install -r requirements-perth.txt` +is equivalent if you install the package without extras. + Mamba's selective-state-update kernel requires shape- and GPU-specific tuning, so an untuned cache can give suboptimal performance. Reuse the same Triton/vLLM cache directories across launches so repeated runs accumulate better kernels; for an explicit sweep, run `python scripts/tune_mamba_ssu.py --model converted_model` and restart. @@ -62,9 +65,11 @@ GPU batch, including outputs produced by incremental streaming requests. The watermarker resamples codec audio to Perth's model rate and back while preserving each output's original sample count. -Watermarking is enabled by default, and engine startup fails if Perth or its -bundled checkpoint cannot be loaded. For controlled quality comparisons only, -set `NEMOTRON_TTS_PERTH_WATERMARK=0` to disable it. +Watermarking is enabled by default after installing the `perth` extra, and +engine startup fails if Perth or its bundled checkpoint cannot be loaded. +`resemble-perth==1.0.1` does not declare its own dependencies, so the extra +also installs `librosa` and `torchaudio` unpinned. For unmarked quality +comparisons only, set `NEMOTRON_TTS_PERTH_WATERMARK=0` to disable it. ### Quick start — offline synthesis diff --git a/tools/easymagpie_vllm_omni/pyproject.toml b/tools/easymagpie_vllm_omni/pyproject.toml index c6d4d8942c93..f27641829db9 100644 --- a/tools/easymagpie_vllm_omni/pyproject.toml +++ b/tools/easymagpie_vllm_omni/pyproject.toml @@ -13,6 +13,13 @@ dependencies = [ # vllm_omni_env"; do not install into NeMo's nemo_virtual_environment. ] +[project.optional-dependencies] +perth = [ + "resemble-perth==1.0.1", + "librosa>=0.11.0", + "torchaudio", +] + [project.entry-points."vllm.general_plugins"] easymagpie_omni = "vllm_plugin_easymagpie_omni:register" diff --git a/tools/easymagpie_vllm_omni/requirements-perth.txt b/tools/easymagpie_vllm_omni/requirements-perth.txt new file mode 100644 index 000000000000..0d48d7e9293f --- /dev/null +++ b/tools/easymagpie_vllm_omni/requirements-perth.txt @@ -0,0 +1,6 @@ +# Optional codec-output watermarking. resemble-perth 1.0.1 does not declare +# Requires-Dist, so list the packages the watermark path actually imports. +# Do not pin librosa/torchaudio transitives; let pip resolve them. +resemble-perth==1.0.1 +librosa>=0.11.0 +torchaudio diff --git a/tools/easymagpie_vllm_omni/requirements.txt b/tools/easymagpie_vllm_omni/requirements.txt index 58f929fad049..7a87192ce2dc 100644 --- a/tools/easymagpie_vllm_omni/requirements.txt +++ b/tools/easymagpie_vllm_omni/requirements.txt @@ -1,20 +1,9 @@ # Runtime dependencies used by EasyMagpie or the speech-serving endpoint. -audioread==3.0.1 -decorator==5.2.1 -joblib==1.5.1 -lazy-loader==0.4 -librosa==0.11.0 numpy>=1.26 -pooch==1.8.2 PyYAML>=6.0 -resemble-perth==1.0.1 safetensors>=0.8.0 -scikit-learn==1.7.1 -soxr==0.5.0.post1 -threadpoolctl==3.6.0 -torchaudio soundfile>=0.13.1 tokenizers transformers>=5.5.3 vllm==0.24.0 # vLLM and vLLM-Omni must use the same major/minor version. -vllm-omni==0.24.0 # vLLM and vLLM-Omni must use the same major/minor version. +vllm-omni==0.24.0 # vLLM and vLLM-Omni must use the same major/minor version. \ No newline at end of file From b5ade35f3149f840d89fe54d6e96f617a79d5f77 Mon Sep 17 00:00:00 2001 From: Yash Hayaran Date: Wed, 9 Sep 2026 00:51:57 +0530 Subject: [PATCH 3/4] feat(watermark): add tapering function to smooth watermark perturbations at chunk edges This update introduces the `_taper_watermark_delta` function, which applies a tapering effect to the watermark perturbations at the edges of independently processed chunks. This enhancement aims to preserve the codec waveform at both edges, preventing audible seams during the concatenation of streaming chunks. The `watermark_waveforms` function has been updated to utilize this new tapering method, improving the overall audio quality of the watermarked outputs. --- .../easymagpie_vllm_omni/watermark.py | 41 ++++++++++++++++++- 1 file changed, 40 insertions(+), 1 deletion(-) diff --git a/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/watermark.py b/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/watermark.py index 2134d55476e1..c38899cccabe 100644 --- a/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/watermark.py +++ b/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/watermark.py @@ -101,6 +101,33 @@ def _restore_length(watermarked: torch.Tensor, original: torch.Tensor) -> torch. return watermarked +def _taper_watermark_delta( + marked: torch.Tensor, + original: torch.Tensor, + *, + edge_samples: int, +) -> torch.Tensor: + """Fade only Perth's perturbation at independently processed chunk edges.""" + samples = int(original.shape[-1]) + edge = min(max(int(edge_samples), 0), samples // 2) + if edge == 0: + return marked + + ramp = torch.sin( + torch.linspace( + 0, + torch.pi / 2, + edge, + device=marked.device, + dtype=marked.dtype, + ) + ).square() + envelope = torch.ones_like(marked) + envelope[..., :edge] = ramp + envelope[..., -edge:] = ramp.flip(0) + return original + (marked - original) * envelope + + def watermark_waveforms( waveforms: list[torch.Tensor], *, @@ -121,6 +148,13 @@ def watermark_waveforms( perth_device = perth_net.device perth_rate = int(perth_net.hp.sample_rate) minimum_samples = int(perth_net.hp.n_fft // 2 + 1) + # Perth reconstructs each chunk independently with an STFT. Tapering its + # perturbation across one half-window preserves the codec waveform at both + # edges and prevents audible seams when streaming chunks are concatenated. + edge_samples = max( + (int(perth_net.hp.n_fft) // 2 * sample_rate + perth_rate - 1) // perth_rate, + 1, + ) resample = None if sample_rate != perth_rate: @@ -161,7 +195,12 @@ def watermark_waveforms( if sample_rate != perth_rate: assert resample is not None marked = resample(marked, perth_rate, sample_rate) - marked = _restore_length(marked, originals).clamp_(-1.0, 1.0) + marked = _restore_length(marked, originals) + marked = _taper_watermark_delta( + marked, + originals, + edge_samples=edge_samples, + ).clamp_(-1.0, 1.0) for batch_index, waveform_index in enumerate(indices): waveforms[waveform_index] = marked[batch_index].reshape(-1) From 010245f388cdea22716e60ff4665b70d4c9b9ce5 Mon Sep 17 00:00:00 2001 From: Yash Hayaran Date: Wed, 9 Sep 2026 11:48:57 +0530 Subject: [PATCH 4/4] feat(watermark): enhance watermarking functionality and update dependencies This commit introduces several improvements to the watermarking functionality, including the addition of a new test for preserving chunk edge samples and updates to existing tests for better accuracy. The Dockerfile and requirements files have been updated to include a new `requirements-perth.txt` for optional dependencies, ensuring that the serving loader uses only necessary packages. The README has also been updated to clarify the use of dependencies and their impact on watermarking. --- .../serving/test_watermark.py | 61 +++++++++++++------ tools/easymagpie_vllm_omni/Dockerfile.ci | 2 + tools/easymagpie_vllm_omni/README.md | 3 +- .../easymagpie_vllm_omni/watermark.py | 51 +++++++++++++--- tools/easymagpie_vllm_omni/pyproject.toml | 1 - .../requirements-perth.txt | 5 +- 6 files changed, 94 insertions(+), 29 deletions(-) diff --git a/tests/collections/tts/easymagpie_vllm_omni/serving/test_watermark.py b/tests/collections/tts/easymagpie_vllm_omni/serving/test_watermark.py index 681e3aadb662..497b92b18e71 100644 --- a/tests/collections/tts/easymagpie_vllm_omni/serving/test_watermark.py +++ b/tests/collections/tts/easymagpie_vllm_omni/serving/test_watermark.py @@ -3,6 +3,7 @@ from __future__ import annotations +import sys import types from unittest import mock @@ -64,8 +65,23 @@ def test_equal_length_waveforms_are_batched_and_length_is_preserved() -> None: assert perth_net.encoder.batch_sizes == [2, 1] assert [tensor.numel() for tensor in result] == [16, 16, 12] - torch.testing.assert_close(result[0], torch.full((16,), 0.125)) - torch.testing.assert_close(result[1], torch.full((16,), 0.625)) + ramp = torch.sin(torch.linspace(0, torch.pi / 2, 4)).square() + envelope = torch.ones(16) + envelope[:4] = ramp + envelope[-4:] = ramp.flip(0) + torch.testing.assert_close(result[0], 0.125 * envelope) + torch.testing.assert_close(result[1], 0.5 + 0.125 * envelope) + + +def test_watermark_delta_taper_preserves_chunk_edge_samples() -> None: + original = torch.linspace(-0.5, 0.5, 16).repeat(2, 1) + marked = original + 0.125 + + result = watermark._taper_watermark_delta(marked, original, edge_samples=4) + + torch.testing.assert_close(result[:, 0], original[:, 0]) + torch.testing.assert_close(result[:, -1], original[:, -1]) + torch.testing.assert_close(result[:, 4:-4], marked[:, 4:-4]) def test_short_waveform_is_left_unchanged() -> None: @@ -101,12 +117,11 @@ def test_explicit_disable_bypasses_initialization() -> None: def test_enabled_initialization_failure_is_fatal() -> None: - broken_perth = types.SimpleNamespace( - PerthImplicitWatermarker=mock.Mock(side_effect=ValueError("bad checkpoint")) - ) with ( mock.patch.dict("os.environ", {"NEMOTRON_TTS_PERTH_WATERMARK": "1"}), - mock.patch.dict("sys.modules", {"perth": broken_perth}), + mock.patch.object( + watermark, "_load_perth_net", side_effect=ValueError("bad checkpoint") + ), pytest.raises(RuntimeError, match="could not be initialized"), ): watermark.initialize_watermarker("cpu") @@ -120,15 +135,11 @@ def test_encoder_is_compiled_on_cuda() -> None: device="cuda:0", ap=_FakeAudioProcessor(), ) - watermarker = types.SimpleNamespace(perth_net=perth_net) compiled = mock.Mock(name="compiled_encoder", side_effect=lambda mag: (mag, None)) - fake_perth = types.SimpleNamespace( - PerthImplicitWatermarker=mock.Mock(return_value=watermarker) - ) with ( mock.patch.dict("os.environ", {"NEMOTRON_TTS_PERTH_WATERMARK": "1"}), - mock.patch.dict("sys.modules", {"perth": fake_perth}), + mock.patch.object(watermark, "_load_perth_net", return_value=perth_net), mock.patch.object(watermark, "_normalized_device", return_value="cuda:0"), mock.patch.object(torch, "compile", return_value=compiled) as compile_fn, mock.patch.object(torch.cuda, "is_available", return_value=False), @@ -146,17 +157,33 @@ def test_encoder_is_not_compiled_on_cpu() -> None: encoder=encoder, device="cpu", ) - watermarker = types.SimpleNamespace(perth_net=perth_net) - fake_perth = types.SimpleNamespace( - PerthImplicitWatermarker=mock.Mock(return_value=watermarker) - ) - with ( mock.patch.dict("os.environ", {"NEMOTRON_TTS_PERTH_WATERMARK": "1"}), - mock.patch.dict("sys.modules", {"perth": fake_perth}), + mock.patch.object(watermark, "_load_perth_net", return_value=perth_net), mock.patch.object(torch, "compile") as compile_fn, ): result = watermark.initialize_watermarker("cpu") compile_fn.assert_not_called() assert result.perth_net.encoder is encoder + + +def test_perth_net_namespace_does_not_execute_librosa_soxr_init( + tmp_path, +) -> None: + perth_root = tmp_path / "perth" + perth_net = perth_root / "perth_net" / "pretrained" + perth_net.mkdir(parents=True) + (perth_root / "__init__.py").write_text("") + (perth_root / "perth_net" / "__init__.py").write_text("import soxr\n") + fake_perth = types.ModuleType("perth") + fake_perth.__file__ = str(perth_root / "__init__.py") + + with mock.patch.dict(sys.modules, {"perth": fake_perth}): + sys.modules.pop("perth.perth_net", None) + sys.modules.pop("soxr", None) + watermark._ensure_perth_net_namespace() + assert "soxr" not in sys.modules + assert sys.modules["perth.perth_net"].__path__ == [ + str(perth_root / "perth_net") + ] diff --git a/tools/easymagpie_vllm_omni/Dockerfile.ci b/tools/easymagpie_vllm_omni/Dockerfile.ci index 114fa7cbee82..8242ab622924 100644 --- a/tools/easymagpie_vllm_omni/Dockerfile.ci +++ b/tools/easymagpie_vllm_omni/Dockerfile.ci @@ -16,9 +16,11 @@ FROM vllm/vllm-openai:v0.24.0 COPY tools/easymagpie_vllm_omni/requirements.txt /tmp/easymagpie-requirements.txt +COPY tools/easymagpie_vllm_omni/requirements-perth.txt /tmp/easymagpie-requirements-perth.txt RUN python3 -m pip install --no-cache-dir \ -r /tmp/easymagpie-requirements.txt \ + -r /tmp/easymagpie-requirements-perth.txt \ coverage \ pytest \ pytest-asyncio \ diff --git a/tools/easymagpie_vllm_omni/README.md b/tools/easymagpie_vllm_omni/README.md index 6d3f7dc04259..f0b133c96c94 100644 --- a/tools/easymagpie_vllm_omni/README.md +++ b/tools/easymagpie_vllm_omni/README.md @@ -76,7 +76,8 @@ preserving each output's original sample count. Watermarking is enabled by default after installing the `perth` extra, and engine startup fails if Perth or its bundled checkpoint cannot be loaded. `resemble-perth==1.0.1` does not declare its own dependencies, so the extra -also installs `librosa` and `torchaudio` unpinned. For unmarked quality +also installs `torchaudio` unpinned. The serving loader uses `PerthNet` +directly and does not import `librosa`/`soxr`. For unmarked quality comparisons only, set `NEMOTRON_TTS_PERTH_WATERMARK=0` to disable it. ### Quick start — offline synthesis diff --git a/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/watermark.py b/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/watermark.py index c38899cccabe..4c7c88f2c0ca 100644 --- a/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/watermark.py +++ b/tools/easymagpie_vllm_omni/easymagpie_vllm_omni/watermark.py @@ -7,6 +7,9 @@ import logging import os +import sys +import types +from pathlib import Path from typing import Any import torch @@ -33,6 +36,45 @@ def watermarking_enabled() -> bool: return value.strip().lower() not in _DISABLED_VALUES +def _perth_pretrained_dir() -> Path: + import perth + + return Path(perth.__file__).resolve().parent / "perth_net" / "pretrained" + + +def _ensure_perth_net_namespace() -> Path: + """Register ``perth.perth_net`` without running its package ``__init__``. + + ``perth.perth_net`` and ``PerthImplicitWatermarker`` import + ``librosa.resample``, which eagerly imports ``soxr`` (LGPLv2.1+). This + serving path only needs ``PerthNet`` plus torchaudio resampling. + """ + import perth + + pretrained_dir = _perth_pretrained_dir() + package_name = "perth.perth_net" + package_path = str(Path(perth.__file__).resolve().parent / "perth_net") + current = sys.modules.get(package_name) + if current is not None and getattr(current, "__path__", None): + return pretrained_dir + + module = types.ModuleType(package_name) + module.__path__ = [package_path] + module.__file__ = str(Path(package_path) / "__init__.py") + module.__package__ = package_name + module.PREPACKAGED_MODELS_DIR = str(pretrained_dir) + sys.modules[package_name] = module + return pretrained_dir + + +def _load_perth_net(device: str) -> Any: + pretrained_dir = _ensure_perth_net_namespace() + from perth.perth_net.perth_net_implicit.model.perth_net import PerthNet + + perth_net = PerthNet.load("implicit", models_dir=str(pretrained_dir)) + return perth_net.to(device).eval() + + def initialize_watermarker(device: str | torch.device) -> Any | None: """Load Perth on ``device`` or fail startup when watermarking is enabled.""" global _INITIALIZED_DEVICE, _WATERMARKER @@ -51,13 +93,8 @@ def initialize_watermarker(device: str | torch.device) -> Any | None: return _WATERMARKER try: - import perth - - watermarker_type = perth.PerthImplicitWatermarker - if watermarker_type is None: - raise ImportError("perth.PerthImplicitWatermarker is unavailable") - _WATERMARKER = watermarker_type(device=requested_device) - perth_net = _WATERMARKER.perth_net + perth_net = _load_perth_net(requested_device) + _WATERMARKER = types.SimpleNamespace(perth_net=perth_net) if torch.device(requested_device).type == "cuda": perth_net.encoder = torch.compile( perth_net.encoder, mode="default", dynamic=True diff --git a/tools/easymagpie_vllm_omni/pyproject.toml b/tools/easymagpie_vllm_omni/pyproject.toml index f27641829db9..47498d142f6c 100644 --- a/tools/easymagpie_vllm_omni/pyproject.toml +++ b/tools/easymagpie_vllm_omni/pyproject.toml @@ -16,7 +16,6 @@ dependencies = [ [project.optional-dependencies] perth = [ "resemble-perth==1.0.1", - "librosa>=0.11.0", "torchaudio", ] diff --git a/tools/easymagpie_vllm_omni/requirements-perth.txt b/tools/easymagpie_vllm_omni/requirements-perth.txt index 0d48d7e9293f..68bc9e8669d8 100644 --- a/tools/easymagpie_vllm_omni/requirements-perth.txt +++ b/tools/easymagpie_vllm_omni/requirements-perth.txt @@ -1,6 +1,5 @@ # Optional codec-output watermarking. resemble-perth 1.0.1 does not declare -# Requires-Dist, so list the packages the watermark path actually imports. -# Do not pin librosa/torchaudio transitives; let pip resolve them. +# Requires-Dist. The serving loader uses PerthNet + torchaudio only; it does +# not import librosa/soxr (LGPLv2.1+). resemble-perth==1.0.1 -librosa>=0.11.0 torchaudio