[EPD] Support metadata-only audio inputs (#57887)

Signed-off-by: Tianyu Guo <guoty@inferact.ai>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
This commit is contained in:
Tianyu Guo
2026-09-23 09:03:50 +08:00
committed by GitHub
co-authored by Cyrus Leung
parent 826e300cde
commit 3ee2906f47
14 changed files with 387 additions and 27 deletions
+13
View File
@@ -61,6 +61,19 @@ and compares the disaggregated outputs against it. Set `MODEL` to use another mo
or `MOONCAKE_EC_PROTOCOL=rdma` to use RDMA on supported hardware. See the script for
additional configuration options.
### Audio inputs
The Python EPD proxy also rewrites audio in `/v1/chat/completions` into
metadata-only references for Qwen2-Audio, AudioFlamingo3, Ultravox,
Qwen2.5-Omni, and Qwen3-Omni. The encoder publishes `audio_num_tokens`
(the placeholder token count of each audio) together with any feature
lengths the model needs; the consumer uses them to reconstruct the audio
placeholders and receives the embeddings through its EC connector without
repeating audio preprocessing.
This covers pure audio inputs, not video with an embedded audio track. It does
not add `/v1/audio/transcriptions` or realtime routes to the example proxy.
---
## 3 Test Script
@@ -480,6 +480,17 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
the grid its processor actually produced), so the grid is never re-derived
here -- a second derivation could disagree with the encoder's.
"""
mm_items = extract_mm_items(req_data)
# Audio processors may consume placeholders from other modalities too.
# Keep the whole request raw if any item cannot be rewritten safely.
control_fields = {"mm_hash", "ec_mm_hash", "transfer_id"}
raw_audio_request = any(
item["type"] in {"audio_url", "input_audio"} for item in mm_items
) and any(
not item_meta.get(i, {}).get("mm_hash")
or not (item_meta.get(i, {}).keys() - control_fields)
for i in range(len(mm_items))
)
rewritten = 0
transfer_items = []
idx = 0
@@ -499,6 +510,18 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
item_uuid = meta.pop("mm_hash", None)
ec_mm_hash = meta.pop("ec_mm_hash", None) or item_uuid
transfer_id = meta.pop("transfer_id", None)
if raw_audio_request:
new_content.append(
{
**item,
"uuid": item_uuid or item.get("uuid") or content_uuid(item),
}
)
if transfer_id is not None:
transfer_items.append(
{"mm_hash": ec_mm_hash, "transfer_id": transfer_id}
)
continue
# Whatever keys the encoder reported are the metadata its model
# declared as needed to size the placeholder range; the proxy does
# not need to know their names.
@@ -532,9 +555,10 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
rewritten += 1
new_messages.append({**msg, "content": new_content})
if not rewritten:
if not rewritten and not raw_audio_request:
return req_data
logger.info("Rewrote %d media item(s) as metadata references", rewritten)
if rewritten:
logger.info("Rewrote %d media item(s) as metadata references", rewritten)
rewritten_request = {**req_data, "messages": new_messages}
if transfer_items:
ec_transfer_params = dict(req_data.get("ec_transfer_params") or {})
@@ -632,7 +656,7 @@ async def fanout_encoder_primer(
content = []
for idx in indices:
item = mm_items[idx]
item_uuid = None if NO_REWRITE else content_uuid(item)
item_uuid = None if NO_REWRITE else (item.get("uuid") or content_uuid(item))
if item_uuid is not None:
item_uuids[idx] = item_uuid
item_transfer_ids[idx] = uuid.uuid4().hex
@@ -756,14 +780,13 @@ async def fanout_encoder_primer(
raise HTTPException(502, "Encoder metadata cannot be matched")
ec_mm_hash, reported = matched
metadata = reported.get("metadata") or {}
if metadata:
item_meta[idx] = {
**metadata,
"mm_hash": item_uuids.get(idx, ec_mm_hash),
"ec_mm_hash": ec_mm_hash,
}
if idx in item_transfer_ids:
item_meta[idx]["transfer_id"] = item_transfer_ids[idx]
item_meta[idx] = {
**metadata,
"mm_hash": item_uuids.get(idx, ec_mm_hash),
"ec_mm_hash": ec_mm_hash,
}
if idx in item_transfer_ids:
item_meta[idx]["transfer_id"] = item_transfer_ids[idx]
# Whatever the encoder reported alongside `metadata` is the
# connector's own handle on the published embedding (for NIXL,
# peer_host/peer_port/size_bytes). The decoder's connector
+5
View File
@@ -144,6 +144,11 @@ def test_qwen3_omni_get_updates_use_audio_in_video(
# Create a mock context
mock_ctx = Mock(spec=InputProcessingContext)
mock_ctx.tokenizer = mock_tokenizer
# `model_config` is an instance attribute, so it is not covered by the
# spec; the data parser reads it for `allow_missing_mm_embeddings`.
mock_ctx.model_config = SimpleNamespace(
multimodal_config=SimpleNamespace(allow_missing_mm_embeddings=False)
)
# Create processing info
info = Qwen3OmniMoeThinkerProcessingInfo(mock_ctx)
@@ -1,13 +1,16 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
from collections.abc import Set as AbstractSet
from functools import partial
import numpy as np
import pytest
import torch
from PIL import Image
from tests.models.utils import build_model_context
from vllm.config import ModelConfig
from vllm.config.multimodal import (
AudioDummyOptions,
@@ -16,10 +19,15 @@ from vllm.config.multimodal import (
MultiModalDummyOptions,
VideoDummyOptions,
)
from vllm.distributed.ec_transfer.ec_connector.utils import (
PlaceholderMetadataResolver,
collect_ec_item_metadata,
)
from vllm.entrypoints.chat_utils import _get_embeds_data
from vllm.inputs import MultiModalDataDict, MultiModalInput
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.cache import MultiModalProcessorOnlyCache
from vllm.multimodal.inputs import batched_tensors_equal
from vllm.multimodal.inputs import MultiModalFeatureSpec, batched_tensors_equal
from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext
from vllm.platforms import current_platform
from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config
@@ -193,6 +201,79 @@ def get_token_prompt(
return inputs.prompt
@pytest.mark.parametrize(
"model_id,durations",
[
("Qwen/Qwen2-Audio-7B-Instruct", [1.0, 2.3]),
("nvidia/audio-flamingo-3-hf", [31.0]),
("fixie-ai/ultravox-v0_5-llama-3_2-1b", [1.0, 31.0]),
("Qwen/Qwen2.5-Omni-3B", [1.0, 2.3]),
("Qwen/Qwen3-Omni-30B-A3B-Instruct", [1.0, 2.3]),
],
)
def test_audio_metadata_only_roundtrip(model_id, durations, monkeypatch):
"""EC metadata preserves the raw-audio prompt without running the HF processor."""
ctx = build_model_context(model_id, limit_mm_per_prompt={"audio": len(durations)})
mm_config = ctx.model_config.multimodal_config
mm_config.enable_mm_embeds = True
mm_config.allow_missing_mm_embeddings = True
processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
sample_rate = processor.info.get_feature_extractor().sampling_rate
audios = [
np.zeros(int(duration * sample_rate), dtype=np.float32)
for duration in durations
]
prompt = processor.dummy_inputs.get_dummy_text({"audio": len(audios)})
raw = processor(
prompt,
mm_items=processor.info.parse_mm_data({"audio": audios}),
hf_processor_mm_kwargs={},
)
features = [
MultiModalFeatureSpec(
data=item,
modality="audio",
identifier=str(index),
mm_position=position,
)
for index, (item, position) in enumerate(
zip(raw["mm_kwargs"]["audio"], raw["mm_placeholders"]["audio"])
)
]
published = collect_ec_item_metadata(
features, PlaceholderMetadataResolver(ctx.model_config)
)
metadata = json.loads(
json.dumps([entry["metadata"] for entry in published.values()])
)
assert len(metadata) == len(audios)
assert all(metadata)
audio_data = _get_embeds_data("audio", metadata, processor)
def unexpected_preprocessing(*args, **kwargs):
pytest.fail("The EC consumer must not preprocess audio")
monkeypatch.setattr(processor, "_call_hf_processor", unexpected_preprocessing)
remote = processor(
prompt,
mm_items=processor.info.parse_mm_data({"audio": audio_data}),
hf_processor_mm_kwargs={},
)
assert remote["prompt_token_ids"] == raw["prompt_token_ids"]
assert remote["mm_placeholders"] == raw["mm_placeholders"]
# Omni also consumes these lengths when constructing M-RoPE positions.
if "audio_feature_lengths" in audio_data:
assert torch.equal(
remote["mm_kwargs"].get_data()["audio_feature_lengths"],
raw["mm_kwargs"].get_data()["audio_feature_lengths"],
)
mm_config.allow_missing_mm_embeddings = False
ordinary_processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
with pytest.raises(ValueError, match="should contain the fields"):
ordinary_processor.info.parse_mm_data({"audio": audio_data})
def random_vision_chunk(
rng: np.random.RandomState,
min_wh: int,
+39
View File
@@ -15,6 +15,45 @@ from vllm.multimodal.parse import (
H, W = 480, 640
class AudioMetadataParser(MultiModalDataParser):
embedding_fields = {
"audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"},
}
@pytest.mark.parametrize("allow_missing", [False, True])
def test_audio_metadata_requires_ec_consumer(allow_missing):
parser = AudioMetadataParser(allow_missing_mm_embeddings=allow_missing)
data = {"audio_num_tokens": torch.tensor([[3], [5]])}
if not allow_missing:
with pytest.raises(ValueError, match="audio_embeds"):
parser.parse_mm_data({"audio": data})
else:
items = parser.parse_mm_data({"audio": data})["audio"]
assert len(items) == 2
assert items.get(1)["audio_num_tokens"].item() == 5
assert items.get_processor_data() == {}
@pytest.mark.parametrize("counts", [[0], [-1], [1.5], [True], [[1, 2]]])
def test_audio_metadata_rejects_invalid_token_counts(counts):
parser = AudioMetadataParser(allow_missing_mm_embeddings=True)
with pytest.raises(ValueError, match="positive integer"):
parser.parse_mm_data({"audio": {"audio_num_tokens": torch.tensor(counts)}})
def test_audio_metadata_checks_supplied_embedding_lengths():
parser = AudioMetadataParser()
data = {
"audio_num_tokens": torch.tensor([3, 5]),
"audio_embeds": [torch.zeros(3, 8), torch.zeros(5, 8)],
}
assert len(parser.parse_mm_data({"audio": data})["audio"]) == 2
data["audio_num_tokens"] = torch.tensor([3, 4])
with pytest.raises(ValueError, match="does not match"):
parser.parse_mm_data({"audio": data})
@pytest.mark.parametrize(
"image",
[
@@ -56,6 +56,22 @@ def test_metadata_preserves_repeated_item_positions():
}
def test_audio_metadata_survives_processor_cache_hit():
resolver = Mock()
resolver.fields_for.return_value = {"audio_num_tokens"}
feature = MultiModalFeatureSpec(
data=None,
modality="audio",
identifier="audio",
mm_position=PlaceholderRange(
offset=0, length=4, is_embed=torch.tensor([False, True, True, False])
),
)
assert collect_ec_item_metadata([feature], resolver) == {
"audio": {"metadata": {"audio_num_tokens": [2]}, "item_indices": [0]}
}
@pytest.fixture
def temp_storage(tmp_path):
"""Fixture providing temporary storage path."""
@@ -87,6 +87,82 @@ class _EncoderSession:
return _EncoderResponse(self._replies.pop(0))
@pytest.mark.parametrize("second_type", ["image_url", "audio_url", "input_audio"])
@pytest.mark.parametrize("transfer", ["example", "nixl", "mooncake"])
def test_partial_audio_metadata_falls_back_with_all_transfers(
proxy, monkeypatch, second_type, transfer
):
"""Never mix metadata-only audio with raw media, or lose EC identities."""
import copy
seen = []
reported = {}
class Session:
async def post(self, url, data=None, headers=None):
body = msgspec.json.decode(data)
seen.append(body)
item = body["messages"][0]["content"][0]
mm_hash = item["uuid"] + "-processed"
entry: dict[str, Any] = {
"metadata": {"audio_feature_lengths": [100]} if len(seen) == 1 else {},
"item_indices": [0],
}
if transfer == "nixl":
entry.update(peer_host="encoder", peer_port=1234, size_bytes=4096)
reported[mm_hash] = entry
return _EncoderResponse({mm_hash: entry})
second = (
{"type": second_type, "input_audio": {"data": "large", "format": "wav"}}
if second_type == "input_audio"
else {"type": second_type, second_type: {"url": "large"}}
)
body = {
"messages": [
{
"role": "user",
"content": [
{
"type": "audio_url",
"audio_url": {"url": "small"},
"uuid": "user-audio",
},
second,
],
}
],
"mm_processor_kwargs": {"sampling_rate": 16000},
}
original = copy.deepcopy(body)
monkeypatch.setattr(proxy, "encode_session", Session())
monkeypatch.setattr(proxy, "NO_REWRITE", False)
consumer = "tcp://consumer:1" if transfer == "mooncake" else None
prepared, _, _ = asyncio.run(
proxy.prepare_for_decode(body, "r", ["http://encoder"], "", consumer)
)
assert body == original
assert prepared["mm_processor_kwargs"] == body["mm_processor_kwargs"]
content = prepared["messages"][0]["content"]
assert [item["type"] for item in content] == ["audio_url", second_type]
assert content[0]["uuid"] == "user-audio"
for item, encoded in zip(content, seen):
assert item == encoded["messages"][0]["content"][0]
params = prepared["ec_transfer_params"]
for mm_hash, entry in reported.items():
assert params[mm_hash] == entry
if transfer == "mooncake":
assert params["ec_items"] == [
{
"mm_hash": mm_hash,
"transfer_id": encoded["ec_transfer_params"]["ec_items"][0][
"transfer_id"
],
}
for mm_hash, encoded in zip(reported, seen)
]
@pytest.mark.parametrize("no_rewrite", [False, True])
@pytest.mark.parametrize(
"batch_size, expected_sizes", [(0, [3, 3]), (1, [1] * 6), (2, [2, 2, 1, 1])]
@@ -509,7 +585,11 @@ def test_video_audio_fallback(proxy, monkeypatch, no_rewrite, transfer):
original = copy.deepcopy(body)
# Collector indices describe two processed features from one video item.
video_metadata = collect_ec_item_metadata(
[SimpleNamespace(identifier=key, data=None) for key in ("video", "audio")], None
[
SimpleNamespace(identifier=key, modality=key, data=None)
for key in ("video", "audio")
],
SimpleNamespace(fields_for=lambda _: set()),
)
if transfer == "handle":
video_metadata["audio"]["transfer_id"] = "reservation"
@@ -60,24 +60,26 @@ def collect_ec_item_metadata(
Keyed by mm_hash, each entry carries a `metadata` dict with whatever
placeholder fields `resolver` says this model needs published for its
modality, so a consumer can skip the image transform. `item_indices`
modality, so a consumer can skip media preprocessing. `item_indices`
identifies every occurrence in `mm_features`, including repeated hashes.
`data` is None for items served from the processor cache, so the metadata is
unavailable here and the consumer has to fall back to processing the
media itself. A connector that also has transfer coordinates to report
Audio token counts come from the placeholder; other metadata requires
`data` and falls back to raw media when it is unavailable on cache hits.
A connector that also has transfer coordinates to report
(e.g. NIXL peer_host/peer_port/size_bytes) merges those in alongside
`metadata`, not into it.
"""
items: dict[str, dict[str, Any]] = {}
for index, feature in enumerate(mm_features):
metadata: dict[str, Any] = {}
wanted = resolver.fields_for(feature.modality)
if "audio_num_tokens" in wanted:
metadata["audio_num_tokens"] = [feature.mm_position.get_num_embeds()]
if feature.data is not None:
wanted = resolver.fields_for(feature.modality)
for key, value in feature.data.get_data().items():
if key not in wanted:
continue
if isinstance(value, torch.Tensor):
metadata[key] = value.tolist()
metadata[key] = torch.atleast_1d(value).tolist()
elif is_list_of(value, (int, float)):
# Some metadata (e.g. Qwen3-VL video timestamps) is
# produced as a plain list rather than a tensor.
+13 -2
View File
@@ -197,6 +197,7 @@ class AudioFlamingo3ProcessingInfo(BaseProcessingInfo):
return AudioFlamingo3MultiModalDataParser(
target_sr=feature_extractor.sampling_rate,
expected_hidden_size=self._get_expected_hidden_size(),
allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
)
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
@@ -237,6 +238,7 @@ def _audioflamingo3_field_config(hf_inputs: Mapping[str, torch.Tensor]):
if chunk_counts is not None:
return dict(
audio_embeds=MultiModalFieldConfig.batched("audio"),
audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
input_features=MultiModalFieldConfig.flat_from_sizes(
"audio", chunk_counts, dim=0
),
@@ -247,6 +249,7 @@ def _audioflamingo3_field_config(hf_inputs: Mapping[str, torch.Tensor]):
)
return dict(
audio_embeds=MultiModalFieldConfig.batched("audio"),
audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
input_features=MultiModalFieldConfig.batched("audio"),
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
@@ -354,11 +357,15 @@ def _count_audio_tokens_from_mask(
class AudioFlamingo3MultiModalDataParser(MultiModalDataParser):
embedding_fields = {
"audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"},
}
def _parse_audio_data(
self,
data: dict[str, torch.Tensor] | ModalityData[Any],
) -> ModalityDataItems[Any, Any] | None:
if isinstance(data, dict):
if isinstance(data, dict) and "audio_num_tokens" not in data:
return DictEmbeddingItems(
data,
modality="audio",
@@ -442,7 +449,11 @@ class AudioFlamingo3MultiModalProcessor(
chunk_counts = out_mm_data.get("chunk_counts")
def get_replacement_audioflamingo3(item_idx: int):
if feature_attention_mask is not None:
if "audio_num_tokens" in out_mm_data:
counts = out_mm_data["audio_num_tokens"]
assert isinstance(counts, torch.Tensor)
num_features = int(counts[item_idx])
elif feature_attention_mask is not None:
num_features = _count_audio_tokens_from_mask(
feature_attention_mask,
chunk_counts,
@@ -282,7 +282,7 @@ def create_qwen2_5_omni_thinker_field_factory(
def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
audio_feature_lengths = hf_inputs.get(
"audio_feature_lengths", torch.empty((0,))
)
).flatten()
image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))
image_pixel_grid_sizes = image_grid_thw.prod(-1)
@@ -332,6 +332,14 @@ def create_qwen2_5_omni_thinker_field_factory(
class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
embedding_fields = {
**Qwen2VLMultiModalDataParser.embedding_fields,
"audio": {
"input_audio_features": "values",
"audio_feature_lengths": "metadata",
},
}
def __init__(self, spatial_merge_size: int, *args, **kwargs):
self._spatial_merge_size = spatial_merge_size
super().__init__(self._spatial_merge_size, *args, **kwargs)
@@ -341,10 +349,15 @@ class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
data: dict[str, torch.Tensor] | ModalityData[ImageItem],
) -> ModalityDataItems[Any, Any]:
if isinstance(data, dict):
required, optional = self.embedding_field_sets("audio")
data = dict(data)
if "audio_feature_lengths" in data:
data["audio_feature_lengths"] = data["audio_feature_lengths"].flatten()
return DictEmbeddingItems(
data,
modality="audio",
required_fields={"input_audio_features", "audio_feature_lengths"},
required_fields=required,
optional_fields=optional,
fields_factory=create_qwen2_5_omni_thinker_field_factory(
self._spatial_merge_size
),
+10 -2
View File
@@ -128,17 +128,22 @@ def _get_feat_extract_output_lengths(input_lengths: torch.Tensor):
def _qwen2audio_field_config(hf_inputs: Mapping[str, torch.Tensor]):
return dict(
audio_embeds=MultiModalFieldConfig.batched("audio"),
audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
input_features=MultiModalFieldConfig.batched("audio"),
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
)
class Qwen2AudioMultiModalDataParser(MultiModalDataParser):
embedding_fields = {
"audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"},
}
def _parse_audio_data(
self,
data: dict[str, torch.Tensor] | ModalityData[AudioItem],
) -> ModalityDataItems[Any, Any] | None:
if isinstance(data, dict):
if isinstance(data, dict) and "audio_num_tokens" not in data:
return DictEmbeddingItems(
data,
modality="audio",
@@ -169,6 +174,7 @@ class Qwen2AudioProcessingInfo(BaseProcessingInfo):
target_sr=feature_extractor.sampling_rate,
target_channels=self.get_target_channels(),
expected_hidden_size=self._get_expected_hidden_size(),
allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
)
def get_target_channels(self) -> int:
@@ -279,7 +285,9 @@ class Qwen2AudioMultiModalProcessor(BaseMultiModalProcessor[Qwen2AudioProcessing
audio_output_lengths = audio_output_lens.tolist()
def get_replacement_qwen2_audio(item_idx: int):
if audio_output_lengths:
if "audio_num_tokens" in out_mm_data:
num_features = int(out_mm_data["audio_num_tokens"][item_idx])
elif audio_output_lengths:
num_features = audio_output_lengths[item_idx]
else:
audio_embeds = out_mm_data["audio_embeds"][item_idx]
@@ -100,6 +100,7 @@ from .qwen2_5_omni_thinker import (
Qwen2_5OmniAudioFeatureInputs,
Qwen2_5OmniConditionalGenerationMixin,
Qwen2_5OmniThinkerDummyInputsBuilder,
Qwen2_5OmniThinkerMultiModalDataParser,
Qwen2_5OmniThinkerMultiModalProcessor,
check_interleaved_audio_video,
merge_interleaved_embeddings,
@@ -1152,6 +1153,13 @@ class Qwen3MoeLLMForCausalLM(Qwen3MoeForCausalLM):
)
class Qwen3OmniMoeThinkerMultiModalDataParser(Qwen2_5OmniThinkerMultiModalDataParser):
# Keep the existing raw-media fallback for vision inputs.
embedding_fields = {
"audio": Qwen2_5OmniThinkerMultiModalDataParser.embedding_fields["audio"],
}
class Qwen3OmniMoeThinkerProcessingInfo(
Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo
):
@@ -1178,6 +1186,15 @@ class Qwen3OmniMoeThinkerProcessingInfo(
assert isinstance(feature_extractor, WhisperFeatureExtractor)
return feature_extractor
def get_data_parser(self):
return Qwen3OmniMoeThinkerMultiModalDataParser(
spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size,
target_sr=self.get_feature_extractor().sampling_rate,
target_channels=self.get_target_channels(),
expected_hidden_size=self._get_expected_hidden_size(),
allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
)
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
return {"audio": None, "image": None, "video": None}
+11 -1
View File
@@ -114,6 +114,12 @@ UltravoxAudioInputs: TypeAlias = (
)
class UltravoxMultiModalDataParser(MultiModalDataParser):
embedding_fields = {
"audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"},
}
class UltravoxProcessingInfo(BaseProcessingInfo):
def get_hf_processor(self, **kwargs: object) -> ProcessorMixin:
config = self.ctx.model_config.hf_config
@@ -145,10 +151,11 @@ class UltravoxProcessingInfo(BaseProcessingInfo):
def get_data_parser(self):
feature_extractor = self.get_feature_extractor()
return MultiModalDataParser(
return UltravoxMultiModalDataParser(
target_sr=feature_extractor.sampling_rate,
target_channels=self.get_target_channels(),
expected_hidden_size=self._get_expected_hidden_size(),
allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
)
def get_target_channels(self) -> int:
@@ -247,6 +254,7 @@ class UltravoxMultiModalProcessor(BaseMultiModalProcessor[UltravoxProcessingInfo
# num_chunks can convert audio_chunked to audio batch dimension
audio_num_chunks=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
audio_embeds=MultiModalFieldConfig.batched("audio"),
audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
)
def _get_prompt_updates(
@@ -273,6 +281,8 @@ class UltravoxMultiModalProcessor(BaseMultiModalProcessor[UltravoxProcessingInfo
)
def get_replacement_ultravox(item_idx: int):
if "audio_num_tokens" in out_mm_data:
return [replacement_id] * int(out_mm_data["audio_num_tokens"][item_idx])
start = chunks_start_idx[item_idx]
end = chunks_start_idx[item_idx + 1]
audio_token_len = out_mm_data["audio_token_len"][start:end].sum()
+43 -1
View File
@@ -664,9 +664,46 @@ class MultiModalDataParser:
assert_never(video)
def _parse_audio_embedding_data(
self, data: dict[str, torch.Tensor]
) -> DictEmbeddingItems:
counts = data.get("audio_num_tokens")
if (
not isinstance(counts, torch.Tensor)
or counts.dtype not in (torch.int32, torch.int64)
or counts.ndim not in (1, 2)
or (counts.ndim == 2 and counts.shape[1] != 1)
or not torch.all(counts > 0)
):
raise ValueError(
"audio_num_tokens must contain one positive integer per audio"
)
counts = counts.flatten()
if "audio_embeds" in data:
embeds = data["audio_embeds"]
if len(embeds) != len(counts) or any(
len(embedding) != count for embedding, count in zip(embeds, counts)
):
raise ValueError("audio_num_tokens does not match audio_embeds")
required, optional = self.embedding_field_sets("audio")
return DictEmbeddingItems(
{**data, "audio_num_tokens": counts},
modality="audio",
required_fields=required,
optional_fields=optional,
fields_factory=lambda _: {
"audio_embeds": MultiModalFieldConfig.batched("audio"),
"audio_num_tokens": MultiModalFieldConfig.batched(
"audio", keep_on_cpu=True
),
},
)
def _parse_audio_data(
self,
data: ModalityData[AudioItem],
data: dict[str, torch.Tensor] | ModalityData[AudioItem],
) -> ModalityDataItems[Any, Any] | None:
if data is None:
return None
@@ -674,6 +711,11 @@ class MultiModalDataParser:
if self.is_embeddings(data):
return AudioEmbeddingItems(data, self.expected_hidden_size)
if isinstance(data, dict) and "audio_num_tokens" in self.embedding_fields.get(
"audio", {}
):
return self._parse_audio_embedding_data(data)
data_items: list[AudioItem | None]
if (
(is_list_of(data, float) and len(data) > 0)