Upload folder using huggingface_hub

README.md

@@ -0,0 +1,114 @@
+---
+library_name: transformers
+tags:
+  - audio
+  - text-to-speech
+  - vocoder
+  - hifi-gan
+  - bigvgan
+  - neural-vocoder
+  - speaker-conditioning
+  - audio-watermarking
+pipeline_tag: text-to-speech
+license: cc-by-4.0
+---
+
+# VocBulwark HiFi-GAN — watermarking neural vocoder (inference export)
+
+Speaker-conditioned **BigVGAN / HiFi-GAN** neural vocoder. It turns an input mel-spectrogram (*what* is said) into a 24 kHz waveform, conditioned on a **precomputed 768-d speaker embedding** (*whose* voice). Every clip it generates carries a **fixed 50-bit provenance watermark** that identifies this specific model instance, see [Watermark](#watermark).
+
+This is the **lean, inference-only** vocoder: the frozen perceptual-loss base models (Whisper / WavLM / Wav2Vec2), the training discriminators, **and the speaker encoder** have all been stripped, you pass the speaker embedding in. Use the companion **speaker-encoder** repo to turn a reference clip into that embedding. The modeling code is bundled, so it loads with `trust_remote_code=True` **without** the training repo.
+
+## Model summary
+
+| | |
+|---|---|
+| Architecture | `HiFiGANArchitecture` (BigVGAN generator, snakebeta activation) |
+| Inputs | log-mel spectrogram (**96 mel channels**) + speaker embedding (**768-d**) |
+| Output | mono waveform, **24 kHz** |
+| Speaker conditioning | precomputed embedding (from the companion speaker encoder) |
+| Generator | initial channels 1536, upsample rates [4, 4, 2, 2, 2, 2] |
+| Watermark | 50-bit fixed VocBulwark signature, always embedded |
+| Framework | 🤗 Transformers, PyTorch, safetensors |
+
+
+## Companion Models
+
+This model is part of a set of 6 repositories:
+
+| Repo | Role |
+|------|------|
+| `mlr2000/vocoder-large` | Large vocoder (this repo) |
+| `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 |
+| `mlr2000/vocoder-small-watermark-detector` | Watermark detector for the small model |
+| `mlr2000/vocoder-small-speaker-encoder` | Speaker encoder for the small model |
+
+
+## Usage
+
+```python
+import torch
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained("mlr2000/vocoder-large", trust_remote_code=True).eval()
+
+mel = torch.randn(1, model.config.hifigan_in_channels, 200)      # [B, mel, T]
+emb = torch.randn(1, model.config.speaker_embedding_size)        # [B, 768] from the speaker encoder
+with torch.no_grad():
+    audio = model(mel_spectrogram=mel, speaker_embedding=emb).audio
+# audio: [B, 1, samples] @ model.config.target_sample_rate
+```
+
+See **`example_roundtrip.ipynb`** in this repo for the full pipeline (reference
+clip → speaker encoder → embedding → vocode → verify watermark).
+
+## Training
+
+| | |
+|---|---|
+| Training data | Multilingual LibriSpeech (8 languages, ~22,200h) and Common Voice (14 languages, ~3,000h) |
+| Training steps | 1,000,000 |
+| Hardware | 2 × NVIDIA H200 GPUs |
+| Training objective | Discriminator-free: mel spectrogram + WavLM + wav2vec 2.0 + Whisper encoder losses |
+| Effective batch size | 32 |
+| Learning rate | 1e-4 |
+
+
+## Watermark
+
+Every clip this model generates carries a **fixed 50-bit provenance watermark** (`config.fixed_watermark`) that identifies this specific model instance. It is embedded automatically inside `forward` and **cannot be disabled or changed** through this interface — there is deliberately no watermark argument to override.
+
+To verify whether a given audio clip was generated by this model, use the companion **detector** repo (`mlr2000/vocoder-large-watermark-detector`), which extracts the embedded bits and compares them to the known fixed code.
+
+
+## Notes
+
+- Inputs: log-mel spectrogram (`config.hifigan_in_channels` channels) and a
+  `[B, config.speaker_embedding_size]` speaker embedding.
+- Output: mono waveform at `config.target_sample_rate`.
+- Use the embedding from the speaker encoder this vocoder was **trained with**, a
+  mismatched encoder will not condition it correctly.
+- Not intended for voice cloning of real individuals without consent, or any
+  deceptive / impersonation use.
+
+## 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_hifigan import HiFiGANConfig
+from .modeling_hifigan import HiFiGANArchitecture, HiFiGANOutput
+
+__all__ = ["HiFiGANConfig", "HiFiGANArchitecture", "HiFiGANOutput"]

alias_free_act.py

@@ -0,0 +1,30 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+import torch.nn as nn
+from .alias_free_resample import UpSample1d, DownSample1d
+
+
+class Activation1d(nn.Module):
+    def __init__(
+        self,
+        activation,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+    ):
+        super().__init__()
+        self.up_ratio = up_ratio
+        self.down_ratio = down_ratio
+        self.act = activation
+        self.upsample = UpSample1d(up_ratio, up_kernel_size)
+        self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+    # x: [B,C,T]
+    def forward(self, x):
+        x = self.upsample(x)
+        x = self.act(x)
+        x = self.downsample(x)
+
+        return x

alias_free_filter.py

@@ -0,0 +1,101 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+if "sinc" in dir(torch):
+    sinc = torch.sinc
+else:
+    # This code is adopted from adefossez's julius.core.sinc under the MIT License
+    # https://adefossez.github.io/julius/julius/core.html
+    #   LICENSE is in incl_licenses directory.
+    def sinc(x: torch.Tensor):
+        """
+        Implementation of sinc, i.e. sin(pi * x) / (pi * x)
+        __Warning__: Different to julius.sinc, the input is multiplied by `pi`!
+        """
+        return torch.where(
+            x == 0,
+            torch.tensor(1.0, device=x.device, dtype=x.dtype),
+            torch.sin(math.pi * x) / math.pi / x,
+        )
+
+
+# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License
+# https://adefossez.github.io/julius/julius/lowpass.html
+#   LICENSE is in incl_licenses directory.
+def kaiser_sinc_filter1d(
+    cutoff, half_width, kernel_size
+):  # return filter [1,1,kernel_size]
+    even = kernel_size % 2 == 0
+    half_size = kernel_size // 2
+
+    # For kaiser window
+    delta_f = 4 * half_width
+    A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+    if A > 50.0:
+        beta = 0.1102 * (A - 8.7)
+    elif A >= 21.0:
+        beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
+    else:
+        beta = 0.0
+    window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+
+    # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio
+    if even:
+        time = torch.arange(-half_size, half_size) + 0.5
+    else:
+        time = torch.arange(kernel_size) - half_size
+    if cutoff == 0:
+        filter_ = torch.zeros_like(time)
+    else:
+        filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
+        """
+        Normalize filter to have sum = 1, otherwise we will have a small leakage of the constant component in the input signal.
+        """
+        filter_ /= filter_.sum()
+        filter = filter_.view(1, 1, kernel_size)
+
+    return filter
+
+
+class LowPassFilter1d(nn.Module):
+    def __init__(
+        self,
+        cutoff=0.5,
+        half_width=0.6,
+        stride: int = 1,
+        padding: bool = True,
+        padding_mode: str = "replicate",
+        kernel_size: int = 12,
+    ):
+        """
+        kernel_size should be even number for stylegan3 setup, in this implementation, odd number is also possible.
+        """
+        super().__init__()
+        if cutoff < -0.0:
+            raise ValueError("Minimum cutoff must be larger than zero.")
+        if cutoff > 0.5:
+            raise ValueError("A cutoff above 0.5 does not make sense.")
+        self.kernel_size = kernel_size
+        self.even = kernel_size % 2 == 0
+        self.pad_left = kernel_size // 2 - int(self.even)
+        self.pad_right = kernel_size // 2
+        self.stride = stride
+        self.padding = padding
+        self.padding_mode = padding_mode
+        filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+        self.register_buffer("filter", filter)
+
+    # Input [B, C, T]
+    def forward(self, x):
+        _, C, _ = x.shape
+
+        if self.padding:
+            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+        out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
+
+        return out

alias_free_resample.py

@@ -0,0 +1,58 @@
+# Adapted from https://github.com/junjun3518/alias-free-torch under the Apache License 2.0
+#   LICENSE is in incl_licenses directory.
+
+import torch.nn as nn
+from torch.nn import functional as F
+from .alias_free_filter import LowPassFilter1d
+from .alias_free_filter import kaiser_sinc_filter1d
+
+
+class UpSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
+        self.stride = ratio
+        self.pad = self.kernel_size // ratio - 1
+        self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+        self.pad_right = (
+            self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        )
+        filter = kaiser_sinc_filter1d(
+            cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size
+        )
+        self.register_buffer("filter", filter)
+
+    # x: [B, C, T]
+    def forward(self, x):
+        _, C, _ = x.shape
+
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
+        x = self.ratio * F.conv_transpose1d(
+            x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C
+        )
+        x = x[..., self.pad_left : -self.pad_right]
+
+        return x
+
+
+class DownSample1d(nn.Module):
+    def __init__(self, ratio=2, kernel_size=None):
+        super().__init__()
+        self.ratio = ratio
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
+        self.lowpass = LowPassFilter1d(
+            cutoff=0.5 / ratio,
+            half_width=0.6 / ratio,
+            stride=ratio,
+            kernel_size=self.kernel_size,
+        )
+
+    def forward(self, x):
+        xx = self.lowpass(x)
+
+        return xx

bigvgan_activations.py

@@ -0,0 +1,129 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+#   Licensed under the MIT license.
+
+# Adapted from https://github.com/NVIDIA/BigVGAN under the MIT license.
+#   LICENSE is in incl_licenses directory.
+
+import torch
+from torch import nn, sin, pow
+from torch.nn import Parameter
+
+
+class Snake(nn.Module):
+    """
+    Implementation of a sine-based periodic activation function
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter
+    References:
+        - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snake(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    """
+
+    def __init__(
+        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+    ):
+        """
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha: trainable parameter
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            alpha will be trained along with the rest of your model.
+        """
+        super(Snake, self).__init__()
+        self.in_features = in_features
+
+        # Initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # Log scale alphas initialized to zeros
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+        else:  # Linear scale alphas initialized to ones
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        """
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        Snake ∶= x + 1/a * sin^2 (xa)
+        """
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # Line up with x to [B, C, T]
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+        x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x
+
+
+class SnakeBeta(nn.Module):
+    """
+    A modified Snake function which uses separate parameters for the magnitude of the periodic components
+    Shape:
+        - Input: (B, C, T)
+        - Output: (B, C, T), same shape as the input
+    Parameters:
+        - alpha - trainable parameter that controls frequency
+        - beta - trainable parameter that controls magnitude
+    References:
+        - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+        https://arxiv.org/abs/2006.08195
+    Examples:
+        >>> a1 = snakebeta(256)
+        >>> x = torch.randn(256)
+        >>> x = a1(x)
+    """
+
+    def __init__(
+        self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False
+    ):
+        """
+        Initialization.
+        INPUT:
+            - in_features: shape of the input
+            - alpha - trainable parameter that controls frequency
+            - beta - trainable parameter that controls magnitude
+            alpha is initialized to 1 by default, higher values = higher-frequency.
+            beta is initialized to 1 by default, higher values = higher-magnitude.
+            alpha will be trained along with the rest of your model.
+        """
+        super(SnakeBeta, self).__init__()
+        self.in_features = in_features
+
+        # Initialize alpha
+        self.alpha_logscale = alpha_logscale
+        if self.alpha_logscale:  # Log scale alphas initialized to zeros
+            self.alpha = Parameter(torch.zeros(in_features) * alpha)
+            self.beta = Parameter(torch.zeros(in_features) * alpha)
+        else:  # Linear scale alphas initialized to ones
+            self.alpha = Parameter(torch.ones(in_features) * alpha)
+            self.beta = Parameter(torch.ones(in_features) * alpha)
+
+        self.alpha.requires_grad = alpha_trainable
+        self.beta.requires_grad = alpha_trainable
+
+        self.no_div_by_zero = 0.000000001
+
+    def forward(self, x):
+        """
+        Forward pass of the function.
+        Applies the function to the input elementwise.
+        SnakeBeta ∶= x + 1/b * sin^2 (xa)
+        """
+        alpha = self.alpha.unsqueeze(0).unsqueeze(-1)  # Line up with x to [B, C, T]
+        beta = self.beta.unsqueeze(0).unsqueeze(-1)
+        if self.alpha_logscale:
+            alpha = torch.exp(alpha)
+            beta = torch.exp(beta)
+        x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+        return x

bigvgan_model.py

@@ -0,0 +1,361 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.
+#   Licensed under the MIT license.
+
+# Adapted from https://github.com/NVIDIA/BigVGAN under the MIT license.
+#   LICENSE is in incl_licenses directory.
+
+from typing import *
+
+import torch
+import torch.nn as nn
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import weight_norm, remove_weight_norm
+
+from . import bigvgan_activations as activations
+from .alias_free_act import Activation1d as TorchActivation1d
+from .temporal_adapter import TemporalAdapterBlock
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def apply_weight_norm(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        weight_norm(m)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+class FiLMBlock(nn.Module):
+    """Feature-wise Linear Modulation for watermark conditioning.
+
+    Maps a binary watermark vector to per-channel scale (γ) and shift (β)
+    parameters that modulate intermediate feature maps.
+    """
+
+    def __init__(self, watermark_bits: int, channels: int, hidden_dim: int = 128):
+        super().__init__()
+        self.scale_net = nn.Sequential(
+            nn.Linear(watermark_bits, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, channels),
+        )
+        self.shift_net = nn.Sequential(
+            nn.Linear(watermark_bits, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, channels),
+        )
+        # Initialize close to identity: γ≈1, β≈0
+        nn.init.zeros_(self.scale_net[-1].weight)
+        nn.init.zeros_(self.scale_net[-1].bias)
+        nn.init.zeros_(self.shift_net[-1].weight)
+        nn.init.zeros_(self.shift_net[-1].bias)
+
+    def forward(self, x, watermark):
+        """
+        Args:
+            x: [B, C, T] feature tensor
+            watermark: [B, watermark_bits] float tensor
+        Returns:
+            [B, C, T] modulated features
+        """
+        gamma = self.scale_net(watermark).unsqueeze(-1)  # [B, C, 1]
+        beta = self.shift_net(watermark).unsqueeze(-1)   # [B, C, 1]
+        return (1 + gamma) * x + beta  # γ≈0 → identity
+
+
+class AMPBlock1(torch.nn.Module):
+    """
+    AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer.
+    AMPBlock1 has additional self.convs2 that contains additional Conv1d layers with a fixed dilation=1 followed by each layer in self.convs1
+
+    Args:
+        h (AttrDict): Hyperparameters.
+        channels (int): Number of convolution channels.
+        kernel_size (int): Size of the convolution kernel. Default is 3.
+        dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5).
+        activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Default is None.
+    """
+
+    def __init__(
+            self,
+            channels: int,
+            kernel_size: int = 3,
+            dilation: tuple = (1, 3, 5),
+            activation: str = None,
+            snake_logscale: bool = True,
+    ):
+        super().__init__()
+
+
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilation
+            ]
+        )
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                )
+                for _ in range(len(dilation))
+            ]
+        )
+        self.convs2.apply(init_weights)
+
+        self.num_layers = len(self.convs1) + len(
+            self.convs2
+        )  # Total number of conv layers
+
+        Activation1d = TorchActivation1d
+
+        # Activation functions
+        if activation == "snake":
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=activations.Snake(
+                            channels, alpha_logscale=snake_logscale
+                        )
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        elif activation == "snakebeta":
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=activations.SnakeBeta(
+                            channels, alpha_logscale=snake_logscale
+                        )
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        else:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+    def forward(self, x):
+        acts1, acts2 = self.activations[::2], self.activations[1::2]
+        for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
+            xt = a1(x)
+            xt = c1(xt)
+            xt = a2(xt)
+            xt = c2(xt)
+            x = xt + x
+
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+
+class BigVGAN(torch.nn.Module):
+    """
+    BigVGAN is a neural vocoder model that applies anti-aliased periodic activation for residual blocks (resblocks).
+    New in BigVGAN-v2: it can optionally use optimized CUDA kernels for AMP (anti-aliased multi-periodicity) blocks.
+
+    Args:
+        h (AttrDict): Hyperparameters.
+        use_cuda_kernel (bool): If set to True, loads optimized CUDA kernels for AMP. This should be used for inference only, as training is not supported with CUDA kernels.
+
+    Note:
+        - The `use_cuda_kernel` parameter should be used for inference only, as training with CUDA kernels is not supported.
+        - Ensure that the activation function is correctly specified in the hyperparameters (h.activation).
+    """
+
+    def __init__(
+            self,
+            num_mels: int = 96,
+            global_channels: int = -1,
+            upsample_initial_channel: int = 1536,
+            resblock_kernel_sizes: List[int] = [3, 7, 11],
+            resblock_dilation_sizes: List[Tuple[int]] = [(1, 3, 5), (1, 3, 5), (1, 3, 5)],
+            upsample_rates: List[int] = [4,4,2,2,2,2],
+            upsample_kernel_sizes: List[int] = [8,8,4,4,4,4],
+            snake_logscale: bool = True,
+            activation: str = 'snakebeta',
+            use_bias_at_final:bool = False,
+            use_tanh_at_final:bool = False,
+            # Watermark FiLM conditioning
+            watermark_bits: int = 0,
+            watermark_film_hidden: int = 128,
+            # VocBulwark Temporal Adapter
+            vocbulwark_bits: int = 0,
+            vocbulwark_ta_hidden: int = 128,
+            vocbulwark_zero_conv_std: float = 0.02,
+    ):
+        super().__init__()
+
+        Activation1d = TorchActivation1d
+
+        self.num_kernels = len(resblock_kernel_sizes)
+        self.num_upsamples = len(upsample_rates)
+
+        # Pre-conv
+        self.conv_pre = weight_norm(
+            Conv1d(num_mels, upsample_initial_channel, 7, 1, padding=3)
+        )
+
+        if global_channels > 0:
+            self.global_conv = torch.nn.Conv1d(global_channels, upsample_initial_channel, 1)
+
+        # Define which AMPBlock to use. BigVGAN uses AMPBlock1 as default
+        resblock_class = AMPBlock1
+
+        # Transposed conv-based upsamplers. does not apply anti-aliasing
+        self.ups = nn.ModuleList()
+        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
+            self.ups.append(
+                nn.ModuleList(
+                    [
+                        weight_norm(
+                            ConvTranspose1d(
+                                upsample_initial_channel // (2 ** i),
+                                upsample_initial_channel // (2 ** (i + 1)),
+                                k,
+                                u,
+                                padding=(k - u) // 2,
+                            )
+                        )
+                    ]
+                )
+            )
+
+        # Residual blocks using anti-aliased multi-periodicity composition modules (AMP)
+        self.resblocks = nn.ModuleList()
+        for i in range(len(self.ups)):
+            ch = upsample_initial_channel // (2 ** (i + 1))
+            for j, (k, d) in enumerate(
+                    zip(resblock_kernel_sizes, resblock_dilation_sizes)
+            ):
+                self.resblocks.append(
+                    resblock_class(ch, k, d, activation=activation, snake_logscale=snake_logscale)
+                )
+
+        # Watermark FiLM layers (one per upsample stage)
+        self.watermark_bits = watermark_bits
+        if watermark_bits > 0:
+            self.film_layers = nn.ModuleList()
+            for i in range(len(self.ups)):
+                ch_film = upsample_initial_channel // (2 ** (i + 1))
+                self.film_layers.append(
+                    FiLMBlock(watermark_bits, ch_film, watermark_film_hidden)
+                )
+
+        # VocBulwark Temporal Adapter layers (one per upsample stage)
+        self.vocbulwark_bits = vocbulwark_bits
+        if vocbulwark_bits > 0:
+            self.ta_layers = nn.ModuleList()
+            for i in range(len(self.ups)):
+                ch_ta = upsample_initial_channel // (2 ** (i + 1))
+                self.ta_layers.append(
+                    TemporalAdapterBlock(vocbulwark_bits, ch_ta, vocbulwark_ta_hidden,
+                                         zero_conv_std=vocbulwark_zero_conv_std)
+                )
+
+        # Post-conv
+        activation_post = (
+            activations.Snake(ch, alpha_logscale=snake_logscale)
+            if activation == "snake"
+            else (
+                activations.SnakeBeta(ch, alpha_logscale=snake_logscale)
+                if activation == "snakebeta"
+                else None
+            )
+        )
+        if activation_post is None:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+        self.activation_post = Activation1d(activation=activation_post)
+
+        # Whether to use bias for the final conv_post. Default to True for backward compatibility
+        self.use_bias_at_final = use_bias_at_final
+        self.conv_post = weight_norm(
+            Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final)
+        )
+
+        # Weight initialization
+        for i in range(len(self.ups)):
+            self.ups[i].apply(init_weights)
+        self.conv_post.apply(init_weights)
+
+        # Final tanh activation. Defaults to True for backward compatibility
+        self.use_tanh_at_final = use_tanh_at_final
+
+    def forward(
+            self,
+            x: torch.Tensor,
+            g: Optional[torch.Tensor] = None,
+            watermark: Optional[torch.Tensor] = None,
+    ):
+        # Pre-conv
+        x = self.conv_pre(x)
+        if g is not None:
+            x = x + self.global_conv(g)
+
+        for i in range(self.num_upsamples):
+            # Upsampling
+            for i_up in range(len(self.ups[i])):
+                x = self.ups[i][i_up](x)
+            # AMP blocks
+            xs = None
+            for j in range(self.num_kernels):
+                if xs is None:
+                    xs = self.resblocks[i * self.num_kernels + j](x)
+                else:
+                    xs += self.resblocks[i * self.num_kernels + j](x)
+            x = xs / self.num_kernels
+
+            # Watermark FiLM conditioning
+            if watermark is not None and self.watermark_bits > 0:
+                x = self.film_layers[i](x, watermark)
+
+            # VocBulwark Temporal Adapter conditioning
+            if watermark is not None and self.vocbulwark_bits > 0:
+                x = self.ta_layers[i](x, watermark)
+
+        # Post-conv
+        x = self.activation_post(x)
+        x = self.conv_post(x)
+        # Final tanh activation
+        if self.use_tanh_at_final:
+            x = torch.tanh(x)
+        else:
+            x = torch.clamp(x, min=-1.0, max=1.0)  # Bound the output to [-1, 1]
+
+        return x
+

config.json

@@ -0,0 +1,193 @@
+{
+  "add_film_watermark": false,
+  "add_mel_loss": false,
+  "add_watermark_vocbulwark": true,
+  "add_wav2vec_loss": false,
+  "add_wavlm_loss": false,
+  "add_whisper_encoder_loss": false,
+  "add_whisper_feature_loss": false,
+  "add_whisper_kl_loss": false,
+  "architectures": [
+    "HiFiGANArchitecture"
+  ],
+  "block_out_channel_mults": [
+    1,
+    2,
+    2,
+    3
+  ],
+  "dtype": "float32",
+  "hifigan_bias": false,
+  "hifigan_channels": 1536,
+  "hifigan_global_channels": 256,
+  "hifigan_in_channels": 96,
+  "hifigan_kernel_size": 7,
+  "hifigan_nonlinear_activation": "snakebeta",
+  "hifigan_nonlinear_activation_params": {
+    "negative_slope": 0.1
+  },
+  "hifigan_out_channels": 1,
+  "hifigan_resblock_dilations": [
+    [
+      1,
+      3,
+      5
+    ],
+    [
+      1,
+      3,
+      5
+    ],
+    [
+      1,
+      3,
+      5
+    ]
+  ],
+  "hifigan_resblock_kernel_sizes": [
+    3,
+    7,
+    11
+  ],
+  "hifigan_upsample_kernel_sizes": [
+    8,
+    8,
+    4,
+    4,
+    4,
+    4
+  ],
+  "hifigan_upsample_scales": [
+    4,
+    4,
+    2,
+    2,
+    2,
+    2
+  ],
+  "hifigan_use_additional_convs": true,
+  "hifigan_use_weight_norm": true,
+  "hop_length": 256,
+  "input_type": "mel",
+  "lambda_adv": 1.0,
+  "lambda_feat_match": 10.0,
+  "lambda_mel": 15.0,
+  "lambda_vocbulwark_ext": 1.0,
+  "lambda_vocbulwark_mel": 0.1,
+  "lambda_vocbulwark_mstft": 0.1,
+  "lambda_vocbulwark_wavlm": 0.0,
+  "lambda_watermark": 1.0,
+  "lambda_wav2vec": 15.0,
+  "lambda_wavlm": 15.0,
+  "lambda_whisper_encoder": 15.0,
+  "lambda_whisper_features": 15.0,
+  "lambda_whisper_kl": 1.0,
+  "mel_channels": 96,
+  "model_type": "hifi_gan",
+  "mpd_reshapes": [
+    2,
+    3,
+    5,
+    7,
+    11
+  ],
+  "mrd_resolutions": [
+    [
+      1024,
+      120,
+      600
+    ],
+    [
+      2048,
+      240,
+      1200
+    ],
+    [
+      512,
+      50,
+      240
+    ]
+  ],
+  "n_fft": 1024,
+  "raw_sample_rate": 22050,
+  "snake_logscale": true,
+  "speaker_embed_checkpoint": null,
+  "speaker_embed_config": null,
+  "target_sample_rate": 24000,
+  "train_discriminator": false,
+  "transformers_version": "4.57.3",
+  "update_speaker_embedding": true,
+  "use_wavlm_features": false,
+  "vocbulwark_acc_tau1": 0.9,
+  "vocbulwark_acc_tau2": 0.95,
+  "vocbulwark_bits": 50,
+  "vocbulwark_bootstrap_threshold": 0.6,
+  "vocbulwark_fidelity_vs_clean": false,
+  "vocbulwark_ta_hidden": 128,
+  "vocbulwark_zero_conv_std": 0.02,
+  "watermark_bits": 32,
+  "watermark_film_hidden": 128,
+  "wav2vec_cnn_frozen": true,
+  "wav2vec_cnn_model_name": "facebook/wav2vec2-base",
+  "wav2vec_cnn_sample_rate": 16000,
+  "wavlm_layer_idx": 6,
+  "wavlm_model_name": "microsoft/wavlm-base-plus",
+  "win_length": 1024,
+  "speaker_embedding_size": 768,
+  "auto_map": {
+    "AutoConfig": "configuration_hifigan.HiFiGANConfig",
+    "AutoModel": "modeling_hifigan.HiFiGANArchitecture"
+  },
+  "fixed_watermark": [
+    0,
+    1,
+    0,
+    0,
+    1,
+    1,
+    1,
+    0,
+    0,
+    0,
+    0,
+    1,
+    0,
+    1,
+    0,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    0,
+    1,
+    1,
+    1,
+    0,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    1,
+    1,
+    0,
+    0,
+    0,
+    0
+  ]
+}

configuration_hifigan.py

@@ -0,0 +1,192 @@
+from transformers import PretrainedConfig
+from .speaker_embedding_config import SpeakerEmbeddingConfig
+
+
+class HiFiGANConfig(PretrainedConfig):
+    model_type = "hifi_gan"
+
+    def __init__(
+            self,
+            # Audio settings
+            raw_sample_rate=22050,
+            target_sample_rate=24000,
+            n_fft=1024,
+            hop_length=256,
+            win_length=1024,
+            mel_channels=96,
+
+            # Speaker embedding
+            speaker_embed_config: SpeakerEmbeddingConfig = None,
+            speaker_embed_checkpoint: str = None,
+            update_speaker_embedding: bool = False,
+
+            # Input feature type: "mel" (default) or "wav2vec_cnn"
+            input_type: str = "mel",
+            wav2vec_cnn_model_name: str = "facebook/wav2vec2-base",
+            wav2vec_cnn_frozen: bool = True,
+            wav2vec_cnn_sample_rate: int = 16000,
+
+            # HiFi-GAN Generator config
+            hifigan_in_channels=96,
+            hifigan_out_channels=1,
+            hifigan_channels=512,
+            hifigan_global_channels=256,  # Speaker conditioning
+            hifigan_kernel_size=7,
+            #hifigan_upsample_scales=[8, 8, 2, 2],
+            hifigan_upsample_kernel_sizes=[8,8,4,4,4,4],
+            #hifigan_upsample_kernel_sizes=[16, 16, 4, 4],
+            hifigan_upsample_scales=[4,4,2,2,2,2],
+            hifigan_resblock_kernel_sizes=[3, 7, 11],
+            hifigan_resblock_dilations=[[1, 3, 5], [1, 3, 5], [1, 3, 5]],
+            hifigan_use_additional_convs=True,
+            hifigan_bias=False,
+            hifigan_nonlinear_activation="snakebeta",
+            snake_logscale=True,
+            hifigan_nonlinear_activation_params={"negative_slope": 0.1},
+            hifigan_use_weight_norm=True,
+
+            mpd_reshapes=[2, 3, 5, 7, 11],
+            mrd_resolutions=[
+                (1024, 120, 600),
+                (2048, 240, 1200),
+                (512, 50, 240),
+            ],
+
+            # VAE config
+            block_out_channel_mults=[1, 2, 2, 3],
+
+            # Loss weights
+            lambda_adv=1.0,  # Adversarial loss weight
+            lambda_feat_match=10.0,  # Feature matching loss weight
+            add_mel_loss=True,
+            lambda_mel=15.0,  # Mel spectrogram loss weight
+            lambda_wavlm=15.0,  # WavLM perceptual loss weight
+            add_wavlm_loss=False,
+            add_wav2vec_loss=False,
+            lambda_wav2vec=15.0,  # Wav2Vec perceptual loss weight
+            add_whisper_encoder_loss=False,
+            add_whisper_kl_loss=False,
+            lambda_whisper_encoder=15.0,  # Whisper encoder hidden states L1 weight
+            lambda_whisper_kl=1.0,  # Whisper decoder KL divergence weight
+            add_whisper_feature_loss=False,
+            lambda_whisper_features=15.0,  # Whisper conv feature extractor L1 weight
+
+
+            use_wavlm_features: bool = False,  # Use WavLM instead of MEL
+            wavlm_model_name: str = "microsoft/wavlm-base-plus",
+            wavlm_layer_idx: int = 6,  # Which layer to extract features from (0-11 for base)
+            
+
+            train_discriminator: bool = True,
+
+            # Watermark (FiLM)
+            add_film_watermark: bool = False,
+            watermark_bits: int = 32,
+            lambda_watermark: float = 1.0,
+            watermark_film_hidden: int = 128,
+
+            # VocBulwark watermarking
+            add_watermark_vocbulwark: bool = False,
+            vocbulwark_bits: int = 100,
+            vocbulwark_ta_hidden: int = 128,
+            lambda_vocbulwark_ext: float = 1.0,
+            lambda_vocbulwark_mel: float = 0.1,
+            lambda_vocbulwark_mstft: float = 0.1,
+            lambda_vocbulwark_wavlm: float = 0.0,
+            vocbulwark_acc_tau1: float = 0.9,
+            vocbulwark_acc_tau2: float = 0.95,
+            vocbulwark_bootstrap_threshold: float = 0.0,
+            vocbulwark_zero_conv_std: float = 0.02,
+            vocbulwark_fidelity_vs_clean: bool = False,
+
+            **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        # Audio
+        self.raw_sample_rate = raw_sample_rate
+        self.target_sample_rate = target_sample_rate
+        self.n_fft = n_fft
+        self.hop_length = hop_length
+        self.win_length = win_length
+
+        # Mel
+        self.mel_channels = mel_channels
+
+        # Speaker
+        self.speaker_embed_config = speaker_embed_config
+        self.speaker_embed_checkpoint = speaker_embed_checkpoint
+        self.update_speaker_embedding = update_speaker_embedding
+
+        # Input type
+        self.input_type = input_type
+        self.wav2vec_cnn_model_name = wav2vec_cnn_model_name
+        self.wav2vec_cnn_frozen = wav2vec_cnn_frozen
+        self.wav2vec_cnn_sample_rate = wav2vec_cnn_sample_rate
+
+        # VAE parameters
+        self.block_out_channel_mults = block_out_channel_mults
+
+        # HiFi-GAN parameters
+        self.hifigan_in_channels = hifigan_in_channels
+        self.hifigan_out_channels = hifigan_out_channels
+        self.hifigan_channels = hifigan_channels
+        self.hifigan_global_channels = hifigan_global_channels
+        self.hifigan_kernel_size = hifigan_kernel_size
+        self.hifigan_upsample_scales = hifigan_upsample_scales
+        self.hifigan_upsample_kernel_sizes = hifigan_upsample_kernel_sizes
+        self.hifigan_resblock_kernel_sizes = hifigan_resblock_kernel_sizes
+        self.hifigan_resblock_dilations = hifigan_resblock_dilations
+        self.hifigan_use_additional_convs = hifigan_use_additional_convs
+        self.hifigan_bias = hifigan_bias
+        self.hifigan_nonlinear_activation = hifigan_nonlinear_activation
+        self.hifigan_nonlinear_activation_params = hifigan_nonlinear_activation_params
+        self.hifigan_use_weight_norm = hifigan_use_weight_norm
+        self.snake_logscale = snake_logscale
+
+        self.mpd_reshapes = mpd_reshapes
+        self.mrd_resolutions = mrd_resolutions
+
+        # Loss weights
+        self.lambda_adv = lambda_adv
+        self.lambda_feat_match = lambda_feat_match
+        self.add_mel_loss = add_mel_loss
+        self.lambda_mel = lambda_mel
+        self.lambda_wavlm = lambda_wavlm
+        self.add_wavlm_loss = add_wavlm_loss
+        self.add_wav2vec_loss = add_wav2vec_loss
+        self.lambda_wav2vec = lambda_wav2vec
+        self.add_whisper_encoder_loss = add_whisper_encoder_loss
+        self.add_whisper_kl_loss = add_whisper_kl_loss
+        self.lambda_whisper_encoder = lambda_whisper_encoder
+        self.lambda_whisper_kl = lambda_whisper_kl
+        self.add_whisper_feature_loss = add_whisper_feature_loss
+        self.lambda_whisper_features = lambda_whisper_features
+
+        self.use_wavlm_features = use_wavlm_features
+        self.wavlm_model_name = wavlm_model_name
+        self.wavlm_layer_idx = wavlm_layer_idx
+
+
+        # train discrimnator or not
+        self.train_discriminator = train_discriminator
+
+        # Watermark (FiLM)
+        self.add_film_watermark = add_film_watermark
+        self.watermark_bits = watermark_bits
+        self.lambda_watermark = lambda_watermark
+        self.watermark_film_hidden = watermark_film_hidden
+
+        # VocBulwark
+        self.add_watermark_vocbulwark = add_watermark_vocbulwark
+        self.vocbulwark_bits = vocbulwark_bits
+        self.vocbulwark_ta_hidden = vocbulwark_ta_hidden
+        self.lambda_vocbulwark_ext = lambda_vocbulwark_ext
+        self.lambda_vocbulwark_mel = lambda_vocbulwark_mel
+        self.lambda_vocbulwark_mstft = lambda_vocbulwark_mstft
+        self.lambda_vocbulwark_wavlm = lambda_vocbulwark_wavlm
+        self.vocbulwark_acc_tau1 = vocbulwark_acc_tau1
+        self.vocbulwark_acc_tau2 = vocbulwark_acc_tau2
+        self.vocbulwark_bootstrap_threshold = vocbulwark_bootstrap_threshold
+        self.vocbulwark_zero_conv_std = vocbulwark_zero_conv_std
+        self.vocbulwark_fidelity_vs_clean = vocbulwark_fidelity_vs_clean

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-large-speaker-encoder\"   # <- speaker encoder\n",
+    "GEN_REPO = \"mlr2000/vocoder-large\"   # <- vocoder (paired with SPK_REPO)\n",
+    "DET_REPO = \"mlr2000/vocoder-large-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:626e334b29898f1b4840552666fd72f431baef10b4386fba4853a3310718e9e7
+size 474989104

modeling_hifigan.py

@@ -0,0 +1,131 @@
+"""Self-contained inference model for the VocBulwark HiFi-GAN / BigVGAN vocoder.
+
+Loadable with `AutoModel.from_pretrained(path, trust_remote_code=True)` without
+the training repository. This is the lean vocoder: it takes a mel-spectrogram and
+a precomputed speaker embedding, and does NOT include the speaker encoder (that
+is a separate model), the training losses, discriminators, or watermark heads.
+"""
+from dataclasses import dataclass
+
+import torch
+import torch.nn.utils.parametrize as parametrize
+from torch.nn.utils import remove_weight_norm
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+
+def _strip_weight_norm(root):
+    """Remove every weight-norm reparametrization (both conventions), leaving
+    plain ``.weight`` parameters that match the folded export weights."""
+    for m in root.modules():
+        if parametrize.is_parametrized(m, "weight"):
+            parametrize.remove_parametrizations(m, "weight", leave_parametrized=True)
+        elif hasattr(m, "weight_g"):
+            try:
+                remove_weight_norm(m)
+            except (ValueError, RuntimeError):
+                pass
+
+from .configuration_hifigan import HiFiGANConfig
+from .bigvgan_model import BigVGAN
+
+# transformers' trust_remote_code loader copies only modules *directly* imported
+# by this entry file, and recognizes the `from .X import Y` form. Import a real
+# symbol from every bundled leaf module so all files get copied into the
+# dynamic-module cache (they are used transitively by the imports above).
+from .bigvgan_activations import Snake as _Snake  # noqa: F401
+from .alias_free_act import Activation1d as _Activation1d  # noqa: F401
+from .alias_free_resample import UpSample1d as _UpSample1d  # noqa: F401
+from .alias_free_filter import LowPassFilter1d as _LowPassFilter1d  # noqa: F401
+from .temporal_adapter import TemporalAdapterBlock as _TemporalAdapterBlock  # noqa: F401
+
+
+@dataclass
+class HiFiGANOutput(ModelOutput):
+    audio: torch.Tensor = None
+
+
+class HiFiGANArchitecture(PreTrainedModel):
+    config_class = HiFiGANConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        # Lean vocoder: the speaker embedding is supplied at inference time, so the
+        # generator only needs its dimension to build the conditioning conv.
+        global_channels = int(getattr(config, "speaker_embedding_size", 0)) or -1
+        self.hifi_gan = BigVGAN(
+            num_mels=config.hifigan_in_channels,
+            upsample_initial_channel=config.hifigan_channels,
+            resblock_kernel_sizes=config.hifigan_resblock_kernel_sizes,
+            resblock_dilation_sizes=config.hifigan_resblock_dilations,
+            upsample_kernel_sizes=config.hifigan_upsample_kernel_sizes,
+            upsample_rates=config.hifigan_upsample_scales,
+            snake_logscale=config.snake_logscale,
+            activation=config.hifigan_nonlinear_activation,
+            use_bias_at_final=config.hifigan_bias,
+            global_channels=global_channels,
+            watermark_bits=config.watermark_bits if config.add_film_watermark else 0,
+            watermark_film_hidden=config.watermark_film_hidden,
+            vocbulwark_bits=config.vocbulwark_bits if config.add_watermark_vocbulwark else 0,
+            vocbulwark_ta_hidden=config.vocbulwark_ta_hidden,
+            vocbulwark_zero_conv_std=getattr(config, "vocbulwark_zero_conv_std", 0.02),
+        )
+
+        # Export folds weight-norm into plain .weight; strip the reparametrization
+        # here so module keys match (and inference is independent of torch's
+        # weight_norm convention).
+        _strip_weight_norm(self)
+
+    @torch.no_grad()
+    def forward(
+        self,
+        mel_spectrogram=None,
+        speaker_embedding=None,
+        input_features=None,
+        return_loss=False,
+        **kwargs,
+    ):
+        """Vocode a (log-)mel spectrogram into a waveform.
+
+        Args:
+            mel_spectrogram:   [B, mel_channels, T] input features.
+            speaker_embedding: [B, speaker_embedding_size] speaker conditioning,
+                               produced by the companion speaker-encoder model.
+        Returns:
+            HiFiGANOutput with `.audio` of shape [B, 1, samples] @ target_sample_rate.
+
+        The model's fixed provenance watermark is embedded in every generated clip
+        and cannot be changed or disabled through this interface.
+        """
+        hifi_gan_device = self.hifi_gan.conv_pre.bias.device
+        if torch.is_autocast_enabled():
+            target_dtype = torch.get_autocast_gpu_dtype()
+        else:
+            target_dtype = self.hifi_gan.conv_pre.bias.dtype
+
+        if input_features is not None:
+            features = input_features.to(device=hifi_gan_device, dtype=target_dtype)
+        else:
+            features = mel_spectrogram.to(device=hifi_gan_device, dtype=target_dtype)
+
+        if speaker_embedding is None:
+            raise ValueError(
+                "speaker_embedding is required: pass a "
+                "[B, config.speaker_embedding_size] tensor from the companion "
+                "speaker-encoder model.")
+        g = speaker_embedding.to(device=hifi_gan_device, dtype=target_dtype)
+        if g.dim() == 2:                       # [B, E] -> [B, E, 1]
+            g = g.unsqueeze(2)
+
+        # Provenance watermark: the model's fixed signature is embedded in every
+        # generated clip. It is intentionally not overridable through this
+        # interface — there is no way to disable or change it here.
+        watermark = None
+        fw = getattr(self.config, "fixed_watermark", None)
+        if fw is not None:
+            watermark = torch.tensor(fw, dtype=target_dtype, device=hifi_gan_device)
+            watermark = watermark.unsqueeze(0).expand(features.shape[0], -1)
+
+        audio = self.hifi_gan(features, g=g, watermark=watermark)
+        return HiFiGANOutput(audio=audio)

speaker_embedding_config.py

@@ -0,0 +1,31 @@
+from transformers.models.wav2vec2.configuration_wav2vec2 import Wav2Vec2Config
+
+
+class SpeakerEmbeddingConfig(Wav2Vec2Config):
+
+    def __init__(
+        self,
+        embedding_size: int = 768,
+        train_batch_speakers: int = 64,
+        train_batch_per_device: int = 64,
+        train_batch_samples_per_speaker: int = 10,
+        disable_positional_embeddings: bool = False,
+        loss_margin: float = 0.2,
+        loss_scale: float = 30.0,
+        use_layer_weights: bool = True,
+        n_projection_layers: int = 3,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.embedding_size = embedding_size
+        self.train_batch_per_device = train_batch_per_device
+        self.train_batch_speakers = train_batch_speakers
+        self.train_batch_samples_per_speaker = train_batch_samples_per_speaker
+        self.disable_positional_embeddings = disable_positional_embeddings
+        self.loss_margin = loss_margin
+        self.loss_scale = loss_scale
+        self.use_layer_weights = use_layer_weights
+        self.n_projection_layers = n_projection_layers
+        #self.conv_stride = (7, 2, 2, 2, 2, 2, 2)
+        #self.conv_kernel = (15, 5, 3, 3, 3, 2, 2)
+        #self.conv_dim = (512, 512, 512, 512, 512, 512, 512)

temporal_adapter.py

@@ -0,0 +1,100 @@
+"""
+Temporal Adapter (TA) from VocBulwark (Appendix B.1).
+
+Lightweight module injected at each upsample stage of the frozen vocoder.
+Embeds watermark bits into acoustic features via:
+  1. Acoustic Feature Alignment: Emb(w) → PFP → w_proj ∈ [B, C, 1]
+  2. Frame-level Temporal Broadcasting: w_proj → w_latent ∈ [B, C, T]
+  3. Adaptive Injection: concat(w_latent, h) → downsample → DSC → zero_conv + h (residual)
+
+Architecture from paper:
+  Embedding: Linear(bits, 2*C) → LeakyReLU → Linear(2*C, 512) → LeakyReLU
+  PFP: Linear(512, rank) → SiLU → Linear(rank, C)
+"""
+
+import torch
+import torch.nn as nn
+
+
+class TemporalAdapterBlock(nn.Module):
+    """Single TA block for one upsample stage."""
+
+    def __init__(self, watermark_bits: int, channels: int, hidden_dim: int = 128,
+                 zero_conv_std: float = 0.02):
+        super().__init__()
+        self.channels = channels
+
+        # Acoustic Feature Alignment (paper B.1):
+        # Two FC layers with two LeakyReLU activations.
+        # First FC projects to 2 * hidden_feature_channels, second to 512.
+        self.embedding = nn.Sequential(
+            nn.Linear(watermark_bits, channels * 2),
+            nn.LeakyReLU(0.2),
+            nn.Linear(channels * 2, 512),
+            nn.LeakyReLU(0.2),
+        )
+
+        # Progressive Feature Projection (PFP) (paper B.1):
+        # Two FC layers with SiLU activation between them.
+        # First FC reduces to pre-defined Rank, second aligns to C.
+        rank = hidden_dim // 2  # pre-defined rank value
+        self.pfp = nn.Sequential(
+            nn.Linear(512, rank),
+            nn.SiLU(),
+            nn.Linear(rank, channels),
+        )
+
+        # Adaptive Injection: concat [w_latent, h] → downsample → DSC → zero_conv
+        # Downsample from 2C to C
+        self.downsample = nn.Conv1d(channels * 2, channels, kernel_size=1)
+
+        # Depth-wise Separable Convolution (DSC)
+        # BatchNorm is load-bearing: it normalizes the DSC output to ~unit variance
+        # before the small-init zero_conv, giving the injected watermark residual a
+        # meaningful scale. Without any norm the residual is ~0, the watermark barely
+        # perturbs the audio, extraction gradients vanish (grad_norm ~0.03) and
+        # training sticks at exactly random (acc 0.5, loss_ext = ln 2). InstanceNorm
+        # also fails (it cancels the time-constant watermark). So: keep BatchNorm.
+        self.dsc = nn.Sequential(
+            # Depth-wise conv
+            nn.Conv1d(channels, channels, kernel_size=3, padding=1, groups=channels),
+            nn.BatchNorm1d(channels),
+            nn.LeakyReLU(0.2),
+            # Point-wise conv
+            nn.Conv1d(channels, channels, kernel_size=1),
+            nn.BatchNorm1d(channels),
+            nn.LeakyReLU(0.2),
+        )
+
+        # Small-init convolution (not zero-init) to break the bootstrap deadlock.
+        # Zero-init prevents any watermark signal from reaching the audio,
+        # so the extractor can never learn. Small random init ensures different
+        # watermarks produce slightly different audio from the start.
+        self.zero_conv = nn.Conv1d(channels, channels, kernel_size=1)
+        nn.init.normal_(self.zero_conv.weight, std=zero_conv_std)
+        nn.init.zeros_(self.zero_conv.bias)
+
+    def forward(self, h: torch.Tensor, watermark: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            h: [B, C, T] hidden features from vocoder upsample stage
+            watermark: [B, watermark_bits] binary watermark vector (float)
+        Returns:
+            [B, C, T] watermarked features (h + residual)
+        """
+        B, C, T = h.shape
+
+        # 1. Acoustic Feature Alignment
+        w_emb = self.embedding(watermark)           # [B, 512]
+        w_proj = self.pfp(w_emb)                    # [B, C]
+
+        # 2. Frame-level Temporal Broadcasting
+        w_latent = w_proj.unsqueeze(-1).expand(B, C, T)  # [B, C, T]
+
+        # 3. Adaptive Injection
+        h_cat = torch.cat([w_latent, h], dim=1)     # [B, 2C, T]
+        h_down = self.downsample(h_cat)              # [B, C, T]
+        h_dsc = self.dsc(h_down)                     # [B, C, T]
+        residual = self.zero_conv(h_dsc)             # [B, C, T]
+
+        return h + residual