From 380e8022ddd562ec4da9c23cb07f1a89e4cfe958 Mon Sep 17 00:00:00 2001
From: tzhouam <tzhouam@connect.ust.hk>
Date: Wed, 9 Sep 2026 16:31:56 +0800
Subject: [PATCH] [Edge][WP5] Normalize DELTA-mode codec snapshots (list of
 cumulative code matrices) before row alignment

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 .../openai_api/test_codec_stream_protocol.py  | 23 +++++++++++++++++
 .../entrypoints/openai/serving_speech.py      | 25 ++++++++++++++++++-
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/tests/entrypoints/openai_api/test_codec_stream_protocol.py b/tests/entrypoints/openai_api/test_codec_stream_protocol.py
index e2029fe9e..f4a11522f 100644
--- a/tests/entrypoints/openai_api/test_codec_stream_protocol.py
+++ b/tests/entrypoints/openai_api/test_codec_stream_protocol.py
@@ -199,3 +199,26 @@ def test_talker_only_deployment_detection():
     assert probe._is_talker_only_deployment() is False
     probe._tts_stage = None
     assert probe._is_talker_only_deployment() is False
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_codec_snapshot_list_is_normalized():
+    import numpy as np
+    import torch
+
+    from vllm_omni.entrypoints.openai.serving_speech import _coerce_codec_snapshots
+
+    assert _coerce_codec_snapshots(None) is None
+    t = torch.arange(32).reshape(2, 16)
+    assert _coerce_codec_snapshots(t) is t
+    # DELTA-mode list of cumulative snapshots -> last one
+    snaps = [torch.arange(16).reshape(1, 16), torch.arange(32).reshape(2, 16), torch.arange(48).reshape(3, 16)]
+    out = _coerce_codec_snapshots(snaps)
+    assert out.shape == (3, 16) and int(out[2, 15]) == 47
+    # list of single-row deltas -> concatenation
+    deltas = [torch.full((1, 16), i) for i in range(4)]
+    out = _coerce_codec_snapshots(deltas)
+    assert out.shape == (4, 16) and list(out[:, 0]) == [0, 1, 2, 3]
+    assert _coerce_codec_snapshots([None]) is None
+    assert isinstance(_coerce_codec_snapshots([np.zeros(16)]), np.ndarray)
diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py
index 1d867771b..cdaa878ee 100644
--- a/vllm_omni/entrypoints/openai/serving_speech.py
+++ b/vllm_omni/entrypoints/openai/serving_speech.py
@@ -278,6 +278,29 @@ def _conditioning_cache_salt(request, tts_params: dict | None = None) -> str:
     return h.hexdigest()[:32]
 
 
+def _coerce_codec_snapshots(codes: Any) -> Any:
+    """Normalize the per-step ``codes.audio`` payload to one ``[rows, codebooks]`` array.
+
+    In DELTA output mode the output processor accumulates the talker's
+    per-step payloads as a list (deferred concatenation). The talker publishes
+    the *cumulative* code matrix each step, so the last snapshot is the current
+    state; a list of single-row entries (a true per-step delta producer) is
+    concatenated instead.
+    """
+    if codes is None or not isinstance(codes, (list, tuple)):
+        return codes
+    items = [c for c in codes if c is not None]
+    if not items:
+        return None
+    import numpy as np
+
+    arrays = [c.detach().cpu().numpy() if hasattr(c, "detach") else np.asarray(c) for c in items]
+    arrays = [a.reshape(1, -1) if a.ndim == 1 else a for a in arrays]
+    if len(arrays) > 1 and all(a.ndim == 2 and a.shape[0] == 1 for a in arrays):
+        return np.concatenate(arrays, axis=0)
+    return arrays[-1]
+
+
 class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
     _diffusion_mode: bool = False
     _media_connector: MediaConnector | None = None
@@ -3280,7 +3303,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         if codes is None:
             return None
         if hasattr(codes, "get"):
-            return codes.get("audio")
+            return _coerce_codec_snapshots(codes.get("audio"))
         return None
 
     async def _generate_codec_chunks(