File size: 2,942 Bytes
a975ec5 | 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 | """Audio encoder: Qwen3-ASR-1.7B audio tower architecture with VeriSpoof weights.
wav (B, 128000) float32 16 kHz -> (B, 104, 25, 1024) bf16 hidden states (input embedding and every layer output).
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file
AUT_REPO = "Qwen/Qwen3-ASR-1.7B-hf" # architecture config and mel filters only
class LogMel(nn.Module):
def __init__(self, mel_filters, n_fft=400, hop=160):
super().__init__()
self.n_fft, self.hop = n_fft, hop
self.register_buffer("mel", torch.as_tensor(mel_filters, dtype=torch.float32).T.contiguous())
self.register_buffer("window", torch.hann_window(n_fft))
@torch.autocast("cuda", enabled=False)
def forward(self, wav):
stft = torch.stft(wav.float(), self.n_fft, self.hop, window=self.window, return_complex=True)
mag = stft[..., :-1].abs() ** 2
log_spec = torch.clamp(self.mel @ mag, min=1e-10).log10()
log_spec = torch.maximum(log_spec, log_spec.amax(dim=(1, 2), keepdim=True) - 8.0)
return (log_spec + 4.0) / 4.0
class AudioEncoder:
def __init__(self, weights, device="cuda"):
from transformers import AutoConfig, AutoFeatureExtractor
from transformers.models.qwen3_asr.modeling_qwen3_asr import Qwen3ASREncoder
self.device = device
acfg = AutoConfig.from_pretrained(AUT_REPO).audio_config
acfg._attn_implementation = "sdpa"
self.encoder = Qwen3ASREncoder(acfg)
self.encoder.load_state_dict(load_file(weights), strict=True)
self.encoder.float()
fe = AutoFeatureExtractor.from_pretrained(AUT_REPO)
self.chunk = 2 * acfg.n_window
self.logmel = LogMel(fe.mel_filters, fe.n_fft, fe.hop_length)
for m in (self.encoder, self.logmel):
m.eval().requires_grad_(False).to(device)
self._hs = []
self.encoder.layers[0].register_forward_pre_hook(
lambda m, a, k: self._hs.append(a[0] if a else k["hidden_states"]), with_kwargs=True)
for layer in self.encoder.layers:
layer.register_forward_hook(lambda m, a, o: self._hs.append(o[0] if isinstance(o, tuple) else o))
@torch.no_grad()
def __call__(self, wav):
wav = wav.to(self.device, torch.float32)
with torch.autocast("cuda", dtype=torch.bfloat16):
mel = self.logmel(wav)
T = mel.shape[-1]
pad = (-T) % self.chunk
mask = torch.ones(mel.shape[0], T + pad, dtype=torch.long, device=mel.device)
mask[:, T:] = 0
self._hs = []
out = self.encoder(input_features=F.pad(mel, (0, pad)), input_features_mask=mask)
hs, self._hs = self._hs, []
hs[-1] = out.last_hidden_state
B = wav.shape[0]
hs = [h.reshape(B, -1, h.shape[-1]) for h in hs]
return torch.stack(hs, 2).to(torch.bfloat16)
|