| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import json |
| 4 | import math |
| 5 | from functools import lru_cache |
| 6 | |
| 7 | import torch |
| 8 | import torch.nn.functional as F |
| 9 | from torch import nn |
| 10 | from transformers import WhisperConfig |
| 11 | from transformers.audio_utils import mel_filter_bank |
| 12 | |
| 13 | from .audio_encoder import DotsSpeechEncoder |
| 14 | |
| 15 | SAMPLE_RATE = 16000 |
| 16 | N_FFT = 400 |
| 17 | HOP_LENGTH = 160 |
| 18 | DEFAULT_CHUNK_LENGTH_S = 60 |
| 19 | DEFAULT_CONV_TEMPORAL_STRIDE = 8 |
| 20 | DEFAULT_MERGE_FACTOR = 1 |
| 21 | N_SAMPLES = DEFAULT_CHUNK_LENGTH_S * SAMPLE_RATE |
| 22 | STRIDE = HOP_LENGTH * DEFAULT_CONV_TEMPORAL_STRIDE * DEFAULT_MERGE_FACTOR |
| 23 | ENCODER_SEQ_LEN = N_SAMPLES // HOP_LENGTH // 2 |
| 24 | _AUDIO_FORWARD_MAX_SEGMENTS = 16 |
| 25 | |
| 26 | |
| 27 | class Dots3NoteAudioConfig: |
| 28 | def __init__(self, **kwargs): |
| 29 | self.encoder_type = kwargs.get("encoder_type", "dots") |
| 30 | self.whisper_config = kwargs.get("whisper_config", {}) |
| 31 | self.whisper_adapter_in_dim = kwargs.get( |
| 32 | "whisper_adapter_in_dim", kwargs.get("adapter_in_dim", 1280) |
| 33 | ) |
| 34 | self.whisper_adapter_out_dim = kwargs.get( |
| 35 | "whisper_adapter_out_dim", kwargs.get("adapter_out_dim", 2048) |
| 36 | ) |
| 37 | self.sampling_rate = kwargs.get("sampling_rate", SAMPLE_RATE) |
| 38 | self.audio_comp_start = kwargs.get("audio_comp_start", "<|audio_comp_start|>") |
| 39 | self.audio_comp_span = kwargs.get("audio_comp_span", "<|audio_comp_pad|>") |
| 40 | self.audio_comp_end = kwargs.get("audio_comp_end", "<|audio_comp_end|>") |
| 41 | self.merge_factor = kwargs.get("merge_factor", DEFAULT_MERGE_FACTOR) |
| 42 | self.chunk_seconds = kwargs.get("chunk_seconds", DEFAULT_CHUNK_LENGTH_S) |
| 43 | self.use_conv2d_stem = kwargs.get("use_conv2d_stem", True) |
| 44 | self.use_rope = kwargs.get("use_rope", True) |
| 45 | self.use_rms_norm = kwargs.get("use_rms_norm", True) |
| 46 | self.use_causal = kwargs.get("use_causal", False) |
| 47 | self.downsample_hidden_size = kwargs.get("downsample_hidden_size", 480) |
| 48 | self.conv_chunksize = kwargs.get("conv_chunksize", 500) |
| 49 | self.conv_stem_gradient_checkpointing = kwargs.get( |
| 50 | "conv_stem_gradient_checkpointing", False |
| 51 | ) |
| 52 | self.conv_bucket_step = kwargs.get("conv_bucket_step") |
| 53 | self.conv_bucket_max_elements = kwargs.get("conv_bucket_max_elements") |
| 54 | self.rope_parameters = kwargs.get( |
| 55 | "rope_parameters", |
| 56 | { |
| 57 | "partial_rotary_factor": 0.5, |
| 58 | "rope_theta": 10000.0, |
| 59 | "rope_type": "default", |
| 60 | }, |
| 61 | ) |
| 62 | |
| 63 | @classmethod |
| 64 | def from_json(cls, path): |
| 65 | with open(path) as f: |
| 66 | return cls(**json.load(f)) |
| 67 | |
| 68 | @property |
| 69 | def conv_temporal_stride(self) -> int: |
| 70 | return 8 if self.use_conv2d_stem else 2 |
| 71 | |
| 72 | @property |
| 73 | def token_stride(self) -> int: |
| 74 | return HOP_LENGTH * self.conv_temporal_stride * self.merge_factor |
| 75 | |
| 76 | @property |
| 77 | def chunk_samples(self) -> int: |
| 78 | return int(self.chunk_seconds * SAMPLE_RATE) |
| 79 | |
| 80 | @property |
| 81 | def chunk_mel_frames(self) -> int: |
| 82 | return int(self.chunk_seconds * 100) |
| 83 | |
| 84 | |
| 85 | def pad_or_trim(array, length=N_SAMPLES, axis=-1): |
| 86 | if array.shape[axis] > length: |
| 87 | array = array.index_select( |
| 88 | dim=axis, index=torch.arange(length, device=array.device) |
| 89 | ) |
| 90 | if array.shape[axis] < length: |
| 91 | pad_widths = [(0, 0)] * array.ndim |
| 92 | pad_widths[axis] = (0, length - array.shape[axis]) |
| 93 | array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes]) |
| 94 | return array |
| 95 | |
| 96 | |
| 97 | @lru_cache(maxsize=4) |
| 98 | def _mel_filters(device, n_mels=128): |
| 99 | filters = mel_filter_bank( |
| 100 | num_frequency_bins=1 + N_FFT // 2, |
| 101 | num_mel_filters=n_mels, |
| 102 | min_frequency=0.0, |
| 103 | max_frequency=float(SAMPLE_RATE) / 2.0, |
| 104 | sampling_rate=SAMPLE_RATE, |
| 105 | norm="slaney", |
| 106 | mel_scale="slaney", |
| 107 | ) |
| 108 | return torch.from_numpy(filters).T.contiguous().float().to(device) |
| 109 | |
| 110 | |
| 111 | @lru_cache(maxsize=4) |
| 112 | def _hann_window(device): |
| 113 | # Generate on CPU (default) then move, to preserve the exact reference |
| 114 | # window values. Direct on-device generation changes numerics slightly. |
| 115 | return torch.hann_window(N_FFT).to(device) |
| 116 | |
| 117 | |
| 118 | def log_mel_spectrogram(audio, n_mels=128): |
| 119 | window = _hann_window(audio.device) |
| 120 | stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True) |
| 121 | magnitudes = stft[..., :-1].abs() ** 2 |
| 122 | filters = _mel_filters(audio.device, n_mels) |
| 123 | mel_spec = filters @ magnitudes |
| 124 | log_spec = torch.clamp(mel_spec, min=1e-10).log10() |
| 125 | log_spec = torch.maximum(log_spec, log_spec.max() - 8.0) |
| 126 | log_spec = (log_spec + 4.0) / 4.0 |
| 127 | return log_spec |
| 128 | |
| 129 | |
| 130 | def batched_log_mel_spectrogram(audio, n_mels=128): |
| 131 | window = _hann_window(audio.device) |
| 132 | stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True) |
| 133 | magnitudes = stft[..., :-1].abs() ** 2 |
| 134 | filters = _mel_filters(audio.device, n_mels) |
| 135 | mel_spec = filters @ magnitudes |
| 136 | log_spec = torch.clamp(mel_spec, min=1e-10).log10() |
| 137 | log_spec = torch.maximum(log_spec, log_spec.amax(dim=(-2, -1), keepdim=True) - 8.0) |
| 138 | return (log_spec + 4.0) / 4.0 |
| 139 | |
| 140 | |
| 141 | def compute_audio_token_length( |
| 142 | num_samples, |
| 143 | *, |
| 144 | sample_rate: int = SAMPLE_RATE, |
| 145 | chunk_seconds: int = DEFAULT_CHUNK_LENGTH_S, |
| 146 | hop_length: int = HOP_LENGTH, |
| 147 | conv_temporal_stride: int = DEFAULT_CONV_TEMPORAL_STRIDE, |
| 148 | merge_factor: int = DEFAULT_MERGE_FACTOR, |
| 149 | ): |
| 150 | stride = hop_length * conv_temporal_stride * merge_factor |
| 151 | total = 0 |
| 152 | time_step = 0 |
| 153 | while time_step * sample_rate < num_samples: |
| 154 | chunk_len = min( |
| 155 | num_samples - time_step * sample_rate, chunk_seconds * sample_rate |
| 156 | ) |
| 157 | total += math.ceil(chunk_len / stride) |
| 158 | time_step += chunk_seconds |
| 159 | return total |
| 160 | |
| 161 | |
| 162 | class DotsEncoderWithMask(nn.Module): |
| 163 | def __init__(self, config: Dots3NoteAudioConfig): |
| 164 | super().__init__() |
| 165 | whisper_config = WhisperConfig(**config.whisper_config) |
| 166 | whisper_config.use_rope = config.use_rope |
| 167 | whisper_config.rope_parameters = config.rope_parameters |
| 168 | whisper_config.use_rms_norm = config.use_rms_norm |
| 169 | whisper_config.use_causal = config.use_causal |
| 170 | whisper_config.use_conv2d_stem = config.use_conv2d_stem |
| 171 | whisper_config.downsample_hidden_size = config.downsample_hidden_size |
| 172 | whisper_config.conv_chunksize = config.conv_chunksize |
| 173 | whisper_config.conv_stem_gradient_checkpointing = ( |
| 174 | config.conv_stem_gradient_checkpointing |
| 175 | ) |
| 176 | whisper_config.conv_bucket_step = config.conv_bucket_step |
| 177 | whisper_config.conv_bucket_max_elements = config.conv_bucket_max_elements |
| 178 | |
| 179 | self.speech_encoder = DotsSpeechEncoder(whisper_config) |
| 180 | self.merge_factor = config.merge_factor |
| 181 | self.chunk_seconds = config.chunk_seconds |
| 182 | self.chunk_samples = config.chunk_samples |
| 183 | self.chunk_mel_frames = config.chunk_mel_frames |
| 184 | self.conv_temporal_stride = config.conv_temporal_stride |
| 185 | |
| 186 | @property |
| 187 | def device(self): |
| 188 | return next(self.speech_encoder.parameters()).device |
| 189 | |
| 190 | def _forward_speech_encoder( |
| 191 | self, |
| 192 | mel_features: torch.Tensor, |
| 193 | input_seq_lens: torch.Tensor, |
| 194 | audio_sample_lens: list[int], |
| 195 | ) -> torch.Tensor: |
| 196 | mel_features = mel_features.to(dtype=torch.bfloat16, device=self.device) |
| 197 | if mel_features.shape[0] <= _AUDIO_FORWARD_MAX_SEGMENTS: |
| 198 | return self.speech_encoder( |
| 199 | mel_features, |
| 200 | return_dict=True, |
| 201 | input_seq_lens=input_seq_lens, |
| 202 | audio_sample_lens=audio_sample_lens, |
| 203 | ).last_hidden_state |
| 204 | |
| 205 | chunks = [] |
| 206 | for start in range(0, mel_features.shape[0], _AUDIO_FORWARD_MAX_SEGMENTS): |
| 207 | end = start + _AUDIO_FORWARD_MAX_SEGMENTS |
| 208 | chunks.append( |
| 209 | self.speech_encoder( |
| 210 | mel_features[start:end], |
| 211 | return_dict=True, |
| 212 | input_seq_lens=input_seq_lens[start:end], |
| 213 | audio_sample_lens=audio_sample_lens[start:end], |
| 214 | ).last_hidden_state |
| 215 | ) |
| 216 | max_seq_len = max(chunk.shape[1] for chunk in chunks) |
| 217 | chunks = [ |
| 218 | F.pad(chunk, (0, 0, 0, max_seq_len - chunk.shape[1])) |
| 219 | if chunk.shape[1] < max_seq_len |
| 220 | else chunk |
| 221 | for chunk in chunks |
| 222 | ] |
| 223 | return torch.cat(chunks, dim=0) |
| 224 | |
| 225 | def encode_waveform(self, audio_waveform: torch.Tensor) -> torch.Tensor: |
| 226 | segments = [] |
| 227 | time_step = 0 |
| 228 | while time_step * SAMPLE_RATE < audio_waveform.shape[0]: |
| 229 | segments.append( |
| 230 | audio_waveform[ |
| 231 | time_step * SAMPLE_RATE : (time_step + self.chunk_seconds) |
| 232 | * SAMPLE_RATE |
| 233 | ] |
| 234 | ) |
| 235 | time_step += self.chunk_seconds |
| 236 | |
| 237 | mel_features = [] |
| 238 | token_lens = [] |
| 239 | audio_sample_lens = [] |
| 240 | for audio_segment in segments: |
| 241 | segment_length = audio_segment.shape[0] |
| 242 | token_len = (segment_length - 1) // ( |
| 243 | HOP_LENGTH * self.conv_temporal_stride * self.merge_factor |
| 244 | ) + 1 |
| 245 | pad_audio = pad_or_trim(audio_segment.flatten(), length=self.chunk_samples) |
| 246 | mel = log_mel_spectrogram(pad_audio) |
| 247 | assert mel.shape[1] == self.chunk_mel_frames |
| 248 | mel_features.append(mel) |
| 249 | token_lens.append(token_len) |
| 250 | audio_sample_lens.append(segment_length) |
| 251 | |
| 252 | mel_features = torch.stack(mel_features, dim=0) |
| 253 | # Keep input_seq_lens on CPU: the conv2d bucket path reads it via |
| 254 | # ``.item()`` and CPU scalars avoid device->host syncs. The encoder's |
| 255 | # varlen path moves it to the device itself. |
| 256 | input_seq_lens = torch.tensor(token_lens, dtype=torch.long) * self.merge_factor |
| 257 | audio_embedding = self._forward_speech_encoder( |
| 258 | mel_features, input_seq_lens, audio_sample_lens |
| 259 | ) |
| 260 | |
| 261 | chunk_embeddings = [] |
| 262 | for idx, token_len in enumerate(token_lens): |
| 263 | chunk_embeddings.append( |
| 264 | audio_embedding[idx, : token_len * self.merge_factor, :] |
| 265 | ) |
| 266 | return torch.cat(chunk_embeddings, dim=0).unsqueeze(0) |
| 267 | |
| 268 | def encode_waveforms(self, audio_waveforms) -> list[torch.Tensor]: |
| 269 | stride = HOP_LENGTH * self.conv_temporal_stride * self.merge_factor |
| 270 | segment_token_lengths = [] |
| 271 | segment_sample_lengths = [] |
| 272 | audio_segment_counts = [] |
| 273 | padded_segments = [] |
| 274 | |
| 275 | for audio_waveform in audio_waveforms: |
| 276 | segment_count = 0 |
| 277 | time_step = 0 |
| 278 | while time_step * SAMPLE_RATE < audio_waveform.shape[0]: |
| 279 | segment = audio_waveform[ |
| 280 | time_step * SAMPLE_RATE : (time_step + self.chunk_seconds) |
| 281 | * SAMPLE_RATE |
| 282 | ] |
| 283 | segment_length = segment.shape[0] |
| 284 | padded_segments.append( |
| 285 | pad_or_trim(segment.flatten(), length=self.chunk_samples) |
| 286 | ) |
| 287 | segment_token_lengths.append((segment_length - 1) // stride + 1) |
| 288 | segment_sample_lengths.append(segment_length) |
| 289 | segment_count += 1 |
| 290 | time_step += self.chunk_seconds |
| 291 | audio_segment_counts.append(segment_count) |
| 292 | |
| 293 | mel_features = batched_log_mel_spectrogram(torch.stack(padded_segments)) |
| 294 | assert mel_features.shape[-1] == self.chunk_mel_frames |
| 295 | input_seq_lens = ( |
| 296 | torch.tensor(segment_token_lengths, dtype=torch.long) * self.merge_factor |
| 297 | ) |
| 298 | audio_embedding = self._forward_speech_encoder( |
| 299 | mel_features, input_seq_lens, segment_sample_lengths |
| 300 | ) |
| 301 | |
| 302 | embeddings = [] |
| 303 | segment_idx = 0 |
| 304 | for segment_count in audio_segment_counts: |
| 305 | chunks = [] |
| 306 | for _ in range(segment_count): |
| 307 | token_length = segment_token_lengths[segment_idx] |
| 308 | chunks.append( |
| 309 | audio_embedding[segment_idx, : token_length * self.merge_factor, :] |
| 310 | ) |
| 311 | segment_idx += 1 |
| 312 | embeddings.append(torch.cat(chunks).unsqueeze(0)) |
| 313 | return embeddings |
| 314 | |
| 315 | |
| 316 | class AudioAdapter(nn.Module): |
| 317 | def __init__(self, in_dim, out_dim): |
| 318 | super().__init__() |
| 319 | self.proj = nn.Sequential( |
| 320 | nn.LayerNorm(in_dim), |
| 321 | nn.Linear(in_dim, out_dim), |
| 322 | nn.GELU(), |
| 323 | nn.Linear(out_dim, out_dim), |
| 324 | ) |
| 325 | |
| 326 | def forward(self, x): |
| 327 | return self.proj(x) |
| 328 | |
| 329 | |
| 330 | class Dots3NoteAudioModel(nn.Module): |
| 331 | def __init__(self, config: Dots3NoteAudioConfig): |
| 332 | super().__init__() |
| 333 | self.config = config |
| 334 | if config.encoder_type != "dots": |
| 335 | raise ValueError("Dots3Note only supports encoder_type='dots'") |
| 336 | self.merge_factor = config.merge_factor |
| 337 | self.audio_adapter = AudioAdapter( |
| 338 | config.whisper_adapter_in_dim, |
| 339 | config.whisper_adapter_out_dim, |
| 340 | ) |
| 341 | |
| 342 | self.dots_encoder = DotsEncoderWithMask(config) |
| 343 | |
| 344 | @property |
| 345 | def device(self): |
| 346 | return self.dots_encoder.device |
| 347 | |
| 348 | def _merge_embeddings(self, embedding: torch.Tensor) -> torch.Tensor: |
| 349 | if self.merge_factor <= 1: |
| 350 | return embedding |
| 351 | return embedding.reshape( |
| 352 | embedding.shape[0], |
| 353 | embedding.shape[1] // self.merge_factor, |
| 354 | embedding.shape[2] * self.merge_factor, |
| 355 | ) |
| 356 | |
| 357 | def _encode_single_audio(self, audio_waveform): |
| 358 | embedding = self.dots_encoder.encode_waveform(audio_waveform) |
| 359 | embedding = self._merge_embeddings(embedding) |
| 360 | embedding = self.audio_adapter(embedding) |
| 361 | return embedding.squeeze(0) |
| 362 | |
| 363 | def _split_waveforms(self, audio_inputs, lengths): |
| 364 | # Convert lengths to CPU ints once so the per-audio slicing below stays |
| 365 | # on the host and does not trigger a device->host sync per audio. |
| 366 | if isinstance(audio_inputs, list): |
| 367 | return audio_inputs |
| 368 | lengths_list = lengths.tolist() |
| 369 | waveforms = [] |
| 370 | audio_start = 0 |
| 371 | for length in lengths_list: |
| 372 | waveforms.append(audio_inputs[audio_start : audio_start + length]) |
| 373 | audio_start += length |
| 374 | return waveforms |
| 375 | |
| 376 | def _forward_batched(self, waveforms): |
| 377 | raw_embeddings = self.dots_encoder.encode_waveforms(waveforms) |
| 378 | merged = [ |
| 379 | self._merge_embeddings(embedding).squeeze(0) for embedding in raw_embeddings |
| 380 | ] |
| 381 | token_lengths = [embedding.shape[0] for embedding in merged] |
| 382 | packed = self.audio_adapter(torch.cat(merged)) |
| 383 | return packed, token_lengths |
| 384 | |
| 385 | def forward(self, audio_inputs, lengths): |
| 386 | waveforms = self._split_waveforms(audio_inputs, lengths) |
| 387 | return self._forward_batched(waveforms) |