48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
from typing import Optional
|
|
|
|
from pydantic import ConfigDict
|
|
from pydantic.dataclasses import dataclass
|
|
from vllm.config import ModelConfig, config
|
|
|
|
import vllm_omni.model_executor.models as me_models
|
|
|
|
|
|
@config
|
|
@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
|
|
class OmniModelConfig(ModelConfig):
|
|
"""Configuration for Omni models, extending the base ModelConfig.
|
|
|
|
This configuration class extends the base vLLM ModelConfig with
|
|
omni-specific fields for multi-stage pipeline processing.
|
|
|
|
Attributes:
|
|
stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
|
|
model_stage: Stage type identifier, e.g., "thinker" or "talker"
|
|
(default: "thinker")
|
|
model_arch: Model architecture name
|
|
(default: "Qwen2_5OmniForConditionalGeneration")
|
|
engine_output_type: Optional output type specification for the engine.
|
|
Used to route outputs to appropriate processors (e.g., "image",
|
|
"audio", "latents"). If None, output type is inferred.
|
|
|
|
Example:
|
|
>>> config = OmniModelConfig(
|
|
... stage_id=0,
|
|
... model_stage="thinker",
|
|
... model_arch="Qwen2_5OmniForConditionalGeneration"
|
|
... )
|
|
"""
|
|
|
|
stage_id: int = 0
|
|
model_stage: str = "thinker"
|
|
model_arch: str = "Qwen2_5OmniForConditionalGeneration"
|
|
engine_output_type: Optional[str] = None
|
|
|
|
@property
|
|
def registry(self):
|
|
return me_models.OmniModelRegistry
|
|
|
|
@property
|
|
def architectures(self) -> list[str]:
|
|
return [self.model_arch]
|