mstar.model.omnivoice.submodules

mstar.model.omnivoice.submodules#

Node submodules for OmniVoice: reference encode, unmask loop, code2wav.

The unmask node is the interesting one. Unlike wan22’s DiT — the only other diffusion node in the tree, which caps its batch at one because each request’s latent grid has its own geometry — OmniVoice’s canvases differ only in length, so several requests genuinely share a step. That is the whole reason to serve this model on M*: the reference’s own server runs one request at a time.

A step packs every request’s two CFG documents end to end into one ragged row and runs the reference’s own fused-kernel forward over it, so the per-step kernel path is the fastest one upstream ships – and the batch spans concurrent requests, which theirs cannot.

Classes

OmniVoiceBackboneSubmodule(backbone, config)

One unmask iteration for every request in the batch.

OmniVoiceCode2WavSubmodule(codec, config)

Decode the finished canvas and apply the reference's output shaping.

OmniVoiceRefEncoderSubmodule(codec, config)

Higgs-Audio-v2 encode of the cloning reference.
class mstar.model.omnivoice.submodules.OmniVoiceBackboneSubmodule(backbone, config)[source]#

Bases: NodeSubmodule

One unmask iteration for every request in the batch.

Loop-carried edges, per request:

audio_tokens [C, T] int64, the live canvas — all MASK at iteration 0 step_index [1] int64, the 0-based iteration counter

The prefix (prefix_ids / prefix_audio_mask) is fixed for the whole loop and arrives persisted from the request seam rather than riding the loop, since rebuilding it per step would re-tokenize the text 32 times.

The reveal schedule is recomputed each step from target_len, num_step and t_shift rather than shipped: it is a few dozen floats and deriving it is cheaper than moving it, the same call wan22 makes with its sigma tables.

Parameters:
can_batch(batch, model_inputs)[source]#

Cap the step by packed tokens, not by row count.

Cost per step is sum(doc_lens) through a bidirectional attention plus a head GEMM over 2 * sum(target_len); eight 30-second requests and eight 2-second ones are two very different steps. Row count alone would let the first case allocate a logits tensor several hundred MB wide.

The cap is deliberately conservative and has room now that the float32 upcast is per request rather than over the whole batch, but raising it is a measurement, not a guess, so it stays where it was.

Parameters:

model_inputs (list[NodeInputs])

Return type:

bool

check_stop(request_id, request_info, outputs)[source]#

End the loop after exactly this request’s num_step iterations.

While iteration k is being postprocessed the counter still reads k, so for N steps the stop must fire at k == N - 1. >= keeps it firing if async scheduling has already advanced the deferred count.

Parameters:
Return type:

set[str]

disable_autocast: bool = True#
disable_torch_compile: bool = True#
forward(graph_walk, engine_inputs, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:
Return type:

dict[str, list[Tensor]]

forward_batched(graph_walk, engine_inputs, items=None, canvas=None, **kwargs)[source]#

One packed forward over the whole batch.

Nothing request-shaped happens here: the canvas arrives built from preprocess and the CFG branches leave as a per-request pair of logit blocks, so the body is a single backbone call.

The float32 upcast happens per slice rather than on the whole packed tensor. Scoring needs float32, but casting first materialises a float32 copy of every request’s logits at once, and that peak is what max_packed_tokens has to be set low enough to survive. Casting after the slice keeps the peak at one batch in the backbone’s dtype plus one request in float32.

Parameters:
Return type:

dict[str, dict[str, list[Tensor]]]

max_batch_size(graph_walk)[source]#
Parameters:

graph_walk (str)

postprocess(request_id, request_info, outputs, inputs=None, **kwargs)[source]#

Turn this request’s logits into its next canvas.

Scoring, the reveal schedule and the in-place write are per-request and data-dependent, which is what keeps them out of forward_batched. The iteration index is the engine’s loop counter (see _loop_step); everything else comes from step metadata.

Parameters:
prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)[source]#
Parameters:
Return type:

NodeInputs | None

preprocess(graph_walk, engine_inputs, inputs)[source]#

Build the step’s packed canvas.

Canvases differ in length per request, so there is no padded collation here; the batch is one packed sequence with a document boundary per CFG branch, and engine_inputs.per_request_info carries the per-request guidance scale that decides each item’s layout.

Parameters:
Return type:

dict

class mstar.model.omnivoice.submodules.OmniVoiceCode2WavSubmodule(codec, config)[source]#

Bases: _SingleRequestMixin, NodeSubmodule

Decode the finished canvas and apply the reference’s output shaping.

Emits audio_output, 1-D int16 PCM at the codec’s sample rate: the speech handler concatenates node output straight into a WAV container, so the conversion belongs here, as it does for orpheus and qwen3_tts.

A canvas that still holds a MASK cell means the loop ended early, which would decode to noise — that is a fault, not a degraded result, so it raises rather than emitting.

Parameters:
disable_autocast: bool = True#
disable_torch_compile: bool = True#
forward(graph_walk, engine_inputs, audio_tokens, ref_rms=None, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:
Return type:

dict[str, list[Tensor]]

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)[source]#
Parameters:
Return type:

NodeInputs

class mstar.model.omnivoice.submodules.OmniVoiceRefEncoderSubmodule(codec, config)[source]#

Bases: _SingleRequestMixin, NodeSubmodule

Higgs-Audio-v2 encode of the cloning reference.

Consumes ref_audio_inputs — one mono waveform already resampled to the codec’s rate by the request seam — and emits ref_audio_tokens [C, T]. The edge is marked persist so the tokens survive from the encode_reference walk into the generation walk of the same request; the encode runs once per request, not once per voice. Caching a voice across requests, the reference’s VoiceClonePrompt, would need a store the engine does not have today.

Parameters:
disable_autocast: bool = True#
disable_torch_compile: bool = True#
forward(graph_walk, engine_inputs, waveform, **kwargs)[source]#

Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.

Parameters:
Return type:

dict[str, list[Tensor]]

prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)[source]#
Parameters:
Return type:

NodeInputs