import sys from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoFeatureExtractor, AutoModel, WhisperModel, WhisperProcessor _THIS_DIR = Path(__file__).resolve().parent if str(_THIS_DIR) not in sys.path: sys.path.insert(0, str(_THIS_DIR)) from _checkpoint_mixin import load_audio_encoder_checkpoint # noqa: E402 def length_to_mask(lengths: torch.Tensor, max_len: int | None = None) -> torch.Tensor: if max_len is None: max_len = int(lengths.max().item()) idx = torch.arange(max_len, device=lengths.device).unsqueeze(0) return (idx < lengths.unsqueeze(1)).long() def _audio_lengths(audio: torch.Tensor, audio_attention_mask: torch.Tensor | None) -> torch.Tensor: if audio_attention_mask is None: return torch.full((audio.shape[0],), audio.shape[-1], device=audio.device, dtype=torch.long) return audio_attention_mask.sum(-1).to(torch.long) def _ceil_div(a: torch.Tensor, b: int) -> torch.Tensor: return (a + (b - 1)) // b def _resample_mask(mask: torch.Tensor, target_len: int) -> torch.Tensor: if mask.shape[1] == target_len: return mask out = F.interpolate(mask.float().unsqueeze(1), size=target_len, mode="nearest") return out.squeeze(1).to(torch.long) def _align_time(x: torch.Tensor, target_len: int) -> torch.Tensor: if x.shape[1] == target_len: return x x_t = x.transpose(1, 2) x_t = F.interpolate(x_t, size=target_len, mode="nearest") return x_t.transpose(1, 2) class DKU_WHU2_Encoder(torch.nn.Module): """Softmax2 token gate + log-mel/STFT residual.""" def __init__( self, dasheng_model_name: str = "mispeech/dasheng-base", whisper_model_name: str = "openai/whisper-base", gate_beta: float = 0.1, gate_align: str = "project", gate_target_dim: int | None = None, target_len: str = "min", temperature: float = 1.0, init_backbones: bool = True, mock_dasheng_dim: int = 768, mock_whisper_dim: int = 768, n_fft: int = 1024, hop_length: int = 640, n_mels: int = 128, use_residual_aux: bool = False, w_recon: float = 1.0, w_decor: float = 1e-3, checkpoint_dir: str | None = "checkpoint-10000", checkpoint_step: int | None = 50000, load_pretrained: bool = True, load_strict: bool = False, router_only: bool = False, verbose: bool = True, ) -> None: super().__init__() self.sampling_rate = 16000 self.hop_size_in_ms = 40 self._target_len = str(target_len) self._gate_align = str(gate_align) self._temperature = float(temperature) self._use_residual_aux = bool(use_residual_aux) self.router_w_recon = float(w_recon) self.router_w_decor = float(w_decor) self._init_backbones = bool(init_backbones) if self._init_backbones: self._dasheng_processor = AutoFeatureExtractor.from_pretrained(dasheng_model_name, trust_remote_code=True) self._dasheng_model = AutoModel.from_pretrained( dasheng_model_name, trust_remote_code=True, low_cpu_mem_usage=False, device_map=None, ) self._whisper_processor = WhisperProcessor.from_pretrained(whisper_model_name) self._whisper_encoder = WhisperModel.from_pretrained( whisper_model_name, low_cpu_mem_usage=False, device_map=None, ).get_encoder() dasheng_dim = self._infer_dasheng_dim() whisper_dim = int(self._whisper_encoder.config.d_model) else: self._dasheng_processor = None self._dasheng_model = None self._whisper_processor = None self._whisper_encoder = None dasheng_dim = int(mock_dasheng_dim) whisper_dim = int(mock_whisper_dim) self._dasheng_dim = int(dasheng_dim) self._whisper_dim = int(whisper_dim) target_dim = int(gate_target_dim) if gate_target_dim is not None else int(min(dasheng_dim, whisper_dim)) if target_dim <= 0: raise ValueError("gate_target_dim must be positive") self._gate_dim = target_dim self.register_buffer("router_beta", torch.tensor(float(gate_beta), dtype=torch.float32), persistent=True) self.router_proj_dasheng: nn.Module | None = None self.router_proj_whisper: nn.Module | None = None if self._gate_align not in ("truncate", "project"): raise ValueError("gate_align must be one of: truncate, project") if self._gate_align == "project": if int(dasheng_dim) != target_dim: self.router_proj_dasheng = nn.Linear(int(dasheng_dim), target_dim, bias=False) if int(whisper_dim) != target_dim: self.router_proj_whisper = nn.Linear(int(whisper_dim), target_dim, bias=False) self.router_gate_linear = nn.Linear(2 * target_dim, 2) self.router_fusion_ln = nn.LayerNorm(target_dim) self.router_sum_ln = nn.LayerNorm(target_dim) # Optional residual-orthogonalized aux branch. self.router_w2d: nn.Module | None = None self.router_pr: nn.Module | None = None if self._use_residual_aux: self.router_w2d = nn.Linear(self._whisper_dim, self._dasheng_dim, bias=False) self.router_pr = nn.Linear(self._dasheng_dim, self._whisper_dim) self._n_fft = int(n_fft) self._hop_length = int(hop_length) self._n_mels = int(n_mels) self.router_spec_proj = nn.Linear(self._n_mels, target_dim) self.router_spec_alpha = nn.Parameter(torch.zeros(1)) self.router_spec_ln = nn.LayerNorm(target_dim) self.output_dim = target_dim self.aux_loss = torch.tensor(0.0) self.router_reg_loss = torch.tensor(0.0) self.recon_loss = torch.tensor(0.0) self.decor_loss = torch.tensor(0.0) self.aux_loss_total = torch.tensor(0.0) self.aux_items: dict[str, torch.Tensor] = {} mel_fb = self._build_mel_filterbank(self._n_fft, self._n_mels, self.sampling_rate) if mel_fb is not None: self.register_buffer("mel_fb", mel_fb, persistent=False) else: self.mel_fb = None if load_pretrained: load_audio_encoder_checkpoint( self, checkpoint_dir, checkpoint_step=checkpoint_step, strict=load_strict, router_only=router_only, verbose=verbose, ) def _build_mel_filterbank(self, n_fft: int, n_mels: int, sr: int) -> torch.Tensor | None: try: import torchaudio.functional as AF # type: ignore fb = AF.melscale_fbanks( n_freqs=n_fft // 2 + 1, f_min=0.0, f_max=float(sr // 2), n_mels=n_mels, sample_rate=sr, ) return fb except Exception: return None def _infer_dasheng_dim(self) -> int: cfg = getattr(self._dasheng_model, "config", None) if cfg is None: raise ValueError("Dasheng model has no config; cannot infer output_dim") for attr in ("hidden_size", "d_model", "embed_dim"): if hasattr(cfg, attr): return int(getattr(cfg, attr)) if hasattr(cfg, "encoder_kwargs") and isinstance(cfg.encoder_kwargs, dict) and "embed_dim" in cfg.encoder_kwargs: return int(cfg.encoder_kwargs["embed_dim"]) if hasattr(cfg, "encoder_kwargs") and isinstance(cfg.encoder_kwargs, dict) and "d_model" in cfg.encoder_kwargs: return int(cfg.encoder_kwargs["d_model"]) raise ValueError("Could not infer Dasheng embedding dim from config") def _dasheng_forward( self, audio: torch.Tensor, audio_attention_mask: torch.Tensor | None ) -> tuple[torch.Tensor, torch.Tensor]: if self._dasheng_model is None or self._dasheng_processor is None: raise RuntimeError("Dasheng backbone not initialized.") features = self._dasheng_processor(audio, return_tensors="pt") model_device = next(self._dasheng_model.parameters()).device features = {k: v.to(model_device) if isinstance(v, torch.Tensor) else v for k, v in features.items()} out = self._dasheng_model(**features) if hasattr(out, "last_hidden_state") and isinstance(out.last_hidden_state, torch.Tensor): feats = out.last_hidden_state elif hasattr(out, "hidden_states"): hs = out.hidden_states feats = hs[-1] if isinstance(hs, (tuple, list)) else hs elif isinstance(out, (tuple, list)) and len(out) > 0 and isinstance(out[0], torch.Tensor): feats = out[0] else: raise ValueError("Unexpected Dasheng model output; cannot get features") t1 = int(feats.shape[1]) hop_samples = int(self.sampling_rate * self.hop_size_in_ms / 1000) lengths = _audio_lengths( audio.to(feats.device), audio_attention_mask.to(feats.device) if audio_attention_mask is not None else None, ) feat_lens = torch.clamp(_ceil_div(lengths, hop_samples), min=1, max=t1) mask = length_to_mask(feat_lens, max_len=t1).to(feats.device) return feats, mask def _whisper_forward( self, audio: torch.Tensor, audio_attention_mask: torch.Tensor | None ) -> tuple[torch.Tensor, torch.Tensor]: if self._whisper_encoder is None or self._whisper_processor is None: raise RuntimeError("Whisper backbone not initialized.") audio_list = [a.detach().cpu().numpy() for a in audio] if audio_attention_mask is None: audio_lens = torch.tensor([a.shape[-1] for a in audio_list], dtype=torch.long) else: audio_lens = audio_attention_mask.sum(-1).detach().cpu().to(torch.long) hop_samples = int(self.sampling_rate * self.hop_size_in_ms / 1000) feature_lengths = torch.clamp(_ceil_div(audio_lens, hop_samples), min=1) trim_length = int(feature_lengths.max().item()) attention_mask = length_to_mask(feature_lengths, max_len=trim_length) feats = self._whisper_processor(audio_list, sampling_rate=self.sampling_rate, return_tensors="pt") model_device = next(self._whisper_encoder.parameters()).device feats = {k: v.to(model_device) if isinstance(v, torch.Tensor) else v for k, v in feats.items()} out = self._whisper_encoder(**feats).last_hidden_state out = out[:, :trim_length, :] attention_mask = attention_mask[:, : out.shape[1]].to(out.device) return out, attention_mask def _align_gate_dims(self, d: torch.Tensor, w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: if self._gate_align == "project": if self.router_proj_dasheng is not None: d = self.router_proj_dasheng(d) if self.router_proj_whisper is not None: w = self.router_proj_whisper(w) return d, w target = int(self._gate_dim) return d[..., :target], w[..., :target] def _softmax2_fuse(self, d: torch.Tensor, w: torch.Tensor) -> torch.Tensor: gate_inp = torch.cat([d, w], dim=-1) logits = self.router_gate_linear(gate_inp) / self._temperature weights = torch.softmax(logits, dim=-1) scaled = weights * self.router_beta w_d = scaled[..., 0:1] w_w = scaled[..., 1:2] fused = w_d * d + w_w * w fused = self.router_fusion_ln(fused) fused = self.router_sum_ln(fused + d + w) return fused def _compute_log_mel(self, audio: torch.Tensor) -> torch.Tensor: window = torch.hann_window(self._n_fft, device=audio.device) spec = torch.stft( audio, n_fft=self._n_fft, hop_length=self._hop_length, win_length=self._n_fft, window=window, center=True, return_complex=True, ) mag = spec.abs() ** 2 if self.mel_fb is not None: mel = torch.matmul(mag.transpose(1, 2), self.mel_fb.to(mag.device)).transpose(1, 2) else: mel = F.interpolate(mag.unsqueeze(1), size=(self._n_mels, mag.shape[2]), mode="nearest").squeeze(1) mel = mel.clamp_min(1e-10).log() return mel.transpose(1, 2) def _masked_mean(self, values: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: weights = mask.to(dtype=values.dtype) denom = weights.sum().clamp_min(1.0) return (values * weights).sum() / denom def _cross_cov_decorrelation_loss( self, h_w: torch.Tensor, r_w: torch.Tensor, mask: torch.Tensor, ) -> torch.Tensor: x = h_w.reshape(-1, h_w.shape[-1]) y = r_w.reshape(-1, r_w.shape[-1]) keep = mask.reshape(-1) > 0 if keep.any(): x = x[keep] y = y[keep] n = int(x.shape[0]) if n <= 1: return h_w.new_zeros((), dtype=torch.float32) x = x.float() y = y.float() x_centered = x - x.mean(dim=0, keepdim=True) y_centered = y - y.mean(dim=0, keepdim=True) cov = torch.matmul(x_centered.transpose(0, 1), y_centered) / float(max(n - 1, 1)) return cov.pow(2).mean() def _compute_residual_aux( self, h_d: torch.Tensor, h_w: torch.Tensor, mask: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: if self.router_w2d is None or self.router_pr is None: return h_w.new_zeros((), dtype=torch.float32), h_w.new_zeros((), dtype=torch.float32) h_d_hat = self.router_w2d(h_w) # [B,T,768] residual_d = h_d - h_d_hat r_w = self.router_pr(residual_d) # [B,T,512] recon_token = residual_d.float().pow(2).mean(dim=-1) recon_loss = self._masked_mean(recon_token, mask) decor_loss = self._cross_cov_decorrelation_loss(h_w, r_w, mask) return recon_loss, decor_loss def forward( self, audio: torch.Tensor, audio_attention_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: if not isinstance(audio, torch.Tensor): raise TypeError("audio must be a torch.Tensor") if audio.ndim != 2: raise ValueError("Expected audio shape [B, T]") dasheng_feat, dasheng_mask = self._dasheng_forward(audio, audio_attention_mask) whisper_feat, whisper_mask = self._whisper_forward(audio, audio_attention_mask) if self._target_len == "whisper": target_len = int(whisper_feat.shape[1]) elif self._target_len == "dasheng": target_len = int(dasheng_feat.shape[1]) else: target_len = int(min(dasheng_feat.shape[1], whisper_feat.shape[1])) dasheng_mask = _resample_mask(dasheng_mask, target_len) whisper_mask = _resample_mask(whisper_mask, target_len) mask = (dasheng_mask & whisper_mask).long() dasheng_feat = dasheng_feat[:, :target_len, :] whisper_feat = whisper_feat[:, :target_len, :] h_d_raw = dasheng_feat h_w_raw = whisper_feat dasheng_feat, whisper_feat = self._align_gate_dims(dasheng_feat, whisper_feat) fused = self._softmax2_fuse(dasheng_feat, whisper_feat) spec = self._compute_log_mel(audio) if spec.ndim != 3: raise ValueError("spec must be [B,T,F]") if spec.shape[1] != fused.shape[1]: spec = _align_time(spec, fused.shape[1]) spec_proj = self.router_spec_proj(spec) spec_proj = self.router_spec_ln(spec_proj) fused = fused + self.router_spec_alpha * spec_proj recon_loss, decor_loss = self._compute_residual_aux(h_d_raw, h_w_raw, mask) decor_w = float(getattr(self, "router_w_decor_runtime", self.router_w_decor)) aux_loss_total = fused.new_zeros((), dtype=torch.float32) if self._use_residual_aux: aux_loss_total = (self.router_w_recon * recon_loss) + (decor_w * decor_loss) self.recon_loss = recon_loss self.decor_loss = decor_loss self.aux_loss_total = aux_loss_total self.aux_items = { "recon_loss": recon_loss, "decor_loss": decor_loss, "aux_loss_total": aux_loss_total, } return fused, mask