[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:
co-authored by
Cyrus Leung
parent
826e300cde
commit
3ee2906f47
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
),
|
||||
|
||||
@@ -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}
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user