{
 "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
}