File size: 3,740 Bytes
8839bbc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """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)
|