Upload folder using huggingface_hub

README.md

@@ -0,0 +1,91 @@
+---
+library_name: transformers
+tags:
+  - audio
+  - watermark
+  - watermark-detection
+  - provenance
+  - vocbulwark
+license: cc-by-4.0
+---
+
+# VocBulwark watermark detector
+
+Standalone detector for the fixed **32-bit** provenance watermark embedded by
+the companion VocBulwark vocoder. Given an audio clip, it extracts the watermark
+bits with the Cage extractor and compares them to the known fixed code, reporting
+how many bits match.
+
+Self-contained: loads with `trust_remote_code=True`, no training repo required.
+
+## Companion Models
+
+This detector is part of a set of 6 repositories:
+
+| Repo | Role |
+|------|------|
+| `mlr2000/vocoder-large` | Large vocoder |
+| `mlr2000/vocoder-large-watermark-detector` | Watermark detector for the large model |
+| `mlr2000/vocoder-large-speaker-encoder` | Speaker encoder for the large model |
+| `mlr2000/vocoder-small` | Small vocoder (generates the watermarked audio) |
+| `mlr2000/vocoder-small-watermark-detector` | Watermark detector (this repo) |
+| `mlr2000/vocoder-small-speaker-encoder` | Speaker encoder for the small model |
+
+## Usage
+
+```python
+import torchaudio
+from transformers import AutoModel
+
+det = AutoModel.from_pretrained("mlr2000/vocoder-small-watermark-detector", trust_remote_code=True).eval()
+
+wav, sr = torchaudio.load("clip.wav")          # [C, T]
+res = det.detect(wav, input_sample_rate=sr)    # resampled to 24 kHz internally
+print(res)
+# {'matches': 32, 'n_bits': 32, 'p_value': 8.9e-16}   # from our model
+```
+
+See **`example_roundtrip.ipynb`** in this repo for an end-to-end example
+(generate with the vocoder → detect here).
+
+## How detection works
+
+`detect()` returns a dict with:
+
+- `matches`: how many of the 32 extracted bits equal the fixed code.
+- `n_bits`: 32.
+- `p_value`: the probability an *unrelated* clip matches at least this well under
+  `Binomial(32, 0.5)`.
+
+A clip generated by our vocoder matches **all 32 bits** (`p_value` ~ 0); an
+unrelated clip matches about half of them (`p_value` ~ 1). There is no built-in
+yes/no threshold — read `matches` / `p_value` and pick whatever operating point
+you need. Requiring all 32 bits gives an astronomically small false-positive
+rate; allowing a few mismatches trades that for robustness to lossy channels.
+
+## Notes
+
+- Detection runs on **24 kHz mono**. Other inputs are resampled when you pass `input_sample_rate`, and multi-channel audio is downmixed to mono.
+- This detector is paired with `mlr2000/vocoder-small`. It will not correctly verify audio generated by the large vocoder (`mlr2000/vocoder-large`), which uses a different fixed 50-bit watermark.
+- A fixed public watermark is for **attribution** ("did our model make this"), not a secret mark.
+- Not intended for surveillance or any use that violates the privacy of individuals.
+
+## Citation
+
+If you use this model, please cite:
+
+```bibtex
+@misc{muletta2026,
+  title  = {Training a Discriminator-Free Foundation Vocoder 
+             with Integrated Audio Watermarking},
+  author = {Muletta, Romolo and Deriu, Jan},
+  year   = {2026},
+  note   = {VT2 Project Report, ZHAW School of Engineering}
+}
+```
+
+## License
+
+`cc-by-4.0`. Trained on MLS (CC-BY-4.0) and Common Voice (CC0). Builds 
+on BigVGAN (MIT) and wav2vec 2.0 (Apache-2.0). Please retain attribution 
+when redistributing or building on this model.

__init__.py

@@ -0,0 +1,4 @@
+from .configuration_detector import CageDetectorConfig
+from .modeling_detector import CageDetector, DetectionOutput
+
+__all__ = ["CageDetectorConfig", "CageDetector", "DetectionOutput"]

cage_config.py

