504698dec8
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com> Co-authored-by: Chenguang ZHENG <645327136@qq.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
import torch
|
|
|
|
from vllm_omni.utils.mm_outputs import partition_flat_payload, partition_payload_list
|
|
|
|
|
|
def test_partition_thinker_latent_payload():
|
|
payload = {
|
|
"hidden_states.layer_0": torch.zeros(2, 4),
|
|
"hidden_states.layer_24": torch.zeros(2, 8),
|
|
"embed.tts_bos": [torch.zeros(1, 1, 4)],
|
|
}
|
|
inter, client = partition_flat_payload(payload)
|
|
assert inter == payload
|
|
assert client == {}
|
|
|
|
|
|
def test_partition_talker_intermediate_codes():
|
|
payload = {
|
|
"codes.audio": torch.zeros(3, 2),
|
|
"hidden": torch.zeros(3, 16),
|
|
}
|
|
inter, client = partition_flat_payload(payload)
|
|
assert inter == payload
|
|
assert client == {}
|
|
|
|
|
|
def test_partition_code2wav_client_audio():
|
|
payload = {
|
|
"model_outputs": torch.zeros(1, 2400),
|
|
"sr": torch.tensor(24000, dtype=torch.int32),
|
|
}
|
|
inter, client = partition_flat_payload(payload)
|
|
assert inter == {}
|
|
assert client == payload
|
|
|
|
|
|
def test_partition_payload_list_preserves_request_alignment():
|
|
payloads = [
|
|
{"hidden_states.layer_0": torch.zeros(1, 2)},
|
|
{"model_outputs": torch.zeros(1, 10)},
|
|
]
|
|
inter_list, client_list = partition_payload_list(payloads)
|
|
assert inter_list == [payloads[0], None]
|
|
assert client_list == [None, payloads[1]]
|