diff --git a/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py b/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py
index dfeb23b67..022cecc17 100644
--- a/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py
+++ b/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py
@@ -1255,6 +1255,13 @@ class Qwen2_5OmniToken2WavDiTModel(nn.Module):
         return generated_mel_spectrogram
 
 
+@auto_docstring(
+    custom_intro="""
+    The full Qwen2.5Omni Token2Wav model. Consists a DiT model take speech
+    tokens as input and predict mel spectrogram and a BigVGAN vocoder take
+    mel spectrogram as input and predict waveform.
+    """
+)
 class Qwen2_5OmniToken2WavModel(nn.Module):
     def __init__(self, *, vllm_config: VllmConfig, prefix=""):
         super().__init__()
diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml
index 54d2b2359..0cb4d34bd 100644
--- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml
+++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml
@@ -60,6 +60,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
+      ode_solver_class: RungeKutta4ODESolver
       worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
       scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
       gpu_memory_utilization: 0.15