SolitaryThinker
c7e0a8e894
Ports vllm-omni's canonical qwen2_5_omni omni-speech cascade as a v2 card: a third weight-sharing topology — three disjoint experts (thinker, talker, vocoder) on three loop types (ar_decode → ar_decode → audio_decode) in one request, with chained cross-stage conditioning and streaming codec→waveform. vllm-omni runs these as three opaque request-scheduled stages; v2 makes every thinker token, talker token, and vocoder chunk a runtime-visible WorkUnit. - models/backend.py: ToyTalker (a genuinely distinct AR expert, weight-salted) + ToyVocoder (streaming code2wav: codec tokens → waveform chunks). - models/omni/vocoder_loop.py: VocoderLoop filling the pre-declared LoopKind.AUDIO_DECODE / WorkUnitKind.AUDIO_CHUNK slot. - models/omni/ar_loop.py: ARDecodeLoop gains a configurable prompt_slot so two chained AR loops don't collide on the prefill slot (thinker vs talker). - models/qwen_omni/: card (3 experts/3 loops) + program (tokenize → thinker → emit_text → thinker→talker full-payload hand-off → talker → talker→vocoder → vocoder → emit_audio). Cross-stage hand-offs are explicit Program nodes, the model-native form of vllm-omni's custom_process_input_func. - _enums.py: Capability.TEXT_TO_SPEECH. - tests/test_thinker_talker.py: 6 tests incl. three-loop interleave parity, cascade conditioning, AUDIO_CHUNK streaming. Full suite 97 passed. designv4.md: §2.3 topology table extended to four topologies; new §9.5 on the cascade; reference-synthesis + package layout updated.
82 lines
3.7 KiB
Python
82 lines
3.7 KiB
Python
"""VocoderLoop — streaming code2wav (design_v3 §4.3 ``LoopKind.AUDIO_DECODE``; Qwen-Omni stage 2).
|
|
|
|
The third loop type in the thinker→talker→vocoder cascade. It consumes the talker's speech-codec
|
|
tokens in fixed-size chunks; **each chunk is one ``AUDIO_CHUNK`` WorkUnit** that synthesizes a
|
|
waveform segment and streams it. This is the differentiation vllm-omni's own docs name: their
|
|
code2wav is one opaque ``LLM_GENERATION`` stage the scheduler never sees inside — here every chunk
|
|
is a runtime-visible, step-scheduled, cancellable WorkUnit, on the *same* driven-loop contract the
|
|
engine uses for denoise steps and AR tokens. All per-request state lives in ``LoopState``.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from ..._enums import ExecutionProfile, WorkUnitKind
|
|
from ...loop.contracts import (
|
|
Done,
|
|
LoopResult,
|
|
LoopState,
|
|
ResourceRequest,
|
|
ShapeSignature,
|
|
StepResult,
|
|
WorkPlan,
|
|
)
|
|
from ...request.streams import StreamChunk
|
|
|
|
|
|
class VocoderLoop:
|
|
def __init__(self, *, loop_id, vocoder_id="vocoder", cost, chunk_tokens=2, speech_slot="speech_tokens"):
|
|
self.loop_id = loop_id
|
|
self.vocoder_id = vocoder_id
|
|
self.cost = cost
|
|
self.chunk_tokens = chunk_tokens
|
|
self.speech_slot = speech_slot
|
|
|
|
def init(self, req, model, ctx) -> LoopState:
|
|
toks = list(ctx.slots.get(self.speech_slot) or [1])
|
|
st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id,
|
|
request_id=req.request_id, profile=ctx.profile)
|
|
st.scratch["tokens"] = toks
|
|
st.scratch["chunks"] = [toks[i:i + self.chunk_tokens]
|
|
for i in range(0, len(toks), self.chunk_tokens)] or [[1]]
|
|
st.scratch["wave"] = []
|
|
st.scratch["stream_audio"] = bool(req.outputs.stream.get("audio"))
|
|
return st
|
|
|
|
def next(self, st: LoopState):
|
|
i = st.step_idx
|
|
chunks = st.scratch["chunks"]
|
|
if i >= len(chunks):
|
|
return Done()
|
|
chunk = chunks[i]
|
|
vid = self.vocoder_id
|
|
|
|
def run(model, override=None):
|
|
seg = model.component(vid).synthesize(chunk) # codec tokens → waveform segment
|
|
return StepResult(output={"segment": np.asarray(seg, dtype="float32")})
|
|
|
|
emits = []
|
|
if st.scratch.get("stream_audio"):
|
|
emits.append(StreamChunk(stream_id=st.request_id, modality="audio", seq=i, preview=False))
|
|
return WorkPlan(
|
|
loop_id=self.loop_id, instance_id=st.instance_id, kind=WorkUnitKind.AUDIO_CHUNK,
|
|
shape_sig=ShapeSignature(WorkUnitKind.AUDIO_CHUNK, dims=(len(chunk),)),
|
|
resources=ResourceRequest(compute_seconds=self.cost.predict(len(chunk)),
|
|
resident_bytes=8 * len(st.scratch["tokens"]),
|
|
peak_activation_bytes=64),
|
|
payload={"chunk": i}, run=run, emits=emits, label=f"vocoder.chunk{i}")
|
|
|
|
def advance(self, st: LoopState, result: StepResult) -> LoopState:
|
|
st.scratch["wave"].append(np.asarray(result.output["segment"], dtype="float32"))
|
|
st.step_idx += 1
|
|
if st.profile == ExecutionProfile.ROLLOUT:
|
|
st.trajectory.append({"chunk": st.step_idx, "samples": int(result.output["segment"].size)})
|
|
return st
|
|
|
|
def finalize(self, st: LoopState) -> LoopResult:
|
|
wave = np.concatenate(st.scratch["wave"]) if st.scratch["wave"] else np.zeros(0, dtype="float32")
|
|
return LoopResult(outputs={"samples": wave.astype("float32")},
|
|
metrics={"audio_chunks": float(st.step_idx),
|
|
"audio_samples": float(wave.size)},
|
|
behavior=st.trajectory or None)
|