Audio conversation pipeline

[Ssofja]

Description

This PR adds an end-to-end multi-speaker conversation data-generation pipeline for audio, wiring together four audio stages into a single runnable tutorial: LLM conversation generation → TTS synthesis → forced alignment → conversation merging.

It integrates three previously separate audio feature branches and adds the glue needed to run them as one pipeline:

• vLLMInference stage (nemo_curator/stages/audio/llm/vllm_inference.py): generates structured multi-turn conversations from topic prompts. It delegates to the existing nemo_curator.models.VLLMModel (already on main) for engine management/inference rather than creating vllm.LLM directly, and emits one AudioTask per conversation turn (shared conversation_id, speaker, utterance, overlap). Includes JSON-schema validation with batch retry and per-turn expansion.
• ChatterboxTTSStage (from the chatterbox_tts work): synthesizes per-turn audio with voice cloning.
• MFAAlignmentStage (from the montreal_forced_aligner work): Montreal Forced Aligner producing word-level TextGrid → RTTM/CTM.
• MergeConversationSDPStage (from the merging_audios work): merges per-turn audio into multi-speaker conversations (per-speaker WAVs, multichannel, mixed mono, combined RTTM/CTM/seglst).

Supporting changes:

• Registered MergeConversationSDPStage and vLLMInference in the audio package lazy exports (nemo_curator/stages/audio/__init__.py) so both are discoverable via nemo_curator.stages.audio without eagerly importing heavy optional deps.
• Added the tutorials/audio/data-generation/ tutorial: main.py (two-phase runner), pipeline.yaml (Hydra config), README.md, plus 2/3/4-speaker prompt templates and sample topic manifests. Registered the tutorial in tutorials/audio/README.md.
• Reuses the existing vllm optional extra and the audio_common extra (which already carries chatterbox-tts and praatio); no new dependency declarations were required.

The pipeline runs in two phases: Phase 1 (GPU) runs ManifestReader → vLLMInference → ChatterboxTTSStage → MFAAlignmentStage via the standard executor (one AudioTask per turn); Phase 2 (CPU) groups turns by conversation_id, runs MergeConversationSDPStage per conversation, and writes the final manifest.

Usage

Install the audio + vLLM extras and MFA models, then run the tutorial from the repo root:

# Install (GPU)
uv sync --extra audio_cuda12 --extra vllm && source .venv/bin/activate

# Download MFA models (one-time)
mfa model download acoustic english_us_arpa
mfa model download dictionary english_us_arpa
mfa model download g2p english_us_arpa

# Run the 2-speaker dialog pipeline
python tutorials/audio/data-generation/main.py \
    --config-path . \
    --config-name pipeline \
    input_manifest=tutorials/audio/data-generation/topics/topics_dialog.jsonl \
    output_dir=/data/conversations \
    reference_voices_dataset=/data/reference_voices \
    prompt_file=tutorials/audio/data-generation/prompts/dialog_prompt.yaml

Using the vLLMInference stage directly in a pipeline:

from nemo_curator.pipeline import Pipeline
from nemo_curator.stages.audio import (
    ChatterboxTTSStage,
    MFAAlignmentStage,
    MergeConversationSDPStage,
    vLLMInference,
)
from nemo_curator.stages.audio.common import ManifestReader

llm = vLLMInference(
    prompt_file="tutorials/audio/data-generation/prompts/dialog_prompt.yaml",
    model={"model": "Qwen/Qwen2.5-7B-Instruct", "max_model_len": 4096, "max_tokens": 4096},
)

pipeline = Pipeline(name="conversation_generation")
pipeline.add_stage(ManifestReader(manifest_path="topics/topics_dialog.jsonl"))
pipeline.add_stage(llm)
pipeline.add_stage(ChatterboxTTSStage(
    output_audio_dir="/data/out/audio",
    reference_voices_dataset="/data/reference_voices",
))
pipeline.add_stage(MFAAlignmentStage(output_dir="/data/out/alignment", text_key="utterance"))
# MergeConversationSDPStage requires all turns of a conversation in one batch;
# see tutorials/audio/data-generation/main.py for the two-phase runner.

Checklist

• I am familiar with the Contributing Guide.
• New or Existing tests cover these changes.
  • Existing stage unit tests pass (tests/stages/audio/tts, tests/stages/audio/alignment, tests/stages/audio/merging — 58 tests). The vLLMInference wrapper and the tutorial runner were verified via import/compile smoke checks; a full end-to-end run requires a GPU plus vLLM/Chatterbox/MFA models and was not run in CI.
