SolitaryThinker
4e331eb7ff
Mechanical conversion of every relative import under v2/ to an absolute v2.* path (from .x / ..x / ...x -> from v2.<pkg>.x) so imports are unambiguous, grep-able, and stable when code is copied/moved between entrypoints (VideoGenerator, CLI, server). Surgical prefix-only rewrite: only the 'from <dots><module>' prefix changed — import names, parentheses, multi-line formatting, comments, and ordering are byte-for-byte preserved (no collapsing, no reorder, no unrelated reformatting). - 471 imports across 125 files; v2/tests/ was already absolute (untouched). - Validated: all 125 files compile, every 'from v2.* import' target resolves to a real module/package, zero relative imports remain (full sweep), CPU mini suite green (202 passed, 2 skipped).
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 v2._enums import ExecutionProfile, WorkUnitKind
|
|
from v2.loop.contracts import (
|
|
Done,
|
|
LoopResult,
|
|
LoopState,
|
|
ResourceRequest,
|
|
ShapeSignature,
|
|
StepResult,
|
|
WorkPlan,
|
|
)
|
|
from v2.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)
|