| |
|
|
| 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) |
| |
| 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 |
|
|
| |
| 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: |
| |
| 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) |
|
|
| |
| log_spec = (log_spec + 4.0) / 4.0 |
|
|
| log_spec_batch.append(log_spec) |
|
|
| return np.array(log_spec_batch) |
|
|
| |
| 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) |
| mel_spec = mel_filters.T @ magnitudes |
|
|
| log_spec = torch.clamp(mel_spec, min=1e-10).log10() |
|
|
| if self.dynamic_log_norm: |
| if waveform.dim() == 2: |
| max_val_t = log_spec.max(dim=2, keepdim=True)[0] |
| max_val_bt = max_val_t.max(dim=1, keepdim=True)[0] |
| 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) |
|
|
