[Model] Add GLM-TTS text-to-speech model support (#3141)
Signed-off-by: BeatSeat <wendavid552@gmail.com> Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
This commit is contained in:
1 parent
748470c1fb
commit
28ce618f0e
38 files changed
+7568
-7
No files matched your search
@@ -0,0 +1,351 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.models.common.nucleus_ras_sampling import ras_sample_one
|
||||
from vllm_omni.model_executor.models.glm_tts import glm_tts as glm_tts_module
|
||||
from vllm_omni.model_executor.models.glm_tts.glm_tts import GLMTTSForConditionalGeneration
|
||||
from vllm_omni.model_executor.models.glm_tts.text_frontend import GLMTTSTextFrontend
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_glm_tts_ras_fallback_matches_official_random_resample() -> None:
|
||||
sampled = ras_sample_one(
|
||||
torch.tensor([0.0, float("-inf")]),
|
||||
decoded_tokens=[0] * 10,
|
||||
top_p=1.0,
|
||||
top_k=1,
|
||||
win_size=10,
|
||||
tau_r=0.1,
|
||||
generator=None,
|
||||
)
|
||||
|
||||
assert sampled == 0
|
||||
|
||||
|
||||
def test_glm_tts_text_frontend_applies_official_symbol_replacements() -> None:
|
||||
frontend = GLMTTSTextFrontend()
|
||||
|
||||
assert frontend.pre_replace("3-2 咯,") == "3减2 喽,"
|
||||
assert frontend.post_replace("α②∈A/B~C²") == "阿尔法二属于A每B到C平方"
|
||||
|
||||
|
||||
def test_glm_tts_call_hf_processor_builds_prompt_text_text_boa(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
processor = object.__new__(glm_tts_module.GLMTTSMultiModalProcessor)
|
||||
processor.info = SimpleNamespace(
|
||||
ctx=SimpleNamespace(
|
||||
get_hf_config=lambda *_args, **_kwargs: SimpleNamespace(),
|
||||
model_config=SimpleNamespace(model="fake-model"),
|
||||
)
|
||||
)
|
||||
processor.text_frontend = object()
|
||||
processor.special_ids = {"boa": 99}
|
||||
processor.speech_tokenizer = object()
|
||||
processor.processor_device = torch.device("cpu")
|
||||
processor.campplus_session = object()
|
||||
|
||||
monkeypatch.setattr(
|
||||
processor,
|
||||
"_ensure_cached_runtime_components",
|
||||
lambda *_args, **_kwargs: None,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
processor,
|
||||
"_encode_text",
|
||||
lambda text: torch.tensor([[11, 12]]) if text.endswith(" ") else torch.tensor([[21, 22, 23]]),
|
||||
)
|
||||
monkeypatch.setattr(processor, "_get_audio", lambda _mm_data: (torch.zeros(8), 24000))
|
||||
monkeypatch.setattr(
|
||||
glm_tts_module,
|
||||
"_normalize_glm_tts_processor_text",
|
||||
lambda _frontend, text, add_trailing_space=False: text + (" " if add_trailing_space else ""),
|
||||
)
|
||||
monkeypatch.setattr(glm_tts_module, "extract_prompt_speech_token", lambda *_args, **_kwargs: [7, 8, 9, 10])
|
||||
monkeypatch.setattr(glm_tts_module, "extract_prompt_feat", lambda *_args, **_kwargs: torch.ones(4, 3))
|
||||
monkeypatch.setattr(glm_tts_module, "extract_spk_embedding", lambda *_args, **_kwargs: [0.1, 0.2, 0.3])
|
||||
|
||||
out = processor._call_hf_processor(
|
||||
prompt="target text",
|
||||
mm_data={"audio": object()},
|
||||
mm_kwargs={"prompt_text": "reference text"},
|
||||
tok_kwargs={},
|
||||
)
|
||||
|
||||
assert out["input_ids"].tolist() == [[11, 12, 21, 22, 23, 99]]
|
||||
assert out["prompt_speech_token"].tolist() == [[7, 8, 9, 10]]
|
||||
assert out["prompt_speech_token_len"][0].item() == 4
|
||||
assert out["glm_tts_text_token_len"][0].item() == 3
|
||||
|
||||
|
||||
def test_glm_tts_prompt_updates_append_single_audio_suffix() -> None:
|
||||
processor = object.__new__(glm_tts_module.GLMTTSMultiModalProcessor)
|
||||
out_mm_kwargs = {
|
||||
"audio": [{"prompt_speech_token_len": SimpleNamespace(data=[torch.tensor([5], dtype=torch.long)])}]
|
||||
}
|
||||
|
||||
updates = processor._get_prompt_updates(
|
||||
mm_items=object(),
|
||||
hf_processor_mm_kwargs={},
|
||||
out_mm_kwargs=out_mm_kwargs,
|
||||
)
|
||||
|
||||
assert len(updates) == 1
|
||||
assert updates[0].modality == "audio"
|
||||
assert updates[0].insertion(0) == [1, 1, 1, 1, 1]
|
||||
|
||||
|
||||
def test_glm_tts_dummy_inputs_forbid_dummy_wav(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
builder = object.__new__(glm_tts_module.GLMTTSDummyInputsBuilder)
|
||||
monkeypatch.setattr(
|
||||
builder,
|
||||
"_get_dummy_audios",
|
||||
lambda *_args, **_kwargs: pytest.fail("GLM-TTS dummy inputs must not fabricate wav audio"),
|
||||
)
|
||||
|
||||
assert builder.get_dummy_mm_data(seq_len=16, mm_counts={"audio": 1}) == {}
|
||||
|
||||
|
||||
def test_glm_tts_processing_info_declares_audio_token_budget() -> None:
|
||||
info = object.__new__(glm_tts_module.GLMTTSMultiModalProcessingInfo)
|
||||
|
||||
assert info.get_mm_max_tokens_per_item(seq_len=4096, mm_counts={"audio": 1}) == {"audio": 1024}
|
||||
assert info.get_mm_max_tokens_per_item(seq_len=512, mm_counts={"audio": 1}) == {"audio": 512}
|
||||
assert info.get_mm_max_tokens_per_item(seq_len=4096, mm_counts={}) == {}
|
||||
|
||||
|
||||
def test_glm_tts_text_only_processor_skips_voice_clone(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
class FakeTokenizer:
|
||||
def encode(self, text: str) -> list[int]:
|
||||
return [11, 12, 13]
|
||||
|
||||
processor = object.__new__(glm_tts_module.GLMTTSMultiModalProcessor)
|
||||
processor.info = SimpleNamespace(
|
||||
ctx=SimpleNamespace(
|
||||
get_hf_config=lambda *_args, **_kwargs: SimpleNamespace(),
|
||||
model_config=SimpleNamespace(model="fake-model", tokenizer=None, trust_remote_code=True),
|
||||
)
|
||||
)
|
||||
|
||||
monkeypatch.setattr(glm_tts_module, "resolve_glm_tts_model_dir", lambda *_args, **_kwargs: "/fake-root")
|
||||
monkeypatch.setattr(glm_tts_module, "load_glm_tts_tokenizer", lambda *_args, **_kwargs: FakeTokenizer())
|
||||
monkeypatch.setattr(glm_tts_module, "get_glm_tts_special_token_ids", lambda _tokenizer: {"boa": 99})
|
||||
monkeypatch.setattr(glm_tts_module, "GLMTTSTextFrontend", lambda: object())
|
||||
monkeypatch.setattr(
|
||||
glm_tts_module,
|
||||
"_normalize_glm_tts_processor_text",
|
||||
lambda _frontend, text, add_trailing_space=False: text,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
glm_tts_module,
|
||||
"load_voice_clone_frontend",
|
||||
lambda *_args, **_kwargs: pytest.fail("text-only dummy path should not load voice clone frontend"),
|
||||
)
|
||||
|
||||
out = processor._call_hf_processor(prompt="target text", mm_data={}, mm_kwargs={}, tok_kwargs={})
|
||||
|
||||
assert out["input_ids"].tolist() == [[11, 12, 13, 99]]
|
||||
assert out["prompt_speech_token"].shape == (1, 1024)
|
||||
assert out["prompt_speech_token_len"][0].item() == 1024
|
||||
assert out["glm_tts_prompt_text_token_len"][0].item() == 0
|
||||
assert out["glm_tts_text_token_len"][0].item() == 3
|
||||
|
||||
fields = processor._get_mm_fields_config(out, {})
|
||||
assert {key: value.modality for key, value in fields.items()} == {
|
||||
"prompt_speech_token": "audio",
|
||||
"prompt_speech_token_len": "audio",
|
||||
"glm_tts_prompt_text_token_len": "audio",
|
||||
"glm_tts_text_token_len": "audio",
|
||||
}
|
||||
|
||||
updates = processor._get_prompt_updates(
|
||||
mm_items=object(),
|
||||
hf_processor_mm_kwargs={},
|
||||
out_mm_kwargs={
|
||||
"audio": [{"prompt_speech_token_len": SimpleNamespace(data=[out["prompt_speech_token_len"][0]])}]
|
||||
},
|
||||
)
|
||||
assert len(updates) == 1
|
||||
assert len(updates[0].insertion(0)) == 1024
|
||||
|
||||
|
||||
def test_glm_tts_embed_replaces_multimodal_positions_per_request() -> None:
|
||||
class FakeModel(torch.nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.marker = torch.nn.Parameter(torch.zeros(()))
|
||||
|
||||
def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return input_ids.to(torch.float32).reshape(-1, 1)
|
||||
|
||||
model = object.__new__(GLMTTSForConditionalGeneration)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.model_stage = "glm_tts"
|
||||
model.model = FakeModel()
|
||||
|
||||
input_ids = torch.tensor([101, 102, 1, 1, 201, 1, 1, 1])
|
||||
is_multimodal = torch.tensor([False, False, True, True, False, True, True, True])
|
||||
mm_embeds = [
|
||||
torch.tensor([[1001.0], [1002.0]]),
|
||||
torch.tensor([[2001.0], [2002.0], [2003.0]]),
|
||||
]
|
||||
|
||||
out = model.embed_input_ids(input_ids, multimodal_embeddings=mm_embeds, is_multimodal=is_multimodal)
|
||||
|
||||
assert out.squeeze(-1).tolist() == [
|
||||
101.0,
|
||||
102.0,
|
||||
1001.0,
|
||||
1002.0,
|
||||
201.0,
|
||||
2001.0,
|
||||
2002.0,
|
||||
2003.0,
|
||||
]
|
||||
|
||||
|
||||
def test_glm_tts_embed_multimodal_keeps_all_batched_requests() -> None:
|
||||
class FakeModel(torch.nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.marker = torch.nn.Parameter(torch.zeros(()))
|
||||
|
||||
def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return input_ids.to(torch.float32).reshape(*input_ids.shape, 1)
|
||||
|
||||
model = object.__new__(GLMTTSForConditionalGeneration)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.model_stage = "glm_tts"
|
||||
model.model = FakeModel()
|
||||
model._ats = 1000
|
||||
|
||||
out = model.embed_multimodal(
|
||||
prompt_speech_token=[
|
||||
torch.tensor([[1, 2]]),
|
||||
torch.tensor([[3, 4, 5]]),
|
||||
]
|
||||
)
|
||||
|
||||
assert out is not None
|
||||
assert len(out) == 2
|
||||
assert out[0].squeeze(-1).tolist() == [1001.0, 1002.0]
|
||||
assert out[1].squeeze(-1).tolist() == [1003.0, 1004.0, 1005.0]
|
||||
|
||||
|
||||
def test_glm_tts_embed_multimodal_returns_one_item_for_single_request() -> None:
|
||||
class FakeModel(torch.nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.marker = torch.nn.Parameter(torch.zeros(()))
|
||||
|
||||
def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return input_ids.to(torch.float32).reshape(*input_ids.shape, 1)
|
||||
|
||||
model = object.__new__(GLMTTSForConditionalGeneration)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.model_stage = "glm_tts"
|
||||
model.model = FakeModel()
|
||||
model._ats = 1000
|
||||
|
||||
out = model.embed_multimodal(prompt_speech_token=torch.tensor([[1, 2, 3]]))
|
||||
|
||||
assert out is not None
|
||||
assert len(out) == 1
|
||||
assert out[0].squeeze(-1).tolist() == [1001.0, 1002.0, 1003.0]
|
||||
|
||||
|
||||
def test_glm_tts_embed_multimodal_treats_scalar_tensor_list_as_one_item() -> None:
|
||||
class FakeModel(torch.nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.marker = torch.nn.Parameter(torch.zeros(()))
|
||||
|
||||
def embed_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||||
return input_ids.to(torch.float32).reshape(*input_ids.shape, 1)
|
||||
|
||||
model = object.__new__(GLMTTSForConditionalGeneration)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.model_stage = "glm_tts"
|
||||
model.model = FakeModel()
|
||||
model._ats = 1000
|
||||
|
||||
out = model.embed_multimodal(prompt_speech_token=[torch.tensor(1), torch.tensor(2)])
|
||||
|
||||
assert out is not None
|
||||
assert len(out) == 1
|
||||
assert out[0].squeeze(-1).tolist() == [1001.0, 1002.0]
|
||||
|
||||
|
||||
def test_glm_tts_preprocess_treats_one_token_text_tail_as_prefill() -> None:
|
||||
model = object.__new__(GLMTTSForConditionalGeneration)
|
||||
model._pad = 0
|
||||
model._ats = 1000
|
||||
model._ate = 2000
|
||||
model._eoa = 3000
|
||||
model._model_dtype = lambda: torch.float32
|
||||
|
||||
input_ids, input_embeds, info = model.preprocess(
|
||||
torch.tensor([123]),
|
||||
torch.tensor([[1.0, 2.0]]),
|
||||
glm_tts_mm_prefill_done=True,
|
||||
)
|
||||
|
||||
assert input_ids.tolist() == [0]
|
||||
assert input_embeds.tolist() == [[1.0, 2.0]]
|
||||
assert info["glm_tts_mm_prefill_done"] is True
|
||||
assert info["speech_tokens"].tolist() == [[-1]]
|
||||
|
||||
|
||||
def test_glm_tts_preprocess_initializes_generation_bounds_before_sampling() -> None:
|
||||
model = object.__new__(GLMTTSForConditionalGeneration)
|
||||
model._pad = 0
|
||||
model._ats = 1000
|
||||
model._ate = 2000
|
||||
model._eoa = 3000
|
||||
model._model_dtype = lambda: torch.float32
|
||||
model.config = SimpleNamespace(min_token_text_ratio=2.0, max_token_text_ratio=20.0)
|
||||
|
||||
_, _, info = model.preprocess(
|
||||
torch.tensor([11, 12]),
|
||||
torch.tensor([[1.0], [2.0]]),
|
||||
glm_tts_text_token_len=[torch.tensor([27])],
|
||||
)
|
||||
|
||||
assert info["glm_tts_text_token_len"][0].item() == 27
|
||||
|
||||
|
||||
def test_glm_tts_make_omni_output_emits_prompt_speech_len_once() -> None:
|
||||
model = object.__new__(GLMTTSForConditionalGeneration)
|
||||
|
||||
out = model.make_omni_output(
|
||||
torch.randn(2, 4),
|
||||
model_intermediate_buffer=[{}],
|
||||
prompt_speech_token_len=[torch.tensor([9])],
|
||||
glm_tts_text_token_len=[torch.tensor([5])],
|
||||
)
|
||||
|
||||
assert out.multimodal_outputs is not None
|
||||
assert out.multimodal_outputs["prompt_speech_token_len"][0].item() == 9
|
||||
assert out.multimodal_outputs["glm_tts_text_token_len"][0].item() == 5
|
||||
|
||||
|
||||
def test_glm_tts_postprocess_caches_prompt_speech_token_len() -> None:
|
||||
model = object.__new__(GLMTTSForConditionalGeneration)
|
||||
model.config = SimpleNamespace(min_token_text_ratio=2.0, max_token_text_ratio=20.0)
|
||||
|
||||
update = model.postprocess(
|
||||
torch.randn(2, 4),
|
||||
multimodal_outputs={
|
||||
"prompt_speech_token_len": [torch.tensor([13])],
|
||||
"glm_tts_text_token_len": [torch.tensor([6])],
|
||||
},
|
||||
)
|
||||
|
||||
assert update["prompt_speech_token_len"].item() == 13
|
||||
assert update["glm_tts_text_token_len"].item() == 6
|
||||
@@ -0,0 +1,278 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.engine.serialization import serialize_additional_information
|
||||
from vllm_omni.model_executor.models.glm_tts.glm_tts_dit_wrapper import (
|
||||
split_request_ids,
|
||||
valid_speech_tokens,
|
||||
)
|
||||
from vllm_omni.model_executor.stage_input_processors.glm_tts import (
|
||||
ar_to_dit,
|
||||
ar_to_dit_async_chunk,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _Request:
|
||||
external_req_id = "req-1"
|
||||
additional_information = None
|
||||
|
||||
def is_finished(self) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def _transfer_manager() -> SimpleNamespace:
|
||||
connector = SimpleNamespace(
|
||||
config={
|
||||
"extra": {
|
||||
"codec_chunk_frames": 25,
|
||||
"codec_left_context_frames": 25,
|
||||
}
|
||||
}
|
||||
)
|
||||
return SimpleNamespace(connector=connector)
|
||||
|
||||
|
||||
def _progressive_transfer_manager() -> SimpleNamespace:
|
||||
connector = SimpleNamespace(
|
||||
config={
|
||||
"extra": {
|
||||
"codec_chunk_frames": [25, 50, 200],
|
||||
"codec_left_context_frames": 25,
|
||||
}
|
||||
}
|
||||
)
|
||||
return SimpleNamespace(connector=connector)
|
||||
|
||||
|
||||
def _source_output(mm: dict, request_id: str = "req-0") -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
request_id=request_id,
|
||||
outputs=[
|
||||
SimpleNamespace(
|
||||
multimodal_output=mm,
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _audio_codes(payload) -> list[int]:
|
||||
assert payload.codes is not None
|
||||
assert payload.codes.audio is not None
|
||||
return payload.codes.audio.tolist()
|
||||
|
||||
|
||||
def test_ar_to_dit_filters_invalid_tokens_and_preserves_conditioning() -> None:
|
||||
prompt_token = torch.tensor([[1, 2, 3]])
|
||||
prompt_feat = torch.randn(1, 4, 80)
|
||||
embedding = torch.randn(1, 192)
|
||||
source_outputs = [
|
||||
_source_output(
|
||||
{
|
||||
"speech_tokens": torch.tensor([[-1, 0, 7, -1, 32767]]),
|
||||
"prompt_token": prompt_token,
|
||||
"prompt_feat": prompt_feat,
|
||||
"embedding": embedding,
|
||||
}
|
||||
)
|
||||
]
|
||||
|
||||
outputs = ar_to_dit(source_outputs)
|
||||
|
||||
assert len(outputs) == 1
|
||||
assert outputs[0]["prompt_token_ids"] == [0, 7, 32767]
|
||||
additional_info = outputs[0]["additional_information"]
|
||||
assert additional_info["speech_tokens"] == [0, 7, 32767]
|
||||
# Voice clone data stored in nested embed dict (matches async_chunk path)
|
||||
embed = additional_info["embed"]
|
||||
assert embed["speech_token"] is prompt_token
|
||||
assert embed["speech_feat"] is prompt_feat
|
||||
assert embed["embedding"] is embedding
|
||||
# Meta dict enables block-causal attention in DiT wrapper
|
||||
meta = additional_info["meta"]
|
||||
assert meta["req_id"] == ["req-0"]
|
||||
assert meta["left_context_size"] == 0
|
||||
assert meta["stream_finished"] is True
|
||||
private = additional_info["kv_metadata"]["glm_tts"]
|
||||
assert private["block_pattern"] == [3] # len([0, 7, 32767])
|
||||
assert private["chunk_sizes_history"] == [3]
|
||||
|
||||
|
||||
def test_dit_explicit_empty_speech_tokens_override_scheduler_placeholder() -> None:
|
||||
placeholder = torch.tensor([0], dtype=torch.long)
|
||||
|
||||
token = valid_speech_tokens([], device=placeholder.device, fallback=placeholder)
|
||||
|
||||
assert token.numel() == 0
|
||||
|
||||
|
||||
def test_ar_to_dit_async_chunk_restores_serialized_conditioning_payload() -> None:
|
||||
transfer = _transfer_manager()
|
||||
request = _Request()
|
||||
prompt_token = torch.tensor([[1, 2, 3]])
|
||||
prompt_feat = torch.randn(1, 4, 80)
|
||||
embedding = torch.randn(1, 192)
|
||||
request.additional_information = serialize_additional_information(
|
||||
{
|
||||
"prompt_speech_token": prompt_token,
|
||||
"prompt_feat": prompt_feat,
|
||||
"embedding": embedding,
|
||||
}
|
||||
)
|
||||
|
||||
payload = ar_to_dit_async_chunk(
|
||||
transfer,
|
||||
{"speech_tokens": torch.tensor([[0]])},
|
||||
request,
|
||||
is_finished=False,
|
||||
)
|
||||
|
||||
assert payload is None
|
||||
state = transfer.request_payload["req-1"]["_glm_tts_async_state"]
|
||||
prompt_payload = state["prompt_payload"]
|
||||
assert torch.equal(prompt_payload["prompt_speech_token"], prompt_token)
|
||||
assert torch.equal(prompt_payload["prompt_feat"], prompt_feat)
|
||||
assert torch.equal(prompt_payload["embedding"], embedding)
|
||||
|
||||
|
||||
def test_split_request_ids_uses_seq_token_counts() -> None:
|
||||
ids = torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.long)
|
||||
|
||||
split = split_request_ids(ids, seq_token_counts=[2, 3, 1])
|
||||
|
||||
assert [part.tolist() for part in split] == [[1, 2], [3, 4, 5], [6]]
|
||||
|
||||
|
||||
def test_ar_to_dit_async_chunk_emits_cumulative_prefixes() -> None:
|
||||
transfer = _transfer_manager()
|
||||
request = _Request()
|
||||
payloads = []
|
||||
|
||||
for token in range(75):
|
||||
payload = ar_to_dit_async_chunk(
|
||||
transfer,
|
||||
{"speech_tokens": torch.tensor([[token]])},
|
||||
request,
|
||||
is_finished=False,
|
||||
)
|
||||
if payload is not None:
|
||||
payloads.append(payload)
|
||||
|
||||
assert len(payloads) == 3
|
||||
assert _audio_codes(payloads[0]) == list(range(25))
|
||||
assert payloads[0].meta.left_context_size == 0
|
||||
assert payloads[0].kv_metadata["glm_tts"]["block_pattern"] == [25]
|
||||
assert _audio_codes(payloads[1]) == list(range(50))
|
||||
assert payloads[1].meta.left_context_size == 25
|
||||
assert payloads[1].kv_metadata["glm_tts"]["block_pattern"] == [25]
|
||||
assert _audio_codes(payloads[2]) == list(range(75))
|
||||
assert payloads[2].meta.left_context_size == 50
|
||||
assert payloads[2].kv_metadata["glm_tts"]["block_pattern"] == [25]
|
||||
|
||||
|
||||
def test_ar_to_dit_async_chunk_keeps_progressive_block_pattern_fixed() -> None:
|
||||
transfer = _progressive_transfer_manager()
|
||||
request = _Request()
|
||||
payloads = []
|
||||
|
||||
for token in range(30):
|
||||
payload = ar_to_dit_async_chunk(
|
||||
transfer,
|
||||
{"speech_tokens": torch.tensor([[token]])},
|
||||
request,
|
||||
is_finished=False,
|
||||
)
|
||||
if payload is not None:
|
||||
payloads.append(payload)
|
||||
|
||||
final_payload = ar_to_dit_async_chunk(
|
||||
transfer,
|
||||
None,
|
||||
request,
|
||||
is_finished=True,
|
||||
)
|
||||
|
||||
assert len(payloads) == 1
|
||||
assert payloads[0].kv_metadata["glm_tts"]["chunk_sizes_history"] == [25]
|
||||
assert payloads[0].kv_metadata["glm_tts"]["block_pattern"] == [25, 50, 200]
|
||||
assert final_payload is not None
|
||||
assert final_payload.kv_metadata["glm_tts"]["chunk_sizes_history"] == [25, 5]
|
||||
assert final_payload.kv_metadata["glm_tts"]["block_pattern"] == [25, 50, 200]
|
||||
|
||||
|
||||
def test_ar_to_dit_async_chunk_terminal_boundary_is_metadata_only() -> None:
|
||||
transfer = _transfer_manager()
|
||||
request = _Request()
|
||||
payloads = []
|
||||
|
||||
for token in range(50):
|
||||
payload = ar_to_dit_async_chunk(
|
||||
transfer,
|
||||
{"speech_tokens": torch.tensor([[token]])},
|
||||
request,
|
||||
is_finished=False,
|
||||
)
|
||||
if payload is not None:
|
||||
payloads.append(payload)
|
||||
|
||||
final_payload = ar_to_dit_async_chunk(
|
||||
transfer,
|
||||
None,
|
||||
request,
|
||||
is_finished=True,
|
||||
)
|
||||
|
||||
assert len(payloads) == 2
|
||||
assert final_payload is not None
|
||||
assert _audio_codes(final_payload) == []
|
||||
assert final_payload.meta.left_context_size == 50
|
||||
assert final_payload.meta.req_id == ["req-1"]
|
||||
assert bool(final_payload.meta.stream_finished.item())
|
||||
|
||||
|
||||
def test_ar_to_dit_async_chunk_finished_flushes_partial_unsent_tail() -> None:
|
||||
transfer = _transfer_manager()
|
||||
request = _Request()
|
||||
payloads = []
|
||||
|
||||
for token in range(30):
|
||||
payload = ar_to_dit_async_chunk(
|
||||
transfer,
|
||||
{"speech_tokens": torch.tensor([[token]])},
|
||||
request,
|
||||
is_finished=False,
|
||||
)
|
||||
if payload is not None:
|
||||
payloads.append(payload)
|
||||
|
||||
final_payload = ar_to_dit_async_chunk(
|
||||
transfer,
|
||||
None,
|
||||
request,
|
||||
is_finished=True,
|
||||
)
|
||||
|
||||
assert len(payloads) == 1
|
||||
assert _audio_codes(payloads[0]) == list(range(25))
|
||||
assert payloads[0].meta.left_context_size == 0
|
||||
assert final_payload is not None
|
||||
assert _audio_codes(final_payload) == list(range(30))
|
||||
assert final_payload.meta.left_context_size == 25
|
||||
assert final_payload.meta.req_id == ["req-1"]
|
||||
assert bool(final_payload.meta.finished.item())
|
||||
assert bool(final_payload.meta.stream_finished.item())
|
||||
assert (
|
||||
ar_to_dit_async_chunk(
|
||||
transfer,
|
||||
None,
|
||||
request,
|
||||
is_finished=True,
|
||||
)
|
||||
is None
|
||||
)
|
||||
Reference in new issue
Block a user