| 228 | |
| 229 | class MelSpec(nn.Module): |
| 230 | def __init__( |
| 231 | self, |
| 232 | n_fft=1024, |
| 233 | hop_length=256, |
| 234 | win_length=1024, |
| 235 | n_mel_channels=100, |
| 236 | target_sample_rate=24_000, |
| 237 | mel_spec_type="vocos", |
| 238 | ): |
| 239 | super().__init__() |
| 240 | assert mel_spec_type in ["vocos", "bigvgan", "bigvgan_qwen", "conformer", "whisper"], print("We only support four extract mel backend: vocos, bigvgan, bigvgan_qwen or conformer") |
| 241 | |
| 242 | self.n_fft = n_fft |
| 243 | self.hop_length = hop_length |
| 244 | self.win_length = win_length |
| 245 | self.n_mel_channels = n_mel_channels |
| 246 | self.target_sample_rate = target_sample_rate |
| 247 | |
| 248 | if mel_spec_type == "vocos": |
| 249 | self.extractor = get_vocos_mel_spectrogram |
| 250 | elif mel_spec_type == "bigvgan": |
| 251 | self.extractor = get_bigvgan_mel_spectrogram |
| 252 | elif mel_spec_type == "bigvgan_qwen": |
| 253 | self.extractor = get_bigvgan_qwen_mel_spectrogram |
| 254 | elif mel_spec_type == "conformer": |
| 255 | self.extractor = get_conformer_mel_spectrogram |
| 256 | elif mel_spec_type == "whisper": |
| 257 | self.extractor = get_whisper_mel_spectrogram |
| 258 | |
| 259 | self.register_buffer("dummy", torch.tensor(0), persistent=False) |
| 260 | |
| 261 | def forward(self, wav): |
| 262 | if self.dummy.device != wav.device: |