[Model][Core] Enable async_chunk streaming pipeline for CosyVoice3 (#1703)
Signed-off-by: linyueqian <linyueqian@outlook.com> Signed-off-by: indevn <indevn@outlook.com> Co-authored-by: linyueqian <linyueqian@outlook.com>
This commit is contained in:
indevnandlinyueqian authored and GitHub committed
1 parent
d92439c155
commit
6fc38e0467
18 files changed
+2094
-131
No files matched your search
@@ -2,6 +2,8 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for CosyVoice3 components."""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
@@ -247,3 +249,32 @@ class TestSDPAFallback:
|
||||
|
||||
assert out.shape == (batch, seq_len, heads, dim)
|
||||
assert out.dtype == torch.float32
|
||||
|
||||
|
||||
def test_code2wav_forward_finalizes_hift_tail():
|
||||
from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav
|
||||
|
||||
class DummyHiFT(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.m_source = SimpleNamespace(l_linear=SimpleNamespace(weight=torch.ones(1, dtype=torch.float32)))
|
||||
self.finalize_calls: list[bool] = []
|
||||
|
||||
def inference(self, speech_feat, finalize=True):
|
||||
self.finalize_calls.append(bool(finalize))
|
||||
return torch.zeros((speech_feat.shape[0], 1, speech_feat.shape[-1]), dtype=speech_feat.dtype), None
|
||||
|
||||
model = object.__new__(CosyVoice3Code2Wav)
|
||||
nn.Module.__init__(model)
|
||||
model.hift = DummyHiFT()
|
||||
model._forward_mel = lambda **_: torch.ones((1, 80, 8), dtype=torch.float32)
|
||||
|
||||
out = model.forward(
|
||||
token=torch.tensor([[1, 2, 3]], dtype=torch.int32),
|
||||
prompt_token=torch.tensor([[4, 5]], dtype=torch.int32),
|
||||
prompt_feat=torch.ones((1, 4, 80), dtype=torch.float32),
|
||||
embedding=torch.ones((1, 192), dtype=torch.float32),
|
||||
)
|
||||
|
||||
assert out.shape == (1, 1, 8)
|
||||
assert model.hift.finalize_calls == [True]
|
||||
@@ -0,0 +1,463 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from threading import Lock
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from vllm.v1.outputs import SamplerOutput
|
||||
from vllm.v1.sample.logits_processor.state import LogitsProcessors
|
||||
from vllm.v1.sample.metadata import SamplingMetadata
|
||||
|
||||
from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import CosyVoice3Model
|
||||
from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _DummyCode2Wav:
|
||||
def __init__(
|
||||
self,
|
||||
vocab_size: int,
|
||||
num_samples: int = 32,
|
||||
outputs: list[tuple[torch.Tensor, dict[str, object] | None]] | None = None,
|
||||
):
|
||||
self.input_embedding = SimpleNamespace(num_embeddings=vocab_size)
|
||||
self.num_samples = num_samples
|
||||
self.outputs = list(outputs or [])
|
||||
self.forward_calls: list[dict[str, object]] = []
|
||||
self.forward_streaming_calls: list[dict[str, object]] = []
|
||||
|
||||
def forward(self, **kwargs):
|
||||
self.forward_calls.append(kwargs)
|
||||
token = kwargs["token"]
|
||||
num_samples = int(token.shape[-1])
|
||||
return torch.linspace(-1.0, 1.0, max(num_samples, 1), dtype=torch.float32).reshape(1, 1, -1)
|
||||
|
||||
def forward_streaming(self, **kwargs):
|
||||
self.forward_streaming_calls.append(kwargs)
|
||||
if self.outputs:
|
||||
return self.outputs.pop(0)
|
||||
|
||||
token = kwargs["token"]
|
||||
num_samples = int(token.shape[-1])
|
||||
audio = torch.linspace(-1.0, 1.0, max(num_samples, 1), dtype=torch.float32).reshape(1, 1, -1)
|
||||
new_state = None
|
||||
if not kwargs.get("finalize", False):
|
||||
new_state = {
|
||||
"mel": torch.ones((1, 80, max(num_samples, 1)), dtype=torch.float32),
|
||||
"speech_offset": audio.shape[-1],
|
||||
}
|
||||
return audio, new_state
|
||||
|
||||
|
||||
def _make_code2wav_model(
|
||||
*,
|
||||
with_stride_cfg: bool = False,
|
||||
num_samples: int = 32,
|
||||
outputs: list[tuple[torch.Tensor, dict[str, object] | None]] | None = None,
|
||||
) -> CosyVoice3Model:
|
||||
model = object.__new__(CosyVoice3Model)
|
||||
nn.Module.__init__(model)
|
||||
model.model_stage = "cosyvoice3_code2wav"
|
||||
hift_cfg = {} if not with_stride_cfg else {"upsample_rates": [8, 5, 3], "istft_params": {"hop_len": 4}}
|
||||
model.config = SimpleNamespace(
|
||||
sample_rate=24000,
|
||||
hift=hift_cfg,
|
||||
token_frame_rate=25 if with_stride_cfg else 0,
|
||||
token_mel_ratio=2 if with_stride_cfg else 0,
|
||||
)
|
||||
model.code2wav = _DummyCode2Wav(vocab_size=4, num_samples=num_samples, outputs=outputs)
|
||||
model.source_cache_len = 4
|
||||
model.speech_window = torch.hamming_window(8, periodic=False)
|
||||
model._stream_audio_cache_by_req = {}
|
||||
model._stream_audio_cache_lock = Lock()
|
||||
model._stream_vocoder_cache_by_req = {}
|
||||
return model
|
||||
|
||||
|
||||
def _make_talker_model() -> CosyVoice3Model:
|
||||
model = object.__new__(CosyVoice3Model)
|
||||
nn.Module.__init__(model)
|
||||
model.model_stage = "cosyvoice3_talker"
|
||||
model.config = SimpleNamespace(
|
||||
llm={
|
||||
"speech_token_size": 6561,
|
||||
"eos_token_id": 6562,
|
||||
"sampling": {
|
||||
"top_p": 0.8,
|
||||
"top_k": 25,
|
||||
"win_size": 10,
|
||||
"tau_r": 0.1,
|
||||
},
|
||||
},
|
||||
vocab_size=151923,
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def _make_sampling_metadata(
|
||||
*,
|
||||
output_token_ids: list[list[int]],
|
||||
repetition_penalty: float = 2.0,
|
||||
) -> SamplingMetadata:
|
||||
return SamplingMetadata(
|
||||
temperature=torch.tensor([1.0], dtype=torch.float32),
|
||||
all_greedy=False,
|
||||
all_random=True,
|
||||
top_p=torch.tensor([0.8], dtype=torch.float32),
|
||||
top_k=torch.tensor([25], dtype=torch.int32),
|
||||
generators={},
|
||||
max_num_logprobs=None,
|
||||
no_penalties=False,
|
||||
prompt_token_ids=None,
|
||||
frequency_penalties=torch.zeros(1, dtype=torch.float32),
|
||||
presence_penalties=torch.zeros(1, dtype=torch.float32),
|
||||
repetition_penalties=torch.tensor([repetition_penalty], dtype=torch.float32),
|
||||
output_token_ids=output_token_ids,
|
||||
allowed_token_ids_mask=None,
|
||||
bad_words_token_ids={},
|
||||
logitsprocs=LogitsProcessors(),
|
||||
)
|
||||
|
||||
|
||||
def test_split_request_ids_uses_seq_token_counts():
|
||||
ids = torch.tensor([10, 11, 12, 13, 14], dtype=torch.long)
|
||||
chunks = CosyVoice3Model._split_request_ids(ids, [2, 2, 2])
|
||||
assert [c.tolist() for c in chunks] == [[10, 11], [12, 13], [14]]
|
||||
|
||||
|
||||
def test_split_request_ids_honors_single_request_seq_token_counts():
|
||||
ids = torch.tensor([10, 11, 12, 13, 14], dtype=torch.long)
|
||||
chunks = CosyVoice3Model._split_request_ids(ids, [3])
|
||||
assert [c.tolist() for c in chunks] == [[10, 11, 12]]
|
||||
|
||||
|
||||
def test_sanitize_codec_tokens_filters_out_of_range():
|
||||
model = _make_code2wav_model()
|
||||
raw = torch.tensor([-1, 0, 3, 4, 99], dtype=torch.long)
|
||||
clean = model._sanitize_codec_tokens(raw)
|
||||
assert clean.tolist() == [0, 3]
|
||||
|
||||
|
||||
def test_forward_prefers_token_offset_when_present():
|
||||
model = _make_code2wav_model()
|
||||
|
||||
runtime_info = [
|
||||
{
|
||||
"speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
|
||||
"speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
|
||||
"embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
|
||||
"token_offset": 2,
|
||||
"left_context_size": 1,
|
||||
}
|
||||
]
|
||||
|
||||
out = model.forward(
|
||||
input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
positions=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
model_intermediate_buffer=runtime_info,
|
||||
seq_token_counts=[3],
|
||||
)
|
||||
|
||||
assert len(out.multimodal_outputs["audio"]) == 1
|
||||
assert out.multimodal_outputs["audio"][0].numel() > 0
|
||||
assert len(model.code2wav.forward_streaming_calls) == 1
|
||||
call = model.code2wav.forward_streaming_calls[0]
|
||||
assert call["token"].shape == (1, 3)
|
||||
assert call["token_offset_tokens"] == 2
|
||||
assert call["finalize"] is False
|
||||
|
||||
|
||||
def test_forward_falls_back_to_left_context_size_for_backward_compat():
|
||||
model = _make_code2wav_model()
|
||||
|
||||
runtime_info = [
|
||||
{
|
||||
"speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
|
||||
"speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
|
||||
"embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
|
||||
"left_context_size": 2,
|
||||
}
|
||||
]
|
||||
|
||||
model.forward(
|
||||
input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
positions=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
model_intermediate_buffer=runtime_info,
|
||||
seq_token_counts=[3],
|
||||
)
|
||||
|
||||
assert model.code2wav.forward_streaming_calls[0]["token_offset_tokens"] == 2
|
||||
|
||||
|
||||
def test_forward_ignores_single_request_padded_tail_tokens():
|
||||
model = _make_code2wav_model(with_stride_cfg=True)
|
||||
runtime_info = [
|
||||
{
|
||||
"speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
|
||||
"speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
|
||||
"embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
|
||||
"token_offset": 0,
|
||||
}
|
||||
]
|
||||
|
||||
out = model.forward(
|
||||
input_ids=torch.tensor([0, 1, 2, 3, 3], dtype=torch.long),
|
||||
positions=torch.tensor([0, 1, 2, 3, 4], dtype=torch.long),
|
||||
model_intermediate_buffer=runtime_info,
|
||||
seq_token_counts=[3],
|
||||
)
|
||||
|
||||
# The padded tail must not contribute to code2wav length.
|
||||
assert out.multimodal_outputs["audio"][0].numel() == 3
|
||||
assert model.code2wav.forward_streaming_calls[0]["token"].tolist() == [[0, 1, 2]]
|
||||
|
||||
|
||||
def test_forward_uses_non_stream_decode_without_chunk_metadata():
|
||||
model = _make_code2wav_model()
|
||||
|
||||
runtime_info = [
|
||||
{
|
||||
"speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
|
||||
"speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
|
||||
"embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
|
||||
"prefix_ids": [101, 102],
|
||||
"generated_len": 3,
|
||||
}
|
||||
]
|
||||
|
||||
out = model.forward(
|
||||
input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
positions=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
model_intermediate_buffer=runtime_info,
|
||||
seq_token_counts=[3],
|
||||
)
|
||||
|
||||
assert out.multimodal_outputs["audio"][0].numel() == 3
|
||||
assert len(model.code2wav.forward_calls) == 1
|
||||
assert len(model.code2wav.forward_streaming_calls) == 0
|
||||
call = model.code2wav.forward_calls[0]
|
||||
assert call["token"].tolist() == [[0, 1, 2]]
|
||||
|
||||
|
||||
def test_forward_reuses_streaming_cache_state_between_chunks():
|
||||
model = _make_code2wav_model(
|
||||
outputs=[
|
||||
(
|
||||
torch.arange(4, dtype=torch.float32).reshape(1, 1, -1),
|
||||
{"mel": torch.ones((1, 80, 3), dtype=torch.float32), "speech_offset": 4},
|
||||
),
|
||||
(
|
||||
torch.full((1, 1, 2), 9.0, dtype=torch.float32),
|
||||
{"mel": torch.ones((1, 80, 5), dtype=torch.float32), "speech_offset": 6},
|
||||
),
|
||||
]
|
||||
)
|
||||
runtime_info = [
|
||||
{
|
||||
"req_id": ["rid-stream"],
|
||||
"speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
|
||||
"speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
|
||||
"embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
|
||||
"token_offset": 0,
|
||||
"stream_finished": torch.tensor(False),
|
||||
}
|
||||
]
|
||||
|
||||
out1 = model.forward(
|
||||
input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
positions=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
model_intermediate_buffer=runtime_info,
|
||||
seq_token_counts=[3],
|
||||
)
|
||||
assert out1.multimodal_outputs["audio"][0].tolist() == [0.0, 1.0, 2.0, 3.0]
|
||||
assert model.code2wav.forward_streaming_calls[0]["cache_state"] is None
|
||||
|
||||
out2 = model.forward(
|
||||
input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
positions=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
model_intermediate_buffer=runtime_info,
|
||||
seq_token_counts=[3],
|
||||
)
|
||||
assert out2.multimodal_outputs["audio"][0].tolist() == [9.0, 9.0]
|
||||
cache_state = model.code2wav.forward_streaming_calls[1]["cache_state"]
|
||||
assert cache_state is not None
|
||||
assert cache_state["speech_offset"] == 4
|
||||
assert "rid-stream" in model._stream_vocoder_cache_by_req
|
||||
|
||||
|
||||
def test_forward_clears_streaming_cache_on_terminal_chunk():
|
||||
model = _make_code2wav_model(
|
||||
outputs=[
|
||||
(
|
||||
torch.arange(4, dtype=torch.float32).reshape(1, 1, -1),
|
||||
{"mel": torch.ones((1, 80, 3), dtype=torch.float32), "speech_offset": 4},
|
||||
),
|
||||
(
|
||||
torch.full((1, 1, 1), 7.0, dtype=torch.float32),
|
||||
None,
|
||||
),
|
||||
]
|
||||
)
|
||||
runtime_info = [
|
||||
{
|
||||
"req_id": ["rid-stream"],
|
||||
"speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
|
||||
"speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
|
||||
"embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
|
||||
"token_offset": 0,
|
||||
"stream_finished": torch.tensor(False),
|
||||
}
|
||||
]
|
||||
|
||||
model.forward(
|
||||
input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
positions=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
model_intermediate_buffer=runtime_info,
|
||||
seq_token_counts=[3],
|
||||
)
|
||||
assert "rid-stream" in model._stream_vocoder_cache_by_req
|
||||
|
||||
runtime_info[0]["stream_finished"] = torch.tensor(True)
|
||||
out = model.forward(
|
||||
input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
positions=torch.tensor([0, 1, 2], dtype=torch.long),
|
||||
model_intermediate_buffer=runtime_info,
|
||||
seq_token_counts=[3],
|
||||
)
|
||||
assert out.multimodal_outputs["audio"][0].tolist() == [7.0]
|
||||
assert "rid-stream" not in model._stream_vocoder_cache_by_req
|
||||
|
||||
|
||||
def test_sample_uses_ras_rejection_for_recent_repetition():
|
||||
model = _make_talker_model()
|
||||
metadata = _make_sampling_metadata(output_token_ids=[[1] * 10])
|
||||
logits = torch.tensor([[-1e9, 10.0, 0.0]], dtype=torch.float32)
|
||||
|
||||
out = model.sample(logits, metadata)
|
||||
|
||||
assert out is not None
|
||||
assert out.sampled_token_ids.tolist() == [[2]]
|
||||
|
||||
|
||||
def test_sample_tolerates_padded_rows_without_history():
|
||||
model = _make_talker_model()
|
||||
metadata = _make_sampling_metadata(output_token_ids=[[1] * 10])
|
||||
logits = torch.tensor(
|
||||
[
|
||||
[-1e9, 10.0, 0.0],
|
||||
[-1e9, 0.0, 10.0],
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
out = model.sample(logits, metadata)
|
||||
|
||||
assert out is not None
|
||||
assert out.sampled_token_ids.shape == (2, 1)
|
||||
|
||||
|
||||
def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in():
|
||||
metadata = _make_sampling_metadata(output_token_ids=[[1, 2, 3]])
|
||||
expected = SamplerOutput(
|
||||
sampled_token_ids=torch.tensor([[7]], dtype=torch.int32),
|
||||
logprobs_tensors=None,
|
||||
)
|
||||
calls: list[torch.Tensor] = []
|
||||
|
||||
class _DummyInputBatch:
|
||||
def __init__(self):
|
||||
self.sampling_metadata = metadata
|
||||
self.updated = False
|
||||
|
||||
def update_async_output_token_ids(self):
|
||||
self.updated = True
|
||||
|
||||
runner = object.__new__(GPUARModelRunner)
|
||||
runner.input_batch = _DummyInputBatch()
|
||||
runner.model = SimpleNamespace(
|
||||
prefer_model_sampler=True,
|
||||
sample=lambda logits, sampling_metadata: calls.append(logits.clone()) or expected,
|
||||
)
|
||||
runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
|
||||
|
||||
out = runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
|
||||
|
||||
assert out is expected
|
||||
assert runner.input_batch.updated is False
|
||||
assert len(calls) == 1
|
||||
|
||||
|
||||
def test_gpu_ar_model_runner_supplies_req_output_history_to_model_sampler():
|
||||
metadata = _make_sampling_metadata(output_token_ids=[])
|
||||
seen_histories: list[list[list[int]]] = []
|
||||
|
||||
class _DummyInputBatch:
|
||||
def __init__(self):
|
||||
self.sampling_metadata = metadata
|
||||
self.req_output_token_ids = [[1, 2, 3]]
|
||||
self.req_ids = ["rid-1"]
|
||||
self.sampled_token_ids_cpu = None
|
||||
self.async_copy_ready_event = None
|
||||
self.prev_req_id_to_index = None
|
||||
|
||||
def update_async_output_token_ids(self):
|
||||
raise AssertionError("fallback async repair should not run for model sampler path")
|
||||
|
||||
runner = object.__new__(GPUARModelRunner)
|
||||
runner.input_batch = _DummyInputBatch()
|
||||
runner.model = SimpleNamespace(
|
||||
prefer_model_sampler=True,
|
||||
sample=lambda logits, sampling_metadata: seen_histories.append(
|
||||
[list(x) for x in sampling_metadata.output_token_ids]
|
||||
)
|
||||
or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None),
|
||||
)
|
||||
runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
|
||||
|
||||
runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
|
||||
|
||||
assert seen_histories == [[[1, 2, 3]]]
|
||||
|
||||
|
||||
def test_gpu_ar_model_runner_repairs_async_placeholders_for_model_sampler():
|
||||
metadata = _make_sampling_metadata(output_token_ids=[])
|
||||
seen_histories: list[list[list[int]]] = []
|
||||
|
||||
class _ReadyEvent:
|
||||
def __init__(self):
|
||||
self.synced = False
|
||||
|
||||
def synchronize(self):
|
||||
self.synced = True
|
||||
|
||||
class _DummyInputBatch:
|
||||
def __init__(self):
|
||||
self.sampling_metadata = metadata
|
||||
self.req_output_token_ids = [[11, -1]]
|
||||
self.req_ids = ["rid-1"]
|
||||
self.sampled_token_ids_cpu = torch.tensor([[29]], dtype=torch.int32)
|
||||
self.async_copy_ready_event = _ReadyEvent()
|
||||
self.prev_req_id_to_index = {"rid-1": 0}
|
||||
|
||||
def update_async_output_token_ids(self):
|
||||
raise AssertionError("fallback async repair should not run for model sampler path")
|
||||
|
||||
runner = object.__new__(GPUARModelRunner)
|
||||
runner.input_batch = _DummyInputBatch()
|
||||
runner.model = SimpleNamespace(
|
||||
prefer_model_sampler=True,
|
||||
sample=lambda logits, sampling_metadata: seen_histories.append(
|
||||
[list(x) for x in sampling_metadata.output_token_ids]
|
||||
)
|
||||
or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None),
|
||||
)
|
||||
runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
|
||||
|
||||
runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
|
||||
|
||||
assert runner.input_batch.async_copy_ready_event.synced is True
|
||||
assert seen_histories == [[[11, 29]]]
|
||||
+267
@@ -0,0 +1,267 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from collections import defaultdict
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import talker2code2wav_async_chunk, text2flow
|
||||
|
||||
|
||||
def _source_output(request_id: str, prompt_ids: list[int], out_ids: list[int], mm: dict):
|
||||
return SimpleNamespace(
|
||||
request_id=request_id,
|
||||
prompt_token_ids=prompt_ids,
|
||||
outputs=[SimpleNamespace(token_ids=out_ids, multimodal_output=mm)],
|
||||
)
|
||||
|
||||
|
||||
def _transfer_manager(
|
||||
*,
|
||||
chunk_frames: int = 2,
|
||||
pre_lookahead_frames: int = 0,
|
||||
stream_scale_factor: int = 1,
|
||||
max_chunk_frames: int | None = None,
|
||||
):
|
||||
if max_chunk_frames is None:
|
||||
max_chunk_frames = chunk_frames
|
||||
return SimpleNamespace(
|
||||
code_prompt_token_ids=defaultdict(list),
|
||||
request_payload={},
|
||||
connector=SimpleNamespace(
|
||||
config={
|
||||
"extra": {
|
||||
"codec_chunk_frames": chunk_frames,
|
||||
"codec_pre_lookahead_frames": pre_lookahead_frames,
|
||||
"codec_max_chunk_frames": max_chunk_frames,
|
||||
"codec_stream_scale_factor": stream_scale_factor,
|
||||
"codec_vocab_size": 6561,
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_text2flow_supports_batched_source_outputs():
|
||||
stage_list = [
|
||||
SimpleNamespace(
|
||||
engine_outputs=[
|
||||
_source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[1, 2]])}),
|
||||
_source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[3, 4]])}),
|
||||
]
|
||||
)
|
||||
]
|
||||
|
||||
outputs = text2flow(stage_list=stage_list, engine_input_source=[0], prompt=None)
|
||||
|
||||
assert len(outputs) == 2
|
||||
assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
|
||||
assert outputs[1]["prompt_token_ids"] == [4, 5]
|
||||
assert outputs[0]["additional_information"]["prefix_ids"] == [10, 11]
|
||||
assert outputs[1]["additional_information"]["prefix_ids"] == [20, 21]
|
||||
|
||||
|
||||
def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
|
||||
transfer_manager = _transfer_manager()
|
||||
request = SimpleNamespace(
|
||||
external_req_id="rid-0",
|
||||
output_token_ids=[1, 2, 6562, 3],
|
||||
additional_information={
|
||||
"speech_token": [torch.tensor([[11, 12, 13]])],
|
||||
"speech_feat": [torch.tensor([[[0.1, 0.2], [0.3, 0.4]]])],
|
||||
"embedding": [torch.tensor([[0.5, 0.6]])],
|
||||
},
|
||||
is_finished=lambda: True,
|
||||
)
|
||||
|
||||
payload = talker2code2wav_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
is_finished=True,
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
assert payload["finished"].item() is True
|
||||
assert payload["code_predictor_codes"] == [1, 2, 3]
|
||||
assert payload["token_offset"] == 0
|
||||
assert payload["left_context_size"] == 0
|
||||
assert payload["req_id"] == ["rid-0"]
|
||||
assert payload["stream_finished"].item() is True
|
||||
assert "speech_token" in payload
|
||||
assert "speech_feat" in payload
|
||||
assert "embedding" in payload
|
||||
|
||||
|
||||
def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes():
|
||||
transfer_manager = _transfer_manager(chunk_frames=25)
|
||||
request = SimpleNamespace(
|
||||
external_req_id="rid-eof",
|
||||
output_token_ids=[6561, 6562], # all filtered out
|
||||
additional_information={},
|
||||
is_finished=lambda: True,
|
||||
)
|
||||
|
||||
payload = talker2code2wav_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
is_finished=True,
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
assert payload["code_predictor_codes"] == []
|
||||
assert payload["finished"].item() is True
|
||||
|
||||
|
||||
def test_talker2code2wav_async_chunk_does_not_reemit_without_new_tokens():
|
||||
transfer_manager = _transfer_manager()
|
||||
request = SimpleNamespace(
|
||||
external_req_id="rid-stable",
|
||||
output_token_ids=[1, 2],
|
||||
additional_information={},
|
||||
is_finished=lambda: False,
|
||||
)
|
||||
|
||||
payload1 = talker2code2wav_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
is_finished=False,
|
||||
)
|
||||
payload2 = talker2code2wav_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
is_finished=False,
|
||||
)
|
||||
|
||||
assert payload1 is not None
|
||||
assert payload1["code_predictor_codes"] == [1, 2]
|
||||
assert payload1["token_offset"] == 0
|
||||
assert payload2 is None
|
||||
|
||||
|
||||
def test_talker2code2wav_async_chunk_waits_for_prelookahead_and_emits_cumulative_prefix():
|
||||
transfer_manager = _transfer_manager(pre_lookahead_frames=1)
|
||||
request = SimpleNamespace(
|
||||
external_req_id="rid-pre",
|
||||
output_token_ids=[1, 2],
|
||||
additional_information={},
|
||||
is_finished=lambda: False,
|
||||
)
|
||||
|
||||
payload_pending = talker2code2wav_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
is_finished=False,
|
||||
)
|
||||
request.output_token_ids = [1, 2, 3]
|
||||
payload_ready = talker2code2wav_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
is_finished=False,
|
||||
)
|
||||
|
||||
assert payload_pending is None
|
||||
assert payload_ready is not None
|
||||
assert payload_ready["code_predictor_codes"] == [1, 2, 3]
|
||||
assert payload_ready["token_offset"] == 0
|
||||
assert payload_ready["finished"].item() is False
|
||||
|
||||
|
||||
def test_talker2code2wav_async_chunk_final_flush_uses_previous_token_offset():
|
||||
transfer_manager = _transfer_manager(pre_lookahead_frames=1)
|
||||
request = SimpleNamespace(
|
||||
external_req_id="rid-tail",
|
||||
output_token_ids=[3, 4, 5],
|
||||
additional_information={},
|
||||
is_finished=lambda: False,
|
||||
)
|
||||
|
||||
payload_stream = talker2code2wav_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
is_finished=False,
|
||||
)
|
||||
request.output_token_ids = [3, 4, 5, 6]
|
||||
payload_final = talker2code2wav_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
is_finished=True,
|
||||
)
|
||||
|
||||
assert payload_stream is not None
|
||||
assert payload_stream["finished"].item() is False
|
||||
assert payload_stream["code_predictor_codes"] == [3, 4, 5]
|
||||
assert payload_stream["token_offset"] == 0
|
||||
assert payload_final is not None
|
||||
assert payload_final["finished"].item() is True
|
||||
assert payload_final["code_predictor_codes"] == [3, 4, 5, 6]
|
||||
assert payload_final["token_offset"] == 2
|
||||
|
||||
|
||||
def test_talker2code2wav_async_chunk_respects_prompt_token_pad_on_first_chunk():
|
||||
transfer_manager = _transfer_manager(pre_lookahead_frames=1)
|
||||
request = SimpleNamespace(
|
||||
external_req_id="rid-pad",
|
||||
output_token_ids=[8, 9, 10],
|
||||
additional_information={
|
||||
"speech_token": [torch.tensor([[1, 2, 3]])],
|
||||
},
|
||||
is_finished=lambda: False,
|
||||
)
|
||||
|
||||
payload_pending = talker2code2wav_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
is_finished=False,
|
||||
)
|
||||
request.output_token_ids = [8, 9, 10, 11]
|
||||
payload_ready = talker2code2wav_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
is_finished=False,
|
||||
)
|
||||
|
||||
assert payload_pending is None
|
||||
assert payload_ready is not None
|
||||
assert payload_ready["code_predictor_codes"] == [8, 9, 10, 11]
|
||||
assert payload_ready["token_offset"] == 0
|
||||
|
||||
|
||||
def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio():
|
||||
transfer_manager = _transfer_manager()
|
||||
request = SimpleNamespace(
|
||||
external_req_id="rid-eof-tail",
|
||||
output_token_ids=[3, 4],
|
||||
additional_information={},
|
||||
is_finished=lambda: False,
|
||||
)
|
||||
|
||||
payload_stream = talker2code2wav_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
is_finished=False,
|
||||
)
|
||||
payload_final = talker2code2wav_async_chunk(
|
||||
transfer_manager=transfer_manager,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
is_finished=True,
|
||||
)
|
||||
|
||||
assert payload_stream is not None
|
||||
assert payload_stream["finished"].item() is False
|
||||
assert payload_stream["code_predictor_codes"] == [3, 4]
|
||||
assert payload_final is not None
|
||||
assert payload_final["finished"].item() is True
|
||||
assert payload_final["code_predictor_codes"] == []
|
||||
Reference in new issue
Block a user