mstar.model.kokoro.submodules#
The Kokoro synthesis node: one sentence chunk per loop iteration, batched across requests, with CUDA graphs per length bucket.
process_prompt hands every request all of its chunks at once (padded
[n_chunks, T] phoneme ids, lengths, one style row per chunk, speed). Each
iteration of the chunk_loop synthesizes chunk i for every request in
the batch and emits its PCM to the client; check_stop ends a request’s
loop after its last chunk.
The forward has one host read, the frame count, between two static halves. Each half is captured as a piecewise CUDA graph per padded length bucket:
text_T<n>: phonemes -> prosody states, text features, durations;frames_F<n>: frame-aligned features -> waveform.
Rows of a batch are grouped by frame bucket before the decoder so a short
sentence is not padded to a long one. A batch larger than the biggest
captured one is cut into captured-size slices; a chunk longer than every
frame bucket runs the same code eagerly, at most max_batch_frames frames
per call, so no call needs more memory than a capture did.
Functions
|
|
|
Row indices grouped by frame bucket, ascending; rows that fit no bucket form a final |
|
Smallest bucket that holds |
|
Classes
|
Phonemes + style -> 24 kHz PCM16 for one chunk of every request. |
- class mstar.model.kokoro.submodules.KokoroSynthSubmodule(model, config)[source]#
Bases:
NodeSubmodulePhonemes + style -> 24 kHz PCM16 for one chunk of every request.
- Parameters:
model (KokoroTTS)
config (KokoroModelConfig)
- can_batch(batch, model_inputs)[source]#
- Parameters:
batch (ExecutingBatch)
model_inputs (list[NodeInputs])
- Return type:
- check_stop(request_id, request_info, outputs)[source]#
Return the set of dynamic-loop names that should stop after this step.
Runs on the worker’s slow-postprocess path after
execute_batchreturns — never insideexecute_batch. Allowed to read tensor values (.item()/.cpu()) because by this point the GPU thread is no longer blocked by it.Stops returned here are deferred by one step: they apply to the worker’s next iter’s fast postprocess. The current in-flight step (already submitted under the assumption that the rid continues) will run for that rid and its output discarded — the standard 1-wasted-step cost for any stop signal.
Default: no stops.
- forward(graph_walk, engine_inputs, input_ids, lengths, style, speed, **kwargs)[source]#
Pure tensor → NameToTensorList computation. Compilable + CUDA-graphable.
- forward_batched(graph_walk, engine_inputs, input_ids, lengths, style, speed, **kwargs)[source]#
Batched form of
forward: a multi-request batch’s outputs, per request or as aBatchedModelOutput. Override whencan_batchreturns True.
- frame_capture_sizes(bucket)[source]#
Batch sizes captured for a frame bucket:
bs * bucket <= max_batch_frames.
- get_piecewise_cuda_graph_configs(device, autocast_dtype, tp_world_size=1, **kwargs)[source]#
One region per text bucket and per frame bucket. Kokoro holds no engine resources, so the regions declare no step.
- prepare_inputs(graph_walk, fwd_info, inputs, **kwargs)[source]#
Select this iteration’s chunk;
Noneonce the request has none left.- Parameters:
- Return type:
NodeInputs | None
- preprocess(graph_walk, engine_inputs, inputs)[source]#
Pad the chunks of the batch to a common phoneme length.
- Parameters:
graph_walk (str)
engine_inputs (ModelInputsFromEngine)
inputs (list[NodeInputs])
- Return type:
- mstar.model.kokoro.submodules.group_by_bucket(sizes, buckets, policy='bucket')[source]#
Row indices grouped by frame bucket, ascending; rows that fit no bucket form a final
(None, rows)group.policy="single"puts every row that fits some bucket into the largest needed bucket instead: one replay per step at the cost of padding.