[Voxtral TTS] Correct decode steps param in Voxtral TTS (#2524)

Signed-off-by: Chen-Yo Sun <chenyo.sun@mistral.ai>

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

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,

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)
 

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,
         )