[Model] Add Gepard-1.0 native-AR FSQ/NanoCodec TTS (offline inference) (#5666)

Signed-off-by: Mengjie Zhao <zmj0129@gmail.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/model_executor/models/test_gepard_window.py

@@ -0,0 +1,363 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""CPU coverage for Gepard's prompt layout, streaming-decode window and config.
+
+The e2e test needs a GPU, NeMo and the checkpoint, so it only runs nightly.
+These pin the parts that are pure logic and break silently: which frames each
+decode covers, which samples survive the lookback trim, the prompt layout, and
+that the packaged deploy config sits where the example looks for it.
+
+The codec is replaced by a shift-invariant stand-in whose output depends only
+on the codes, which makes "chunked emissions concatenate to one whole-history
+decode" an exact assertion — the contract lookback exists to provide.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn as nn
+
+import vllm_omni
+from vllm_omni.config.stage_config import load_deploy_config
+from vllm_omni.model_executor.models.gepard.configuration_gepard import GepardConfig
+from vllm_omni.model_executor.models.gepard.gepard_talker import (
+    CHUNK_FRAMES,
+    FIRST_CHUNK_FRAMES,
+    GepardTalkerForConditionalGeneration,
+    _GepardState,
+)
+from vllm_omni.model_executor.models.gepard.nanocodec import (
+    LOOKBACK_FRAMES,
+    NanoCodec,
+    apply_end_of_speech_tail,
+)
+from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+SAMPLES_PER_FRAME = 1024
+NUM_HEADS = 32
+SAMPLE_RATE = 22050
+# Text-token count at or above which the layout is left unrepeated, read from
+# the checkpoint defaults so the boundary cases follow the config.
+APPLY_BELOW = GepardConfig().text_repetition_apply_below
+
+
+class _ShiftInvariantCodec:
+    """Frame tagged ``v`` -> samples ``v*SPF .. v*SPF+SPF-1``, wherever it sits."""
+
+    def __init__(self) -> None:
+        self.windows: list[int] = []  # frame count of every decode call
+
+    def decode_from_codes(self, codes: torch.Tensor, codes_len: torch.Tensor):
+        assert codes.shape[0] == 1, "streaming decode is per-request"
+        assert codes.shape[1] == NUM_HEADS, f"expected {NUM_HEADS} code channels"
+        tags = codes[0, 0, :].to(torch.int64)  # (T,)
+        self.windows.append(int(tags.shape[0]))
+        offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64)
+        audio = (tags.unsqueeze(1) * SAMPLES_PER_FRAME + offsets).reshape(1, -1)
+        return audio.to(torch.float32), codes_len * SAMPLES_PER_FRAME
+
+
+def _make_codec() -> tuple[NanoCodec, _ShiftInvariantCodec]:
+    codec = NanoCodec(codec_id="test/fake-codec", sample_rate=SAMPLE_RATE)
+    fake = _ShiftInvariantCodec()
+    codec._codec = fake  # bypasses load(); NeMo is never imported
+    assert codec.is_loaded
+    return codec, fake
+
+
+def _make_talker(codec: NanoCodec) -> GepardTalkerForConditionalGeneration:
+    """A real talker with only the audio-path attributes populated.
+
+    ``__init__`` needs a VllmConfig and builds the backbone; the methods under
+    test touch only these four attributes.
+    """
+    talker = GepardTalkerForConditionalGeneration.__new__(GepardTalkerForConditionalGeneration)
+    nn.Module.__init__(talker)  # sets up _parameters/_buffers/_modules only
+    talker._codec = codec
+    talker._audio_queue = []
+    talker._active_states = {}
+    talker._deferred_cleanup_ids = set()
+    return talker
+
+
+def _frame(tag: int) -> torch.Tensor:
+    """One committed frame: 32 per-dimension codes, all carrying the tag."""
+    return torch.full((NUM_HEADS,), tag, dtype=torch.long)
+
+
+def _expected_samples(num_frames: int) -> torch.Tensor:
+    offsets = torch.arange(SAMPLES_PER_FRAME, dtype=torch.int64)
+    tags = torch.arange(num_frames, dtype=torch.int64).unsqueeze(1)
+    return (tags * SAMPLES_PER_FRAME + offsets).reshape(-1).to(torch.float32)
+
+
+def _drain(talker) -> torch.Tensor:
+    """Concatenate and clear the queued deltas, as make_omni_output does."""
+    if not talker._audio_queue:
+        return torch.empty(0)
+    out = torch.cat([a.reshape(-1) for _, a in talker._audio_queue])
+    talker._audio_queue.clear()
+    return out
+
+
+def _generate(talker, state, num_frames: int, *, stop_at_end: bool) -> None:
+    """Drive the commit loop the way forward() does, one frame per step."""
+    for tag in range(num_frames):
+        state.frames.append(_frame(tag))
+        state.frame_count += 1
+        state.past_first_step = True
+        talker._emit_audio(state, is_final=stop_at_end and tag == num_frames - 1)
+
+
+# Counts straddling both cadence thresholds, where an off-by-one hides.
+@pytest.mark.parametrize(
+    "num_frames",
+    [
+        1,
+        2,
+        FIRST_CHUNK_FRAMES - 1,
+        FIRST_CHUNK_FRAMES,
+        FIRST_CHUNK_FRAMES + 1,
+        CHUNK_FRAMES - 1,
+        CHUNK_FRAMES,
+        CHUNK_FRAMES + 1,
+        FIRST_CHUNK_FRAMES + CHUNK_FRAMES,
+        79,
+        200,
+    ],
+)
+def test_chunked_emissions_equal_a_whole_history_decode(num_frames: int) -> None:
+    codec, _fake = _make_codec()
+    talker = _make_talker(codec)
+    state = _GepardState(request_id="req-0")
+
+    _generate(talker, state, num_frames, stop_at_end=True)
+
+    torch.testing.assert_close(_drain(talker), _expected_samples(num_frames))
+    assert state.emitted_frames == num_frames
+
+
+def test_every_frame_is_emitted_exactly_once() -> None:
+    """The trim must not drop or duplicate samples at a chunk boundary."""
+    codec, _fake = _make_codec()
+    talker = _make_talker(codec)
+    state = _GepardState(request_id="req-0")
+    num_frames = FIRST_CHUNK_FRAMES + 2 * CHUNK_FRAMES + 3
+
+    _generate(talker, state, num_frames, stop_at_end=True)
+    audio = _drain(talker)
+
+    assert audio.numel() == num_frames * SAMPLES_PER_FRAME
+    # Tags are globally unique, so a duplicated or dropped frame shows up here.
+    assert len(torch.unique(audio)) == audio.numel()
+
+
+def test_decode_windows_carry_lookback_context() -> None:
+    """Every decode after the first must re-decode LOOKBACK_FRAMES of context."""
+    codec, fake = _make_codec()
+    talker = _make_talker(codec)
+    state = _GepardState(request_id="req-0")
+    num_frames = FIRST_CHUNK_FRAMES + 2 * CHUNK_FRAMES
+
+    _generate(talker, state, num_frames, stop_at_end=True)
+
+    assert len(fake.windows) >= 3, "expected a first chunk plus steady-state chunks"
+    assert fake.windows[0] == FIRST_CHUNK_FRAMES, "the first chunk has no left context"
+    for window in fake.windows[1:]:
+        assert window == CHUNK_FRAMES + LOOKBACK_FRAMES
+
+
+def test_codec_runs_once_per_chunk_not_once_per_frame() -> None:
+    codec, fake = _make_codec()
+    talker = _make_talker(codec)
+    state = _GepardState(request_id="req-0")
+    num_frames = FIRST_CHUNK_FRAMES + 3 * CHUNK_FRAMES
+
+    _generate(talker, state, num_frames, stop_at_end=True)
+
+    assert len(fake.windows) == 4
+    assert sum(fake.windows) < num_frames * 2  # far below one decode per frame
+
+
+def test_final_emit_drains_a_partial_chunk() -> None:
+    """A request can end mid-chunk (max_tokens, STOP); the frames committed
+    since the last cadence boundary must still reach the queue.
+
+    Arithmetic only: this drives ``_emit_audio`` directly, so it says nothing
+    about *whether* the engine reaches it on the request's last step. That
+    wiring lives in ``test_gepard_wiring.py`` — a bug there leaves this green.
+    """
+    codec, _fake = _make_codec()
+    talker = _make_talker(codec)
+    state = _GepardState(request_id="req-0")
+    num_frames = FIRST_CHUNK_FRAMES + CHUNK_FRAMES + 5
+
+    _generate(talker, state, num_frames, stop_at_end=False)
+    assert state.emitted_frames < num_frames, "a partial chunk should still be pending"
+
+    talker._emit_audio(state, is_final=True)
+
+    torch.testing.assert_close(_drain(talker), _expected_samples(num_frames))
+    assert state.emitted_frames == num_frames
+
+
+def test_emitting_with_nothing_pending_is_a_no_op() -> None:
+    codec, fake = _make_codec()
+    talker = _make_talker(codec)
+    state = _GepardState(request_id="req-0")
+
+    _generate(talker, state, 4, stop_at_end=True)
+    calls_after_flush = len(fake.windows)
+    talker._audio_queue.clear()
+
+    talker._emit_audio(state, is_final=True)
+
+    assert len(fake.windows) == calls_after_flush, "codec ran with no new frames"
+    assert talker._audio_queue == []
+
+
+def test_concurrent_requests_keep_separate_windows() -> None:
+    codec, _fake = _make_codec()
+    talker = _make_talker(codec)
+    states = [_GepardState(request_id=f"req-{i}") for i in range(2)]
+
+    for state in states:
+        _generate(talker, state, FIRST_CHUNK_FRAMES, stop_at_end=True)
+
+    by_req: dict[str, list[torch.Tensor]] = {}
+    for req_id, audio in talker._audio_queue:
+        by_req.setdefault(req_id, []).append(audio.reshape(-1))
+    assert set(by_req) == {"req-0", "req-1"}
+    for chunks in by_req.values():
+        torch.testing.assert_close(torch.cat(chunks), _expected_samples(FIRST_CHUNK_FRAMES))
+
+
+def test_end_of_speech_tail_fades_and_pads() -> None:
+    audio = torch.ones(SAMPLE_RATE)  # 1 s of full-scale signal
+    out = apply_end_of_speech_tail(audio, SAMPLE_RATE, fade_ms=10.0, silence_ms=20.0)
+
+    fade = int(SAMPLE_RATE * 10.0 / 1000.0)
+    pad = int(SAMPLE_RATE * 20.0 / 1000.0)
+    assert out.numel() == SAMPLE_RATE + pad
+    assert torch.all(out[-pad:] == 0.0), "trailing pad must be silent"
+    assert out[SAMPLE_RATE - fade - 1] == 1.0, "the fade must not reach back past fade_ms"
+    # Monotonically ramped to zero — the pad alone would leave the click.
+    ramp = out[SAMPLE_RATE - fade : SAMPLE_RATE]
+    assert torch.all(ramp[:-1] >= ramp[1:])
+    assert ramp[-1] == 0.0
+
+
+def test_end_of_speech_tail_is_a_no_op_by_default() -> None:
+    audio = torch.ones(128)
+    out = apply_end_of_speech_tail(audio, SAMPLE_RATE, fade_ms=0.0, silence_ms=0.0)
+    assert out is audio
+
+
+def _repeat_count(ids: list[int], cfg: GepardConfig) -> int:
+    """How many ``[SOT ... EOT]`` copies the layout carries."""
+    return ids.count(cfg.start_of_text)
+
+
+def test_gepard_prompt_layout() -> None:
+    """The prompt carries the speaker slots and exactly one SOS, at the end."""
+    cfg = GepardConfig()
+    ids = build_gepard_prompt_ids([1, 2, 3], config=cfg)
+
+    slots = [cfg.speaker_token_base + i for i in range(cfg.num_speaker_prefix)]
+    assert ids[: cfg.num_speaker_prefix] == slots
+    assert ids.count(cfg.start_of_speech) == 1, "only the canonical copy may carry SOS"
+    assert ids[-1] == cfg.start_of_speech, "the first frame is sampled from the SOS position"
+
+
+def test_short_text_repeats_to_reach_the_target_token_mass() -> None:
+    """A short text is repeated until its region holds ~target_text_tokens."""
+    cfg = GepardConfig()
+    text = [1, 2, 3, 4]  # target 16 / 4 -> 4 copies
+    ids = build_gepard_prompt_ids(text, config=cfg)
+
+    assert _repeat_count(ids, cfg) == 4
+    assert ids.count(cfg.start_of_speech) == 1, "only the last copy opens the speech region"
+    assert ids[-1] == cfg.start_of_speech
+
+
+@pytest.mark.parametrize(
+    ("num_text_tokens", "repeated"),
+    [
+        (APPLY_BELOW - 1, True),
+        (APPLY_BELOW, False),
+        (APPLY_BELOW + 1, False),
+    ],
+)
+def test_apply_below_is_the_repetition_boundary(num_text_tokens: int, repeated: bool) -> None:
+    """Both sides of the threshold: at apply_below the text is left alone.
+
+    The layout has to match training here or WER collapses, and the boundary
+    is where an off-by-one would be invisible in a single happy-path clip.
+    """
+    cfg = GepardConfig()
+    ids = build_gepard_prompt_ids(list(range(1, num_text_tokens + 1)), config=cfg)
+
+    assert (_repeat_count(ids, cfg) > 1) is repeated
+
+
+def test_max_repeats_caps_a_one_token_text() -> None:
+    """target/1 would ask for 16 copies; the cap keeps the prefill bounded."""
+    cfg = GepardConfig()
+    ids = build_gepard_prompt_ids([1], config=cfg)
+
+    assert _repeat_count(ids, cfg) == cfg.text_repetition_max_repeats
+
+
+def test_repetition_thresholds_are_read_from_the_config() -> None:
+    """The numbers come from the checkpoint sidecar, not from literals.
+
+    A checkpoint whose text_repetition block differs must change the layout;
+    duplicating the values in the builder would silently ignore it.
+    """
+    cfg = GepardConfig(text_repetition={"enabled": True, "target_text_tokens": 9, "apply_below": 8, "max_repeats": 2})
+    assert _repeat_count(build_gepard_prompt_ids([1, 2, 3], config=cfg), cfg) == 2  # ceil(9/3)=3, capped at 2
+    assert _repeat_count(build_gepard_prompt_ids(list(range(8)), config=cfg), cfg) == 1  # at apply_below
+
+
+def test_repetition_can_be_disabled() -> None:
+    cfg = GepardConfig(text_repetition={"enabled": False})
+    ids = build_gepard_prompt_ids([1], config=cfg)
+
+    assert _repeat_count(ids, cfg) == 1
+    assert ids[-1] == cfg.start_of_speech
+
+
+def test_empty_text_is_rejected() -> None:
+    """An empty text would still assemble a valid layout and be voiced.
+
+    ``[slots, SOT, EOT, SOS]`` parses, prefills and generates: the failure is
+    audible, not raised, which is the worst shape for a caller to debug.
+    """
+    with pytest.raises(ValueError, match="at least one text token"):
+        build_gepard_prompt_ids([], config=GepardConfig())
+
+
+def test_packaged_deploy_config_is_where_the_example_looks_for_it() -> None:
+    """The example resolves its default off ``vllm_omni.__file__``; tests reach
+    the same YAML through ``tests.helpers.stage_config``, so nothing covered the
+    packaged location until this.
+
+    ``pipeline: gepard`` is the load-bearing key: the checkpoint self-identifies
+    as qwen3_5_text, so without the pin the architectures fallback routes it to
+    the diffusion registry.
+
+    This asserts the location, not the example's constant — importing the
+    example would pull ``Omni`` and NeMo into a CPU test. Running the quick
+    start verbatim stays a release check.
+    """
+    path = Path(vllm_omni.__file__).resolve().parent / "deploy" / "gepard.yaml"
+    assert path.is_file(), f"packaged deploy config missing at {path}"
+
+    deploy = load_deploy_config(path)
+    assert deploy.pipeline == "gepard"
+    assert deploy.stages, "deploy config declares no stages"