41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
from dataclasses import dataclass
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from vllm.outputs import RequestOutput
|
|
from vllm.v1.outputs import ModelRunnerOutput
|
|
|
|
|
|
class OmniModelRunnerOutput(ModelRunnerOutput):
|
|
"""Model runner output for omni models.
|
|
|
|
Extends the base ModelRunnerOutput with support for multimodal outputs
|
|
that may be produced by non-autoregressive stages.
|
|
|
|
Attributes:
|
|
multimodal_outputs: Optional dictionary mapping modality names to
|
|
output tensors (e.g., {"image": tensor, "audio": tensor})
|
|
"""
|
|
|
|
multimodal_outputs: Optional[dict[str, torch.Tensor]] = None
|
|
|
|
|
|
@dataclass
|
|
class OmniRequestOutput(RequestOutput):
|
|
"""Request output for omni pipeline stages.
|
|
|
|
Wraps a standard RequestOutput with stage-specific metadata,
|
|
indicating which stage produced the output and what type of
|
|
final output it represents.
|
|
|
|
Attributes:
|
|
stage_id: Identifier of the stage that produced this output
|
|
final_output_type: Type of final output (e.g., "text", "image",
|
|
"audio", "latents")
|
|
request_output: The underlying RequestOutput from the stage
|
|
"""
|
|
|
|
stage_id: int
|
|
final_output_type: str
|
|
request_output: RequestOutput
|