diff --git a/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py b/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py index 9d5c266b4..3137a878e 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py @@ -10,18 +10,10 @@ import torch import torch.nn as nn import torch.nn.functional as F from transformers import PretrainedConfig, PreTrainedModel, Qwen2Config, Qwen2Model +from transformers.utils import is_flash_attn_2_available from vllm.logger import init_logger logger = init_logger(__name__) -try: - import flash_attn # noqa: F401 -except (ImportError, ModuleNotFoundError): - flash_attn = None - logger.warning( - "flash_attn is not available, the model may not yield the " - "exactly same result as the transformers implementation " - "in the audio tower part." - ) class AudioVAEConfig(PretrainedConfig): @@ -200,8 +192,13 @@ class Decoder(nn.Module): def __init__(self, decoder_args, output_dim=320, latent_dim=64, patch_size=-1): super().__init__() config = Qwen2Config.from_dict(config_dict=decoder_args) - if flash_attn is None: + if is_flash_attn_2_available(): + config._attn_implementation_autoset = True + config._attn_implementation = "flash_attention_2" + else: config._attn_implementation = "sdpa" + + logger.info("AudioVAE Decoder: using attn_implementation=%r", config._attn_implementation) self.decoder = Qwen2Model(config) self.output_dim = output_dim self.latent_dim = latent_dim @@ -276,8 +273,13 @@ class Encoder(nn.Module): def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, patch_size=-1): super().__init__() config = Qwen2Config.from_dict(config_dict=encoder_args) - if flash_attn is None: + if is_flash_attn_2_available(): + config._attn_implementation_autoset = True + config._attn_implementation = "flash_attention_2" + else: config._attn_implementation = "sdpa" + + logger.info("AudioVAE Encoder: using attn_implementation=%r", config._attn_implementation) self.encoder = Qwen2Model(config) self.input_dim = input_dim self.hop_size = hop_size diff --git a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py index ae34eb865..ee2b6b2a1 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py @@ -885,6 +885,10 @@ class MingAudioGenerator: self._use_cuda_graphs = use_cuda_graphs + # For FA2, let it see a full-length seq Q + # trailing latent frames prepended on each decode call + self._vae_decode_pad_frames = 32 + @cached_property def _sampler_pool(self) -> CFMGraphExecutorPool | None: device = next(self._model.parameters()).device @@ -1087,21 +1091,33 @@ class MingAudioGenerator: # VAE streaming decode def _stream_decode(self, latents: list[torch.Tensor]) -> torch.Tensor: sr = int(self._audio_vae.config.sample_rate) - vae_cache = {"past_key_values": None, "stream_state": (None, None, None)} + decode_pad: torch.Tensor | None = None sil_cache: dict | None = None wav_chunks: list[torch.Tensor] = [] for i, lat in enumerate(latents): last_chunk = i == (len(latents) - 1) - speech, stream_state, past_key_values = self._audio_vae.decode( - lat, - past_key_values=vae_cache["past_key_values"], - use_cache=True, - stream_state=vae_cache["stream_state"], - last_chunk=last_chunk, + + if decode_pad is not None: + vae_input = torch.cat([decode_pad, lat], dim=1) + pad_frames = decode_pad.shape[1] + else: + vae_input = lat + pad_frames = 0 + + # Stateless, no KV cache accum intentionally. + speech, _, _ = self._audio_vae.decode( + vae_input, + use_cache=False, + stream_state=(None, None, None), + last_chunk=True, ) - vae_cache = {"past_key_values": past_key_values, "stream_state": stream_state} - speech_chunk = speech[0].detach().float() + + total_frames = vae_input.shape[1] + dcs = speech.shape[-1] // total_frames + + # keep only the new audio. + speech_chunk = speech[:, :, pad_frames * dcs :][0].detach().float() speech_chunk, sil_cache = silence_holder( speech_chunk, sr, @@ -1111,6 +1127,9 @@ class MingAudioGenerator: if speech_chunk.numel() > 0: wav_chunks.append(speech_chunk) + # Advance the sliding buffer + decode_pad = vae_input[:, -self._vae_decode_pad_frames :, :].detach() + if not wav_chunks: device = next(self._model.parameters()).device dtype = next(self._model.parameters()).dtype