diff --git a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py index 6f072944d..847adae06 100644 --- a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py +++ b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py @@ -78,6 +78,13 @@ except Exception: AudioSpecialTokens = _mod2.AudioSpecialTokens +class SyntheticAcousticTransformerArgs: + """Mimics AcousticTransformerArgs interface.""" + + def __init__(self): + self.n_decoding_steps = 7 + + class SyntheticModelArgs: """Mimics MultimodalAudioModelArgs interface.""" @@ -96,6 +103,7 @@ class SyntheticAcousticTransformer(nn.Module): def __init__(self): super().__init__() self.model_args = SyntheticModelArgs() + self.acoustic_transformer_args = SyntheticAcousticTransformerArgs() self.acoustic_embeddings_levels = ACOUSTIC_EMBEDDINGS_LEVELS # semantic_codebook_output: hidden_dim -> padded_codebook_size diff --git a/vllm_omni/model_executor/models/voxtral_tts/configuration_voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/configuration_voxtral_tts.py index d32a882e7..0f22c764a 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/configuration_voxtral_tts.py +++ b/vllm_omni/model_executor/models/voxtral_tts/configuration_voxtral_tts.py @@ -48,6 +48,15 @@ class VoxtralTTSConfigParser(MistralConfigParser): audio_tokenizer_args = config_dict["multimodal"].pop("audio_tokenizer_args", None) audio_config = {} if encoder_args is not None: + # Default n_decoding_steps if not provided + acoustic_args = encoder_args.get("acoustic_transformer_args", {}) + if acoustic_args.get("n_decoding_steps") is None: + logger.warning( + "n_decoding_steps not provided in acoustic_transformer_args, defaulting to 7. " + "Please add 'n_decoding_steps' to params.json under acoustic_transformer_args." + ) + acoustic_args["n_decoding_steps"] = 7 + audio_config = { "sampling_rate": encoder_args["audio_encoding_args"]["sampling_rate"], "codec_args": audio_tokenizer_args, diff --git a/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py b/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py index a4d58df5b..ff053342d 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py +++ b/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py @@ -49,7 +49,7 @@ class CUDAGraphAcousticTransformerWrapper: self.acoustic_embeddings_levels = self.acoustic_transformer.acoustic_embeddings_levels self.cfg_alpha = 1.2 - self.n_steps = 8 + self.n_steps = self.acoustic_transformer.acoustic_transformer_args.n_decoding_steps # Graph storage self.graphs: dict[int, CUDAGraph] = {} @@ -73,7 +73,7 @@ class CUDAGraphAcousticTransformerWrapper: ) # Pre-create persistent buffers - self.timesteps = torch.linspace(0, 1, self.n_steps, device=device, dtype=dtype) + self.timesteps = torch.linspace(0, 1, self.n_steps + 1, device=device, dtype=dtype) self.fake_eos_one = torch.tensor(1.0, dtype=dtype, device=device) self.fake_eos_zero = torch.tensor(0.0, dtype=dtype, device=device) diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py index b5d116173..4041a53e5 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py @@ -108,6 +108,7 @@ class AcousticTransformerArgs: use_biases: bool = False norm_eps: float = 1e-5 sigma: float = 1e-5 # was 0.01 in beta version + n_decoding_steps: int | None = None # Number of Euler ODE steps for flow matching @dataclass @@ -436,14 +437,13 @@ class FlowMatchingAudioTransformer(nn.Module): self._empty_audio_token_id = AudioSpecialTokens.id(AudioSpecialTokens.empty_audio) # Flow matching constants - # TODO(chenyo): hardcoded, need to fix - self._acoustic_decode_iters = 8 + self._n_steps = args.n_decoding_steps # TODO(chenyo): hardcoded, need to fix self._cfg_alpha = 1.2 self._noise_scale = 1.0 self.register_buffer( "_timesteps", - torch.linspace(0, 1, self._acoustic_decode_iters), + torch.linspace(0, 1, self._n_steps + 1), persistent=False, )