• The documentation is up to date with these changes.
  • Added tutorials/audio/data-generation/README.md and registered the tutorial in tutorials/audio/README.md (tutorial selection, data availability, and system dependency tables).

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR wires four previously separate audio feature branches into a single end-to-end multi-speaker conversation generation pipeline: LLM conversation generation (vLLMInference) → TTS synthesis (ChatterboxTTSStage) → forced alignment (MFAAlignmentStage) → conversation merging (MergeConversationSDPStage). A two-phase tutorial runner handles the split between per-turn GPU stages and the conversation-batched CPU merge stage.

• vLLMInference: new stage that delegates to VLLMModel for vLLM engine management, applies JSON-schema validation with batch retry, and emits one AudioTask per conversation turn.
• ChatterboxTTSStage: synthesises per-turn audio with voice cloning from reference WAVs (wavs/ or MLS layout), supporting English and multilingual Chatterbox models.
• MFAAlignmentStage: runs mfa align in batch, converts TextGrids to RTTM/CTM, copies models to node-local storage to avoid NFS race conditions.
• MergeConversationSDPStage: merges per-turn audio into per-speaker WAVs, multichannel, mixed mono, combined RTTM/CTM, and a seglst JSON; shares a single _compute_timeline for timestamp consistency.

Confidence Score: 3/5

Two defects on the hot path will cause failures for any user who follows the documented API: an AttributeError when a pre-initialized VLLMModel is passed, and missing output fields propagated to downstream stages for empty utterances.

The vLLMInference stage documents that callers may pass a pre-configured VLLMModel instance, but _ensure_model() uses the private _llm field as the sole initialization check. When _llm is already set, setup() is skipped, self.tokenizer remains None, and the first get_entry_prompt() call raises AttributeError. Separately, ChatterboxTTSStage.process_batch re-emits the raw input task (without audio_filepath, duration, or reference_voice) when utterance text is empty, silently propagating an incomplete task that will break MFAAlignmentStage. Both defects affect the core data path and need to be addressed before the pipeline can be used reliably.

nemo_curator/stages/audio/llm/vllm_inference.py (tokenizer/ensure-model logic) and nemo_curator/stages/audio/tts/chatterbox_tts.py (empty-text skip path)

Important Files Changed

Filename	Overview
nemo_curator/stages/audio/llm/vllm_inference.py	New LLM conversation generation stage; tokenizer is never set when a pre-initialized VLLMModel is passed, causing AttributeError in process_batch. Retry count hard-coded. Accessing private _llm attribute crosses class boundaries.
nemo_curator/stages/audio/tts/chatterbox_tts.py	New TTS stage; skips tasks with empty text but emits the unmodified input task without required output fields (audio_filepath, duration, reference_voice), breaking downstream stages.
nemo_curator/stages/audio/alignment/mfa_alignment.py	New MFA batch alignment stage; _is_node_local_mfa_root() can return True for shared paths when copy_models_to_local=True but setup_on_node was skipped, risking deletion of shared command_history.yaml. Otherwise well-structured with fallback handling.
nemo_curator/stages/audio/merging/merge_conversation.py	New conversation merging stage; uses a shared _compute_timeline for RTTM+CTM consistency, handles missing audio gracefully. No blocking issues found.
tutorials/audio/data-generation/main.py	Two-phase tutorial runner; phase-2 stage selection uses hard-coded indices that silently break if stage count/order in pipeline.yaml changes.
nemo_curator/stages/audio/init.py	Registers new stages in the lazy-export table; straightforward and consistent with existing pattern.
pyproject.toml	Adds praatio>=6.0 and chatterbox-tts>=0.1.4 to audio_common extra; consistent with existing conditional platform markers.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant MR as ManifestReader
    participant LLM as vLLMInference
    participant TTS as ChatterboxTTSStage
    participant MFA as MFAAlignmentStage
    participant Merge as MergeConversationSDPStage
    participant MW as ManifestWriterStage

    note over MR,MFA: Phase 1 (GPU — standard executor)
    MR->>LLM: AudioTask(topic)
    LLM->>LLM: VLLMModel.generate() + JSON validate
    LLM-->>TTS: AudioTask×N (one per turn, shared conversation_id)
    TTS->>TTS: ChatterboxTTS.generate() per turn
    TTS-->>MFA: AudioTask + audio_filepath
    MFA->>MFA: mfa align subprocess (batch)
    MFA-->>Merge: AudioTask + rttm_filepath + ctm_filepath

    note over Merge,MW: Phase 2 (CPU — grouped by conversation_id)
    Merge->>Merge: group turns, _merge_audio_files()
    Merge->>Merge: _compute_timeline() shared RTTM+CTM offsets
    Merge->>Merge: _merge_rttm_files() + _merge_ctm_files()
    Merge-->>MW: AudioTask(mixed.wav, all.rttm, all.ctm, seglst)
    MW-->>MW: write merged_manifest.jsonl

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant MR as ManifestReader
    participant LLM as vLLMInference
    participant TTS as ChatterboxTTSStage
    participant MFA as MFAAlignmentStage
    participant Merge as MergeConversationSDPStage
    participant MW as ManifestWriterStage

    note over MR,MFA: Phase 1 (GPU — standard executor)
    MR->>LLM: AudioTask(topic)
    LLM->>LLM: VLLMModel.generate() + JSON validate
    LLM-->>TTS: AudioTask×N (one per turn, shared conversation_id)
    TTS->>TTS: ChatterboxTTS.generate() per turn
    TTS-->>MFA: AudioTask + audio_filepath
    MFA->>MFA: mfa align subprocess (batch)
    MFA-->>Merge: AudioTask + rttm_filepath + ctm_filepath

    note over Merge,MW: Phase 2 (CPU — grouped by conversation_id)
    Merge->>Merge: group turns, _merge_audio_files()
    Merge->>Merge: _compute_timeline() shared RTTM+CTM offsets
    Merge->>Merge: _merge_rttm_files() + _merge_ctm_files()
    Merge-->>MW: AudioTask(mixed.wav, all.rttm, all.ctm, seglst)
    MW-->>MW: write merged_manifest.jsonl

