From 1c8a31821f65d4f7a358c9fff32822ee207d7864 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Tue, 18 Nov 2025 23:56:22 +0800 Subject: [PATCH] expose ode_solver_class Signed-off-by: Isotr0py --- vllm_omni/config/model.py | 1 + vllm_omni/engine/arg_utils.py | 2 ++ vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py | 4 ++-- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 67bcc5fad..7bd77dc0b 100644 --- a/vllm_omni/config/model.py +++ b/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 diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 64df79c7b..6f1c1d277 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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) 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 72b35d57b..774d3704a 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py +++ b/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]