vllm.models.dots3_note.nvidia.audio ¶

class DotsEncoderWithMask(nn.Module):
    def __init__(self, config: Dots3NoteAudioConfig):
        super().__init__()
        whisper_config = WhisperConfig(**config.whisper_config)
        whisper_config.use_rope = config.use_rope
        whisper_config.rope_parameters = config.rope_parameters
        whisper_config.use_rms_norm = config.use_rms_norm
        whisper_config.use_causal = config.use_causal
        whisper_config.use_conv2d_stem = config.use_conv2d_stem
        whisper_config.downsample_hidden_size = config.downsample_hidden_size
        whisper_config.conv_chunksize = config.conv_chunksize
        whisper_config.conv_stem_gradient_checkpointing = (
            config.conv_stem_gradient_checkpointing
        )
        whisper_config.conv_bucket_step = config.conv_bucket_step
        whisper_config.conv_bucket_max_elements = config.conv_bucket_max_elements

        self.speech_encoder = DotsSpeechEncoder(whisper_config)
        self.merge_factor = config.merge_factor
        self.chunk_seconds = config.chunk_seconds
        self.chunk_samples = config.chunk_samples
        self.chunk_mel_frames = config.chunk_mel_frames
        self.conv_temporal_stride = config.conv_temporal_stride

    @property
    def device(self):
        return next(self.speech_encoder.parameters()).device

    def _forward_speech_encoder(
        self,
        mel_features: torch.Tensor,
        input_seq_lens: torch.Tensor,
        audio_sample_lens: list[int],
    ) -> torch.Tensor:
        """Run the eager speech encoder without server-side slicing/batching."""
        mel_features = mel_features.to(dtype=torch.bfloat16, device=self.device)
        return self.speech_encoder(
            mel_features,
            return_dict=True,
            input_seq_lens=input_seq_lens,
            audio_sample_lens=audio_sample_lens,
        ).last_hidden_state

    def encode_waveform(self, audio_waveform: torch.Tensor) -> torch.Tensor:
        segments = []
        time_step = 0
        while time_step * SAMPLE_RATE < audio_waveform.shape[0]:
            segments.append(
                audio_waveform[
                    time_step * SAMPLE_RATE : (time_step + self.chunk_seconds)
                    * SAMPLE_RATE
                ]
            )
            time_step += self.chunk_seconds

        mel_features = []
        token_lens = []
        audio_sample_lens = []
        for audio_segment in segments:
            segment_length = audio_segment.shape[0]
            token_len = (segment_length - 1) // (
                HOP_LENGTH * self.conv_temporal_stride * self.merge_factor
            ) + 1
            pad_audio = pad_or_trim(audio_segment.flatten(), length=self.chunk_samples)
            mel = log_mel_spectrogram(pad_audio)
            assert mel.shape[1] == self.chunk_mel_frames
            mel_features.append(mel)
            token_lens.append(token_len)
            audio_sample_lens.append(segment_length)

        mel_features = torch.stack(mel_features, dim=0)
        # Keep input_seq_lens on CPU: the conv2d bucket path reads it via
        # ``.item()`` and CPU scalars avoid device->host syncs. The encoder's
        # varlen path moves it to the device itself.
        input_seq_lens = torch.tensor(token_lens, dtype=torch.long) * self.merge_factor
        audio_embedding = self._forward_speech_encoder(
            mel_features, input_seq_lens, audio_sample_lens
        )

        chunk_embeddings = []
        for idx, token_len in enumerate(token_lens):
            chunk_embeddings.append(
                audio_embedding[idx, : token_len * self.merge_factor, :]
            )
        return torch.cat(chunk_embeddings, dim=0).unsqueeze(0)