@@ -0,0 +1,23 @@
+from transformers import PretrainedConfig
+
+
+class CageExtractorConfig(PretrainedConfig):
+    model_type = "cage_extractor"
+
+    def __init__(
+        self,
+        watermark_bits: int = 100,
+        in_channels: int = 1,
+        fine_kernel: int = 3,
+        mid_kernel: int = 5,
+        coarse_kernel: int = 7,
+        sample_rate: int = 24000,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.watermark_bits = watermark_bits
+        self.in_channels = in_channels
+        self.fine_kernel = fine_kernel
+        self.mid_kernel = mid_kernel
+        self.coarse_kernel = coarse_kernel
+        self.sample_rate = sample_rate

cage_extractor.py

@@ -0,0 +1,138 @@
+"""
+Coarse-to-Fine Gated Extractor (Cage) from VocBulwark (Appendix B.2).
+
+Architecture from paper:
+  - Three branches (fine/mid/coarse), each with 4 GSCMs
+  - Channel progression: 1 → 32 → 64 → 128 → 128
+  - Fine: k=3, s=2, p=1; Mid: k=5, s=2, p=2; Coarse: k=7, s=2, p=3
+  - Each GSCM has InstanceNorm + LeakyReLU
+  - Each branch output → 1D adaptive average pooling
+  - Fusion: GSCM(384→256, k=1, s=1, p=0) → GSCM(256, k=3, s=1, p=1)
+  - Dual-path pooling (AAP + AMP) → arithmetic mean
+  - Decoder: Linear(256, 512) → Linear(512, watermark_bits)
+"""
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+
+from .cage_config import CageExtractorConfig
+
+
+class GatedSeparableConvModule(nn.Module):
+    """GSCM: Gated Separable Convolution Module (Sec 4.3 + Appendix B.2).
+
+    Depthwise separable conv with gating mechanism.
+    Uses InstanceNorm and LeakyReLU as per paper.
+    Supports stride and different input/output channels.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 3,
+        stride: int = 1,
+        padding: int = 1,
+    ):
+        super().__init__()
+
+        # Content branch (DSC): depthwise → pointwise → InstanceNorm → LeakyReLU
+        self.content = nn.Sequential(
+            nn.Conv1d(in_channels, in_channels, kernel_size,
+                      stride=stride, padding=padding, groups=in_channels),
+            nn.Conv1d(in_channels, out_channels, 1),
+            nn.InstanceNorm1d(out_channels),
+            nn.LeakyReLU(0.2),
+        )
+
+        # Gating branch: same structure but Sigmoid output
+        self.gate = nn.Sequential(
+            nn.Conv1d(in_channels, in_channels, kernel_size,
+                      stride=stride, padding=padding, groups=in_channels),
+            nn.Conv1d(in_channels, out_channels, 1),
+            nn.InstanceNorm1d(out_channels),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.content(x) * self.gate(x)
+
+
+def _make_branch(kernel_size: int, stride: int, padding: int):
+    """Build a 4-GSCM branch with channel progression 1 → 32 → 64 → 128 → 128."""
+    return nn.Sequential(
+        GatedSeparableConvModule(1, 32, kernel_size, stride, padding),
+        GatedSeparableConvModule(32, 64, kernel_size, stride, padding),
+        GatedSeparableConvModule(64, 128, kernel_size, stride, padding),
+        GatedSeparableConvModule(128, 128, kernel_size, stride, padding),
+    )
+
+
+class CageExtractor(PreTrainedModel):
+    config_class = CageExtractorConfig
+
+    def __init__(self, config: CageExtractorConfig):
+        super().__init__(config)
+
+        # Three branches processing raw audio directly (paper B.2)
+        # Fine-grained: k=3, s=2, p=1
+        self.fine_branch = _make_branch(
+            kernel_size=config.fine_kernel, stride=2, padding=1)
+        # Mid-grained: k=5, s=2, p=2
+        self.mid_branch = _make_branch(
+            kernel_size=config.mid_kernel, stride=2, padding=2)
+        # Coarse-grained: k=7, s=2, p=3
+        self.coarse_branch = _make_branch(
+            kernel_size=config.coarse_kernel, stride=2, padding=3)
+
+        # Fusion: concat 3 branches (3*128=384) → reduce to 256 → refine
+        # GSCM (k=1, s=1, p=0) to reduce channel dimension to 256
+        self.fusion_reduce = GatedSeparableConvModule(
+            384, 256, kernel_size=1, stride=1, padding=0)
+        # GSCM (k=3, s=1, p=1) that maintains the dimensions
+        self.fusion_refine = GatedSeparableConvModule(
+            256, 256, kernel_size=3, stride=1, padding=1)
+
+        # Decoder: two FC layers (paper B.2)
+        self.decoder = nn.Sequential(
+            nn.Linear(256, 512),
+            nn.LeakyReLU(0.2),
+            nn.Linear(512, config.watermark_bits),
+        )
+
+    def forward(self, audio: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            audio: [B, 1, T] waveform (or [B, T] which gets unsqueezed)
+        Returns:
+            logits: [B, watermark_bits] raw logits
+        """
+        if audio.dim() == 2:
+            audio = audio.unsqueeze(1)
+
+        # Multi-scale branches (each processes raw audio)
+        x_fine = self.fine_branch(audio)      # [B, 128, T']
+        x_mid = self.mid_branch(audio)        # [B, 128, T'']
+        x_coarse = self.coarse_branch(audio)  # [B, 128, T''']
+
+        # 1D adaptive average pooling to standardize temporal dimension
+        pool_t = min(x_fine.shape[-1], x_mid.shape[-1], x_coarse.shape[-1])
+        x_fine = nn.functional.adaptive_avg_pool1d(x_fine, pool_t)
+        x_mid = nn.functional.adaptive_avg_pool1d(x_mid, pool_t)
+        x_coarse = nn.functional.adaptive_avg_pool1d(x_coarse, pool_t)
+
+        # Fusion
+        x_fuse = torch.cat([x_fine, x_mid, x_coarse], dim=1)  # [B, 384, pool_t]
+        x_fuse = self.fusion_reduce(x_fuse)   # [B, 256, pool_t]
+        x_fuse = self.fusion_refine(x_fuse)   # [B, 256, pool_t]
+
+        # Dual-path Pooling: AAP + AMP → arithmetic mean
+        x_aap = nn.functional.adaptive_avg_pool1d(x_fuse, 1).squeeze(-1)  # [B, 256]
+        x_amp = nn.functional.adaptive_max_pool1d(x_fuse, 1).squeeze(-1)  # [B, 256]
+        x_pool = (x_aap + x_amp) / 2
+
+        # Decode
+        logits = self.decoder(x_pool)  # [B, watermark_bits]
+
+        return logits

config.json

@@ -0,0 +1,50 @@
+{
+  "model_type": "cage_detector",
+  "architectures": [
+    "CageDetector"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_detector.CageDetectorConfig",
+    "AutoModel": "modeling_detector.CageDetector"
+  },
+  "watermark_bits": 32,
+  "in_channels": 1,
+  "fine_kernel": 3,
+  "mid_kernel": 5,
+  "coarse_kernel": 7,
+  "sample_rate": 24000,
+  "fixed_watermark": [
+    1,
+    0,
+    1,
+    1,
+    0,
+    0,
+    1,
+    0,
+    1,
+    1,
+    0,
+    1,
+    0,
+    0,
+    1,
+    1,
+    0,
+    1,
+    1,
+    0,
+    1,
+    0,
+    0,
+    1,
+    1,
+    1,
+    0,
+    0,
+    1,
+    0,
+    1,
+    1
+  ]
+}

configuration_detector.py

@@ -0,0 +1,11 @@
+"""Config for the standalone Cage watermark detector."""
+from .cage_config import CageExtractorConfig
+
+
+class CageDetectorConfig(CageExtractorConfig):
+    model_type = "cage_detector"
+
+    def __init__(self, fixed_watermark=None, **kwargs):
+        super().__init__(**kwargs)
+        # The fixed signature this detector checks extracted bits against.
+        self.fixed_watermark = fixed_watermark

example_roundtrip.ipynb

@@ -0,0 +1,125 @@
+{
+ "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
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1e9d11e749c65af36893011c5871409ffd5f74e3a42ef72c7387a3daac0bfb9
+size 2614304

modeling_detector.py

@@ -0,0 +1,91 @@
+"""Standalone VocBulwark Cage watermark detector.
+
+Loadable with `AutoModel.from_pretrained(path, trust_remote_code=True)` without
+the training repository. Wraps the Cage extractor and compares the extracted
+bits to a fixed provenance watermark baked into the config.
+"""
+from dataclasses import dataclass
+from math import comb
+
+import torch
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+from .configuration_detector import CageDetectorConfig
+from .cage_config import CageExtractorConfig
+from .cage_extractor import CageExtractor
+
+
+@dataclass
+class DetectionOutput(ModelOutput):
+    bits: torch.Tensor = None
+    logits: torch.Tensor = None
+
+
+class CageDetector(PreTrainedModel):
+    config_class = CageDetectorConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        ext_cfg = CageExtractorConfig(
+            watermark_bits=config.watermark_bits,
+            in_channels=config.in_channels,
+            fine_kernel=config.fine_kernel,
+            mid_kernel=config.mid_kernel,
+            coarse_kernel=config.coarse_kernel,
+            sample_rate=config.sample_rate,
+        )
+        self.cage_extractor = CageExtractor(ext_cfg)
+
+    def _prep(self, audio, input_sample_rate):
+        """Coerce input to [B, T] mono at the detector sample rate."""
+        if not torch.is_tensor(audio):
+            audio = torch.as_tensor(audio)
+        p = next(self.parameters())
+        audio = audio.to(device=p.device, dtype=p.dtype)
+        if audio.dim() == 1:                       # [T] -> [1, T]
+            audio = audio.unsqueeze(0)
+        elif audio.dim() == 3:                     # [B, C, T] -> [B, T]
+            audio = audio.mean(dim=1) if audio.shape[1] > 1 else audio.squeeze(1)
+        if input_sample_rate is not None and input_sample_rate != self.config.sample_rate:
+            import torchaudio
+            audio = torchaudio.functional.resample(
+                audio, input_sample_rate, self.config.sample_rate)
+        return audio
+
+    @staticmethod
+    def _p_value(matches, n):
+        """P(Binomial(n, 0.5) >= matches): chance an unrelated clip matches this well."""
+        return sum(comb(n, i) for i in range(int(matches), n + 1)) / (2 ** n)
+
+    @torch.no_grad()
+    def extract(self, audio, input_sample_rate=None):
+        """Return (bits, logits). bits: [B, n] in {0,1}; logits: [B, n]."""
+        audio = self._prep(audio, input_sample_rate)
+        logits = self.cage_extractor(audio)
+        bits = (logits > 0).to(torch.int64)
+        return bits, logits
+
+    @torch.no_grad()
+    def detect(self, audio, input_sample_rate=None):
+        """Detect the fixed watermark. Returns a dict (or list of dicts for B>1):
+        {matches, n_bits, p_value}. `matches` is how many of the model's fixed
+        bits the clip carries; `p_value` is the chance an unrelated clip matches
+        at least this well under Binomial(n_bits, 0.5). A clip from our model
+        matches all n_bits (p_value ~ 0); an unrelated clip matches about half."""
+        bits, _ = self.extract(audio, input_sample_rate)
+        wstar = torch.tensor(self.config.fixed_watermark,
+                             device=bits.device, dtype=bits.dtype)
+        matches = (bits == wstar).sum(dim=-1).tolist()
+        n = int(self.config.watermark_bits)
+        results = [{
+            "matches": int(m),
+            "n_bits": n,
+            "p_value": self._p_value(m, n),
+        } for m in matches]
+        return results[0] if len(results) == 1 else results
+
+    @torch.no_grad()
+    def forward(self, audio=None, input_sample_rate=None, **kwargs):
+        bits, logits = self.extract(audio, input_sample_rate)
+        return DetectionOutput(bits=bits, logits=logits)