| """Ported from github.com/ProSDD/codes @ bdb80440cd946368e77f4b51f6ddf26c7443dc33 (no LICENSE file): |
| model_stage2realfake.py (ProSDDStage2), main_eval.py (inference_forward), data_utils_eval.py (pad). |
| |
| Inference-only subset of ProSDDStage2. Edits for the Arena wrapper: |
| 1. `ssl` is built from the vendored facebook/wav2vec2-xls-r-300m config (xlsr300m_config.json, HF @ |
| 1a640f32ac3e39899438a2931f9924c02f080a54) instead of Wav2Vec2Model.from_pretrained(...): no Hub |
| download; the checkpoint overwrites every backbone tensor (strict load, see the wrapper). |
| 2. Training-only members are not built: mask_embed, pros_ln, final_proj (1024->448 Stage-1 head) and the |
| masked/contrastive pass. classifier_pool is fixed to "mean" (the default; the released checkpoints have |
| no attn/attn_q tensors). |
| 3. forward(wav) = main_eval.py:inference_forward verbatim (CNN feature extractor -> feature projection -> |
| latent truncated / zero-padded to T_target=200 frames -> encoder (clean pass, no attention mask) -> |
| mean over the 200 frames -> cls_head), returning the (B, 2) logits. |
| 4. pad_np = numpy port of data_utils_eval.py:pad (64000 samples, centre crop / centred zero-pad). |
| """ |
| import os |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| from transformers import Wav2Vec2Config, Wav2Vec2Model |
|
|
| _DIR = os.path.dirname(os.path.abspath(__file__)) |
| XLSR_CONFIG = os.path.join(_DIR, "xlsr300m_config.json") |
| TARGET_SAMPLES = 64000 |
| T_TARGET = 200 |
|
|
|
|
| def pad_np(wav, max_len=TARGET_SAMPLES): |
| """data_utils_eval.py:pad -- longer: centre crop s=(n-max_len)//2; shorter: zero-pad left=total//2, right=rest.""" |
| wav = np.asarray(wav, dtype=np.float32).reshape(-1) |
| n = wav.shape[0] |
| if n >= max_len: |
| s = (n - max_len) // 2 |
| return np.ascontiguousarray(wav[s:s + max_len]) |
| left = (max_len - n) // 2 |
| out = np.zeros(max_len, dtype=np.float32) |
| out[left:left + n] = wav |
| return out |
|
|
|
|
| class ProSDDNet(nn.Module): |
| def __init__(self, num_classes=2, T_target=T_TARGET): |
| super().__init__() |
| self.T_target = int(T_target) |
| self.ssl = Wav2Vec2Model(Wav2Vec2Config.from_json_file(XLSR_CONFIG)) |
| self.hidden_dim = self.ssl.config.hidden_size |
| self.cls_head = nn.Sequential( |
| nn.Linear(self.hidden_dim, 512), |
| nn.ReLU(inplace=True), |
| nn.Dropout(0.1), |
| nn.Linear(512, num_classes), |
| ) |
|
|
| def forward(self, wav): |
| z = self.ssl.feature_extractor(wav) |
| z = z.transpose(1, 2) |
| z = self.ssl.feature_projection(z) |
| if isinstance(z, (tuple, list)): |
| z = z[0] |
| B, T, H = z.shape |
| Tt = self.T_target |
| if T > Tt: |
| z = z[:, :Tt, :] |
| elif T < Tt: |
| z = torch.cat([z, z.new_zeros(B, Tt - T, H)], dim=1) |
| out = self.ssl.encoder(z, attention_mask=None, output_hidden_states=False, return_dict=True) |
| pooled = out.last_hidden_state.mean(dim=1) |
| return self.cls_head(pooled) |
|
|
