Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| from typing import Dict, Optional, Tuple | |
| import torch.nn.functional as F | |
| import torch | |
| import torch.nn as nn | |
| class WMDetector(nn.Module): | |
| """ | |
| Detect watermarks in an audio signal using a Transformer architecture, | |
| where the watermark bits are split into bytes (8 bits each). | |
| We assume nbits is a multiple of 8. | |
| """ | |
| def __init__( | |
| self, input_channels: int, nbits: int, nchunk_size: int, d_model: int = 512 | |
| ): | |
| """ | |
| Args: | |
| input_channels (int): Number of input channels in the audio feature (e.g., mel channels). | |
| nbits (int): Total number of bits in the watermark, must be a multiple of 8. | |
| d_model (int): Embedding dimension for the Transformer. | |
| """ | |
| super().__init__() | |
| self.nchunk_size = nchunk_size | |
| assert nbits % nchunk_size == 0, "nbits must be a multiple of 8!" | |
| self.nbits = nbits | |
| self.d_model = d_model | |
| # Number of bytes | |
| self.nchunks = nbits // nchunk_size | |
| # 1D convolution to map the input channels to d_model | |
| self.embedding = nn.Conv1d(input_channels, d_model, kernel_size=1) | |
| # Transformer encoder block | |
| self.transformer = nn.TransformerEncoder( | |
| nn.TransformerEncoderLayer( | |
| d_model=d_model, | |
| nhead=1, | |
| dim_feedforward=d_model * 2, | |
| activation="gelu", | |
| batch_first=True, | |
| ), | |
| num_layers=8, | |
| ) | |
| # A linear head for watermark presence detection (binary) | |
| self.watermark_head = nn.Linear(d_model, 1) | |
| # For each byte, we perform a 256-way classification | |
| self.message_heads = nn.ModuleList( | |
| nn.Linear(d_model, 2**nchunk_size) for _ in range(self.nchunks) | |
| ) | |
| # Learnable embeddings for each byte chunk (instead of per bit) | |
| # Shape: [nchunks, d_model] | |
| self.nchunk_embeddings = nn.Parameter(torch.randn(self.nchunks, d_model)) | |
| def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """ | |
| Forward pass of the detector. | |
| Returns: | |
| logits (torch.Tensor): Watermark detection logits of shape [batch, seq_len]. | |
| chunk_logits (torch.Tensor): Byte-level classification logits of shape [batch, nchunks, 256]. | |
| """ | |
| batch_size, input_channels, time_steps = x.shape | |
| # 1) Map [batch, in_channels, time_steps] → [batch, time_steps, d_model] | |
| x = self.embedding(x).permute(0, 2, 1) # [batch, time_steps, d_model] | |
| # 2) Prepend chunk embeddings at the beginning of the sequence | |
| # [nchunks, d_model] → [1, nchunks, d_model] → [batch, nchunks, d_model] | |
| nchunk_embeds = self.nchunk_embeddings.unsqueeze(0).expand(batch_size, -1, -1) | |
| # Concatenate along the time dimension: [batch, nchunks + time_steps, d_model] | |
| x = torch.cat([nchunk_embeds, x], dim=1) | |
| # 3) Pass through the Transformer | |
| x = self.transformer(x) | |
| # x has shape [batch, nchunks + time_steps, d_model] | |
| # (a) Watermark presence detection: skip the first nchunks | |
| detection_part = x[:, self.nchunks :] # [batch, time_steps, d_model] | |
| logits = self.watermark_head(detection_part).squeeze(-1) # [batch, time_steps] | |
| # (b) Message decoding: use the first nchunks | |
| message_part = x[:, : self.nchunks] # [batch, nchunks, d_model] | |
| chunk_logits_list = [] | |
| for i, head in enumerate(self.message_heads): | |
| # message_part[:, i, :] has shape [batch, d_model] | |
| # each head outputs [batch, 256] | |
| chunk_vec = message_part[:, i, :] | |
| chunk_logits_list.append(head(chunk_vec).unsqueeze(1)) # [batch, 1, 256] | |
| # Concatenate along the 'nchunks' dimension → [batch, nchunks, 256] | |
| chunk_logits = torch.cat(chunk_logits_list, dim=1) | |
| return logits, chunk_logits | |
| def detect_watermark( | |
| self, | |
| x: torch.Tensor, | |
| sample_rate: Optional[int] = None, | |
| threshold: float = 0.5, | |
| ) -> Tuple[float, torch.Tensor, torch.Tensor]: | |
| """ | |
| A convenience function for inference. | |
| Returns: | |
| detect_prob (float): Probability that the audio is watermarked. | |
| binary_message (torch.Tensor): The recovered message of shape [batch, nbits] (binary). | |
| detected (torch.Tensor): The sigmoid values of the per-timestep watermark detection. | |
| """ | |
| logits, chunk_logits = self.forward(x) | |
| # logits: [batch, seq_len] → raw logits for watermark presence detection | |
| # chunk_logits: [batch, nchunks, 256] → classification logits for each byte | |
| # (1) Compute watermark detection probability | |
| detected = torch.sigmoid(logits) # [batch, seq_len] | |
| detect_prob = detected.mean(dim=-1).cpu().item() | |
| # (2) Decode the message: chunk_logits has shape [batch, nchunks, 256] | |
| chunk_probs = F.softmax(chunk_logits, dim=-1) # [batch, nchunks, 256] | |
| chunk_indices = torch.argmax( | |
| chunk_probs, dim=-1 | |
| ) # [batch, nchunks], each in [0..255] | |
| # (3) Convert each byte back to 8 bits | |
| # Finally, assemble into a [batch, nbits] binary tensor | |
| binary_message = [] | |
| for i in range(self.nchunks): | |
| chunk_val = chunk_indices[:, i] # [batch] | |
| # Extract 8 bits from the integer (0..255) | |
| chunk_bits = [] | |
| for b in range(self.nchunk_size): | |
| bit_b = (chunk_val >> b) & 1 # get bit b | |
| chunk_bits.append(bit_b.unsqueeze(-1)) | |
| # Concatenate bits to shape [batch, 8] | |
| chunk_bits = torch.cat(chunk_bits, dim=-1) | |
| binary_message.append(chunk_bits) | |
| # Concatenate all bytes → [batch, nbits] | |
| binary_message = torch.cat(binary_message, dim=-1) | |
| return detect_prob, binary_message, detected | |
| class WMEmbedder(nn.Module): | |
| """ | |
| A class that takes a secret message, processes it into chunk embeddings | |
| (as a small sequence), and uses a TransformerDecoder to do cross-attention | |
| between the original hidden (target) and the watermark tokens (memory). | |
| """ | |
| def __init__( | |
| self, | |
| nbits: int, # total bits in the secret message | |
| input_dim: int, # the input dimension (e.g. audio feature dimension) | |
| nchunk_size: int, | |
| hidden_dim: int = 256, | |
| num_heads: int = 1, | |
| num_layers: int = 4, | |
| ): | |
| super().__init__() | |
| self.nchunk_size = nchunk_size | |
| assert nbits % nchunk_size == 0, "nbits must be a multiple of nchunk_size!" | |
| self.nbits = nbits | |
| self.nchunks = nbits // nchunk_size # how many chunks | |
| # Each chunk (0..2^nchunk_size - 1) maps to an embedding of size [hidden_dim] | |
| self.msg_embeddings = nn.ModuleList( | |
| nn.Embedding(2**nchunk_size, hidden_dim) for _ in range(self.nchunks) | |
| ) | |
| # Linear to project [input_dim] -> [hidden_dim] | |
| self.input_projection = nn.Linear(input_dim, hidden_dim) | |
| # TransformerDecoder for cross-attention | |
| # d_model=hidden_dim, so the decoder expects [b, seq_len, hidden_dim] as tgt | |
| # and [b, memory_len, hidden_dim] as memory | |
| decoder_layer = nn.TransformerDecoderLayer( | |
| d_model=hidden_dim, | |
| nhead=num_heads, | |
| dim_feedforward=2 * hidden_dim, | |
| activation="gelu", | |
| batch_first=True, # so shape is [batch, seq, feature] | |
| ) | |
| self.transformer_decoder = nn.TransformerDecoder( | |
| decoder_layer, num_layers=num_layers | |
| ) | |
| # Project [hidden_dim] -> [input_dim] | |
| # self.output_projection1 = nn.Linear(hidden_dim * 2, hidden_dim) | |
| self.output_projection = nn.Linear(hidden_dim, input_dim) | |
| def forward(self, hidden: torch.Tensor, msg: torch.Tensor) -> torch.Tensor: | |
| """ | |
| Args: | |
| hidden: [batch, input_dim, seq_len] | |
| msg: [batch, nbits] | |
| Returns: | |
| A tensor [batch, input_dim, seq_len] with watermark injected. | |
| """ | |
| b, in_dim, seq_len = hidden.shape | |
| # 1) Project input features to [b, seq_len, hidden_dim] | |
| hidden_projected = self.input_projection( | |
| hidden.permute(0, 2, 1) | |
| ) # => [b, seq_len, hidden_dim] | |
| # 2) Convert the msg bits into a sequence of chunk embeddings | |
| # We keep each chunk as one token => [b, nchunks, hidden_dim] | |
| chunk_emb_list = [] | |
| for i in range(self.nchunks): | |
| # msg[:, i*nchunk_size : (i+1)*nchunk_size] => shape [b, nchunk_size] | |
| chunk_bits = msg[:, i * self.nchunk_size : (i + 1) * self.nchunk_size] | |
| chunk_val = torch.zeros_like(chunk_bits[:, 0]) # shape [b] | |
| for bit_idx in range(self.nchunk_size): | |
| # shift bits | |
| chunk_val += chunk_bits[:, bit_idx] << bit_idx | |
| # embedding => [b, hidden_dim] | |
| chunk_emb = self.msg_embeddings[i](chunk_val) | |
| chunk_emb_list.append(chunk_emb.unsqueeze(1)) # => [b,1,hidden_dim] | |
| # Concat => [b, nchunks, hidden_dim] | |
| chunk_emb_seq = torch.cat(chunk_emb_list, dim=1) # [b, nchunks, hidden_dim] | |
| # 3) Use chunk_emb_seq as memory, hidden_projected as target for TransformerDecoder | |
| # | |
| # TransformerDecoder forward signature: | |
| # transformer_decoder(tgt, memory, ...) | |
| # => [b, seq_len, hidden_dim] | |
| x_decoded = self.transformer_decoder( | |
| tgt=hidden_projected, # [b, seq_len, hidden_dim] | |
| memory=chunk_emb_seq, # [b, nchunks, hidden_dim] | |
| ) | |
| # 4) Project back to input_dim => [b, seq_len, input_dim] | |
| x_output = self.output_projection(x_decoded) | |
| # 5) permute back to [b, input_dim, seq_len] | |
| x_output = x_output.permute(0, 2, 1) # => [b, input_dim, seq_len] | |
| # 6) (Optional) Residual with original hidden | |
| x_output = x_output + hidden | |
| return x_output | |
| from speechtokenizer import SpeechTokenizer | |
| class SBW(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.nbits = 16 | |
| config_path = ( | |
| "speechtokenizer/pretrained_model/speechtokenizer_hubert_avg_config.json" | |
| ) | |
| ckpt_path = "speechtokenizer/pretrained_model/SpeechTokenizer.pt" | |
| self.st_model = SpeechTokenizer.load_from_checkpoint(config_path, ckpt_path) | |
| self.msg_processor = WMEmbedder( | |
| nbits=16, | |
| input_dim=1024, | |
| nchunk_size=4, | |
| ) | |
| self.detector = WMDetector( | |
| 1024, | |
| 16, | |
| nchunk_size=4, | |
| ) | |
| def detect_watermark( | |
| self, x: torch.Tensor, return_logits=False | |
| ) -> Tuple[float, torch.Tensor]: | |
| embedding = self.st_model.forward_feature(x) | |
| if return_logits: | |
| return self.detector(embedding) | |
| return self.detector.detect_watermark(embedding) | |
| def forward( | |
| self, | |
| speech_input: torch.Tensor, | |
| message: Optional[torch.Tensor] = None, | |
| ) -> Dict[str, torch.Tensor]: | |
| recon, recon_wm, acoustic, acoustic_wm = self.st_model( | |
| speech_input, msg_processor=self.msg_processor, message=message | |
| ) | |
| wav_length = min(speech_input.size(-1), recon_wm.size(-1)) | |
| speech_input = speech_input[..., :wav_length] | |
| recon = recon[..., :wav_length] | |
| recon_wm = recon_wm[..., :wav_length] | |
| return { | |
| "recon": recon, | |
| "recon_wm": recon_wm, | |
| } | |