Loading

Comments Outside Diff (2)

1. nemo_curator/stages/audio/llm/vllm_inference.py, line 972-975 (link)

   Tokenizer left None when pre-initialized VLLMModel is passed

   _ensure_model guards initialization by checking the private _vllm_model._llm attribute. When a caller passes an already-initialized VLLMModel (i.e. _llm is already set), this guard evaluates to False, so setup() is never called and self.tokenizer stays None. The very next call to get_entry_prompt() then raises AttributeError: 'NoneType' object has no attribute 'apply_chat_template'. The docstring explicitly lists a pre-configured VLLMModel instance as a supported input, making this a reachable failure path.

2. nemo_curator/stages/audio/llm/vllm_inference.py, line 998-999 (link)

   Retry count hard-coded in process_batch

   generate_batch_with_retry accepts a configurable max_retry_rounds parameter, but process_batch always passes 5. If users need to tune retry behaviour (e.g. fewer retries for speed, more for quality), there is no way to do so without modifying the source. Consider surfacing this as a __init__ parameter or at least using a named constant that can be overridden.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

_{Reviews (1): Last reviewed commit: "feat(audio): wire full conversation data..." | Re-trigger Greptile}

[greptile-apps]

Missing required output fields on skipped tasks

When text is empty the method appends the raw input task, which lacks the declared output keys audio_filepath, duration, and reference_voice. Any downstream stage that reads these keys (e.g. MFAAlignmentStage expecting audio_filepath) will raise a KeyError or FileNotFoundError for every such task.

Suggested change

	if not text or not text.strip():
	logger.warning(f"Skipping task {task.task_id}: no text")
	output_tasks.append(task)
	continue
	if not text or not text.strip():
	logger.warning(f"Skipping task {task.task_id}: no text")
	out_data = dict(data)
	out_data.setdefault("audio_filepath", "")
	out_data.setdefault("duration", 0.0)
	out_data.setdefault("reference_voice", "")
	output_tasks.append(
	AudioTask(
	data=out_data,
	task_id=task.task_id,
	dataset_name=task.dataset_name,
	)
	)
	continue

[greptile-apps]

_is_node_local_mfa_root returns True without checking the actual path

When copy_models_to_local=True the method immediately returns True, regardless of whether _mfa_root was actually set to the local copy. If setup_on_node was never called (e.g. single-node local run) and setup() fell back to _effective_mfa_root (the shared path), this method returns True for a shared path. Consequently _run_mfa_align may delete command_history.yaml from the shared MFA root, corrupting state for any parallel worker still using it.

[greptile-apps]

Stage indices hard-coded; silent breakage if pipeline.yaml adds/removes stages

merge_cfg = stage_cfgs[_PHASE1_STAGE_COUNT] and writer_cfg = stage_cfgs[_PHASE1_STAGE_COUNT + 1] work only when phase-1 has exactly 4 stages in positions 0–3 and the merge/writer stages are at fixed positions. Adding a stage to phase 1 (e.g. a normalisation step) or reordering stages in pipeline.yaml silently pulls the wrong stage config with no error until runtime. Consider identifying stage configs by _target_ class name rather than by index.