diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py index 3b3f0d506..c34d21c49 100644 --- a/vllm_omni/config/config_factory.py +++ b/vllm_omni/config/config_factory.py @@ -9,7 +9,7 @@ from dataclasses import asdict from pathlib import Path from typing import Any -from transformers import PreTrainedConfig +from transformers import PretrainedConfig from vllm.logger import init_logger from vllm.transformers_utils.config import get_config from vllm.transformers_utils.repo_utils import get_hf_file_to_dict @@ -329,7 +329,7 @@ class StageConfigFactory: return build_stage_runtime_overrides(stage.stage_id, cli_overrides) @staticmethod - def resolve_pipeline_config(model_type: str, hf_config: PreTrainedConfig | None = None) -> PipelineConfig | None: + def resolve_pipeline_config(model_type: str, hf_config: PretrainedConfig | None = None) -> PipelineConfig | None: """Given a model type, resolve to the pipeline to be used. If the pipeline maps to a callable we resolve based on the HF config.""" if model_type not in OMNI_PIPELINES: diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index dc5843a26..a3f7f6e16 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -27,7 +27,7 @@ from __future__ import annotations from collections.abc import Callable from typing import TypeAlias -from transformers import PreTrainedConfig +from transformers import PretrainedConfig from vllm.logger import init_logger from vllm_omni.config.stage_config import ( @@ -78,7 +78,7 @@ from vllm_omni.model_executor.models.voxtral_tts.pipeline import VOXTRAL_TTS_PIP logger = init_logger(__name__) -PipelineResolverFunc: TypeAlias = Callable[[PreTrainedConfig | None], PipelineConfig | None] +PipelineResolverFunc: TypeAlias = Callable[[PretrainedConfig | None], PipelineConfig | None] # --- Multi-stage omni pipelines (LLM-centric; audio / video I/O) --- OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {