File size: 4,995 Bytes
1fdc661 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | {
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# VocBulwark watermark round-trip\n",
"\n",
"The full pipeline, in three self-contained models loaded from the Hub:\n",
"\n",
"1. **Speaker encoder** — a reference clip → a fixed-size speaker embedding (*whose* voice).\n",
"2. **Vocoder** — a mel spectrogram (*what* is said) + that embedding → a 24 kHz waveform, with the fixed provenance watermark embedded automatically.\n",
"3. **Detector** — the waveform → how many watermark bits match.\n",
"\n",
"In a real TTS system you usually already have the mel (from an acoustic model) and the speaker embedding (precomputed once per speaker), so step 1 is only needed to *make* an embedding from audio. Here we reconstruct one clip (use it as both content and speaker reference) to exercise the watermark.\n",
"\n",
"Needs only `torch`, `transformers`, `soundfile`, `torchaudio`. Put one or more `.wav` files in an `example_audio/` folder next to this notebook. If the repos are private, run `huggingface-cli login` first."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from pathlib import Path\n",
"\n",
"SPK_REPO = \"mlr2000/vocoder-small-speaker-encoder\" # <- speaker encoder\n",
"GEN_REPO = \"mlr2000/vocoder-small\" # <- vocoder (paired with SPK_REPO)\n",
"DET_REPO = \"mlr2000/vocoder-small-watermark-detector\" # <- matching detector\n",
"TGT_SR = 24000 # vocoder output SR\n",
"MAX_SAMPLES = 16 * TGT_SR\n",
"\n",
"AUDIO_DIR = next(p for p in [Path(\"example_audio\"), Path(\"../example_audio\")] if p.exists())"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"from transformers import AutoModel\n",
"\n",
"enc = AutoModel.from_pretrained(SPK_REPO, trust_remote_code=True).eval()\n",
"gen = AutoModel.from_pretrained(GEN_REPO, trust_remote_code=True).eval()\n",
"det = AutoModel.from_pretrained(DET_REPO, trust_remote_code=True).eval()\n",
"RAW_SR = enc.config.raw_sample_rate # SR the speaker encoder expects\n",
"print(f\"embedding dim: {enc.config.embedding_size} | watermark: {len(gen.config.fixed_watermark)} bits\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import soundfile as sf\n",
"import torchaudio.functional as AF\n",
"from transformers import WhisperFeatureExtractor\n",
"\n",
"# Mel front-end — must match the trained vocoder (n_fft=1024, hop=256, mel channels from config).\n",
"N_MELS = gen.config.hifigan_in_channels\n",
"mel_fe = WhisperFeatureExtractor(sampling_rate=TGT_SR, n_fft=1024, feature_size=N_MELS, hop_length=256)\n",
"mel_fe.n_samples = MAX_SAMPLES\n",
"mel_fe.chunk_length = MAX_SAMPLES / TGT_SR\n",
"\n",
"def load_inputs(path):\n",
" \"\"\"A clip -> (mel = content, raw = speaker reference @ RAW_SR).\"\"\"\n",
" w, sr = sf.read(str(path))\n",
" w = torch.tensor(w, dtype=torch.float32)\n",
" if w.dim() == 2:\n",
" w = w.mean(1) # mono\n",
" raw = AF.resample(w, sr, RAW_SR) if sr != RAW_SR else w # speaker ref @ RAW_SR\n",
" tgt = AF.resample(raw, RAW_SR, TGT_SR) # content @ 24 kHz\n",
" mel = mel_fe(tgt.numpy(), sampling_rate=TGT_SR, padding=\"longest\",\n",
" return_tensors=\"pt\")[\"input_features\"]\n",
" return mel, raw.unsqueeze(0)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from IPython.display import Audio, display\n",
"\n",
"for path in sorted(AUDIO_DIR.glob(\"*.wav\")):\n",
" mel, raw = load_inputs(path)\n",
" with torch.no_grad():\n",
" emb = enc.embed(raw) # [1, embedding_size]\n",
" # The vocoder embeds its fixed watermark automatically.\n",
" audio = gen(mel_spectrogram=mel, speaker_embedding=emb).audio.squeeze(1)\n",
" r = det.detect(audio)\n",
" print(f\"{path.name}: {r['matches']}/{r['n_bits']} bits matched p={r['p_value']:.1e}\")\n",
" display(Audio(audio[0].numpy(), rate=TGT_SR))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Sanity check: an **unrelated** clip (here, random noise) is not from our model, so it matches only about half the bits with a large `p_value`."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(\"random noise:\", det.detect(torch.randn(1, TGT_SR)))"
]
}
],
"metadata": {
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
|