expose ode_solver_class

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>

vllm_omni/config/model.py

@@ -15,6 +15,7 @@ class OmniModelConfig(ModelConfig):
     stage_id: int = 0
     model_stage: str = "thinker"
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
+    ode_solver_class: Optional[str] = None
     engine_output_type: Optional[str] = None
 
     @property

vllm_omni/engine/arg_utils.py

@@ -12,6 +12,7 @@ class OmniEngineArgs(EngineArgs):
     stage_id: int = 0
     model_stage: str = "thinker"
     model_arch: str = "Qwen2_5OmniForConditionalGeneration"
+    ode_solver_class: Optional[str] = None
     engine_output_type: Optional[str] = None
 
     def create_model_config(self) -> OmniModelConfig:
@@ -27,6 +28,7 @@ class OmniEngineArgs(EngineArgs):
         config_dict["model_stage"] = self.model_stage
         config_dict["model_arch"] = self.model_arch
         config_dict["engine_output_type"] = self.engine_output_type
+        config_dict["ode_solver_class"] = self.ode_solver_class
 
         # Create and return the OmniModelConfig instance
         omni_config = OmniModelConfig(**config_dict)

vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py

@@ -34,7 +34,7 @@ from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
-from vllm_omni.model_executor.ode_solver import RungeKutta4ODESolver
+from vllm_omni.model_executor.ode_solver import get_ode_solver_class
 
 
 # Provide a no-op auto_docstring decorator to satisfy annotations if missing
@@ -1058,7 +1058,7 @@ class Qwen2_5OmniToken2WavDiTModel(nn.Module):
         self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(config.hidden_size)  # final modulation
         self.proj_out = nn.Linear(config.hidden_size, config.mel_dim)
 
-        self.ode_solver = RungeKutta4ODESolver(function=None, initial_value=None)  # placeholder
+        self.ode_solver = get_ode_solver_class(vllm_config.model_config.ode_solver_class)(function=None, initial_value=None)  # placeholder
 
     def _create_block_diff(self, hidden_states):
         batch, seq_len = hidden_states.shape[0], hidden_states.shape[1]