File size: 6,412 Bytes
7df3ebb | 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | # -*- coding: utf-8 -*-
import copy
from typing import List
from typing import Optional
import numpy as np
import torch
from transformers.audio_utils import spectrogram
from transformers.audio_utils import window_function
from transformers.models.whisper.feature_extraction_whisper import WhisperFeatureExtractor
class MiniCPMAAudioProcessor(WhisperFeatureExtractor):
"""
在 WhisperFeatureExtractor 基础上：
- 支持 dynamic_log_norm（原版 max-8dB，可调 dynamic_range_db）
- 或固定对数下限 log_floor_db（例如 -10dB）
- 这是由于我们需要做 streaming 方案，在这个方案下，没办法进行动态设置
- 这个 可以在中途修改，通过 set_dynamic_log_norm
两条路径（torch / numpy）保持一致的裁剪与缩放顺序：
log10 -> (动态/固定下限裁剪) -> (+4)/4
"""
def __init__(
self,
*args,
dynamic_log_norm: bool = True,
dynamic_range_db: float = 8.0,
log_floor_db: float = -10.0,
**kwargs,
):
super().__init__(*args, **kwargs)
self.dynamic_log_norm = bool(dynamic_log_norm)
self.dynamic_range_db = float(dynamic_range_db)
self.log_floor_db = float(log_floor_db)
def set_spac_log_norm(
self,
dynamic_range_db: Optional[float] = None,
log_floor_db: Optional[float] = None,
*,
inplace: bool = True,
) -> "MiniCPMAAudioProcessor":
"""
热更新动态/固定下限策略。
Args:
enabled: True=使用动态阈值(max - dynamic_range_db)，False=使用固定下限 log_floor_db。
传 None 表示保持不变。
dynamic_range_db: 动态范围（dB），仅在 enabled=True 时生效。None 表示保持不变。
log_floor_db: 固定对数下限（dB，通常<=0），仅在 enabled=False 时使用。None 表示保持不变。
inplace: True 直接修改当前实例；False 返回一个浅拷贝并在其上修改。
Returns:
self 或新的实例（当 inplace=False 时）。
"""
target = self if inplace else copy.copy(self)
if dynamic_range_db is not None:
val = float(dynamic_range_db)
if val < 0:
raise ValueError("dynamic_range_db must be >= 0.")
target.dynamic_log_norm = True # 给出显式提示：设置了该值就走动态模式
target.dynamic_range_db = val
if log_floor_db is not None:
val = float(log_floor_db)
# 一般 log10(mel) 最大不超过 ~0dB，floor 设为<=0更合理；这里做宽松校验
if val > 0:
raise ValueError("log_floor_db should be <= 0 (log10 scale).")
target.dynamic_log_norm = False # 设置了该值就走固定下限模式
target.log_floor_db = val
return target
# ------------ NumPy 实现：覆写 ----------
def _np_extract_fbank_features(self, waveform_batch: np.ndarray, device: str) -> np.ndarray:
"""
与上游一致的 NumPy 版本，但将 max-8dB 替换为可配置的 动态/固定 下限裁剪。
"""
if device != "cpu":
raise ValueError(
f"Got device `{device}` for feature extraction, but feature extraction on CUDA accelerator "
"devices requires torch. Set device='cpu' or install torch."
)
log_spec_batch: List[np.ndarray] = []
for waveform in waveform_batch:
# 生成 log10 Mel
log_spec = spectrogram(
waveform,
window_function(self.n_fft, "hann"),
frame_length=self.n_fft,
hop_length=self.hop_length,
power=2.0,
dither=self.dither,
mel_filters=self.mel_filters,
log_mel="log10",
)
# 与上游保持一致：去掉最后一帧
log_spec = log_spec[:, :-1]
# 动态/固定裁剪
if self.dynamic_log_norm:
threshold = log_spec.max() - self.dynamic_range_db
log_spec = np.maximum(log_spec, threshold)
else:
log_spec = np.maximum(log_spec, self.log_floor_db)
# 与 Whisper 一致的线性缩放
log_spec = (log_spec + 4.0) / 4.0
log_spec_batch.append(log_spec)
return np.array(log_spec_batch)
# ------------ PyTorch 实现：覆写 ----------
def _torch_extract_fbank_features(self, waveform: np.ndarray, device: str = "cpu") -> np.ndarray:
if torch is None:
raise RuntimeError("PyTorch 未安装，无法在 GPU 上计算 STFT。")
waveform = torch.from_numpy(waveform).to(device, torch.float32)
window = torch.hann_window(self.n_fft, device=device)
if self.dither != 0.0:
waveform = waveform + self.dither * torch.randn_like(waveform)
stft = torch.stft(waveform, n_fft=self.n_fft, hop_length=self.hop_length, window=window, return_complex=True)
magnitudes = stft[..., :-1].abs() ** 2
mel_filters = torch.from_numpy(self.mel_filters).to(device, torch.float32) # [n_mels, 1+n_fft//2]
mel_spec = mel_filters.T @ magnitudes # [..., n_mels, T]
log_spec = torch.clamp(mel_spec, min=1e-10).log10() # ≤ 0
if self.dynamic_log_norm:
if waveform.dim() == 2:
max_val_t = log_spec.max(dim=2, keepdim=True)[0] # over T
max_val_bt = max_val_t.max(dim=1, keepdim=True)[0] # over mel
threshold = max_val_bt - self.dynamic_range_db
log_spec = torch.maximum(log_spec, threshold)
else:
threshold = log_spec.max() - self.dynamic_range_db
log_spec = torch.maximum(log_spec, threshold)
else:
floor_tensor = torch.tensor(self.log_floor_db, dtype=log_spec.dtype, device=log_spec.device)
log_spec = torch.maximum(log_spec, floor_tensor)
log_spec = (log_spec + 4.0) / 4.0
if device != "cpu":
log_spec = log_spec.detach().cpu()
return log_spec.numpy()
def process(self, *args, **kwargs):
"""Alias of __call__ for convenience."""
return self.__call__(*args, **kwargs)
|