31 lines
1004 B
Python
31 lines
1004 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from typing import NamedTuple
|
|
|
|
import torch
|
|
from vllm.sequence import IntermediateTensors
|
|
|
|
from vllm_omni.data_entry_keys import OmniPayload
|
|
|
|
|
|
class OwnedBatchTensor(NamedTuple):
|
|
"""A batch tensor whose producer transfers ownership to the buffer.
|
|
|
|
The producer must have freshly allocated the tensor (e.g. an index_select
|
|
result) and must not write it afterwards. The buffer keeps row views
|
|
instead of taking a second snapshot. Tensors without this marker are
|
|
treated as borrowed (graph output / shared scratch) and still get a clone.
|
|
"""
|
|
|
|
tensor: torch.Tensor
|
|
|
|
|
|
class OmniOutput(NamedTuple):
|
|
"""Output from the merged Omni model containing both text and audio."""
|
|
|
|
text_hidden_states: torch.Tensor
|
|
multimodal_outputs: OmniPayload | None = None
|
|
intermediate_tensors: IntermediateTensors | None = None
|
|
next_token_id: torch.Tensor | None = None
|