Skip to content

feat(audio-segmentation): in-memory segment handoff + agent-ready contracts #2338

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shubhamNvidia wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(audio-segmentation): in-memory segment handoff + agent-ready contracts #2338

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
159 changes: 146 additions & 13 deletions nemo_curator/stages/audio/segmentation/speaker_separation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
)
"""

import os
from dataclasses import dataclass, field
from typing import Any

Expand All @@ -44,7 +45,13 @@
SortformerEncLabelModel = None

from nemo_curator.backends.base import WorkerMetadata
from nemo_curator.stages.audio.common import resolve_waveform_from_item
from nemo_curator.stages.audio._agent._agent_ready import AgentReady, Gates, IOSpec, StageContract
from nemo_curator.stages.audio._agent._residency import (
InputResidency,
accepts_for_residency,
resolve_audio,
write_audio_stable,
)
from nemo_curator.stages.audio.segmentation.speaker_separation_module.speaker_sep import SpeakerSeparator
from nemo_curator.stages.base import ProcessingStage
from nemo_curator.stages.resources import Resources
Expand All @@ -61,7 +68,7 @@ def _pydub_to_waveform_sr(seg: AudioSegment) -> tuple[torch.Tensor, int]:


@dataclass
class SpeakerSeparationStage(ProcessingStage[AudioTask, AudioTask]):
class SpeakerSeparationStage(AgentReady, ProcessingStage[AudioTask, AudioTask]):
"""
Speaker separation stage using NeMo SortFormer diarization model.

Expand All @@ -75,8 +82,28 @@ class SpeakerSeparationStage(ProcessingStage[AudioTask, AudioTask]):
min_duration: Minimum segment duration in seconds
gap_threshold: Gap threshold for merging speaker segments
buffer_time: Buffer time around speaker segments
audio_filepath_key: Key in data dict for the input audio file path.
waveform_key: Key in data dict for the in-memory waveform tensor.
sample_rate_key: Key in data dict for the waveform sample rate.
speaker_id_key: Key where each child task's speaker id is written.
num_speakers_key: Key where the detected speaker count is written.
duration_key: Key where each child's speech duration in seconds is written.
diar_segments_key: Key where each child's diarization segments are written.
input_residency: Which input to use — "waveform" (in-memory only), "file"
(audio_filepath only), or "auto" (waveform first, file fallback; default).
keep_waveform_in_task: Keep each per-speaker waveform in the task (default True,
today's behavior). Set False to emit only on-disk paths (requires write_to_disk).
write_to_disk: Also write each per-speaker track to ``separated_audio_dir`` and set
``audio_filepath_key`` to it, so file-based downstream stages can consume it.
Defaults to False (in-memory only, unchanged).
separated_audio_dir: Directory for per-speaker WAVs (required when write_to_disk=True).

Note:
By default (write_to_disk=False) per-speaker child tasks DROP the parent's
audio_filepath (it points at the full multi-speaker file) and carry
``original_file`` for provenance; downstream consumes the per-speaker waveform.
With write_to_disk=True, each child instead gets its own audio_filepath.

GPU assignment is handled by the executor via _resources.
Use .with_(resources=Resources(gpus=X)) to configure GPU allocation.
"""
Expand All @@ -86,6 +113,18 @@ class SpeakerSeparationStage(ProcessingStage[AudioTask, AudioTask]):
min_duration: float = 0.8
gap_threshold: float = 0.1
buffer_time: float = 0.5
audio_filepath_key: str = "audio_filepath"
waveform_key: str = "waveform"
sample_rate_key: str = "sample_rate"
speaker_id_key: str = "speaker_id"
num_speakers_key: str = "num_speakers"
duration_key: str = "duration"
diar_segments_key: str = "diar_segments"
input_residency: InputResidency = "auto"
# Output residency (both default to today's behavior: in-memory waveform only, no disk).
keep_waveform_in_task: bool = True
write_to_disk: bool = False
separated_audio_dir: str | None = None

name: str = "SpeakerSeparation"
batch_size: int = 1
Expand All @@ -94,12 +133,64 @@ class SpeakerSeparationStage(ProcessingStage[AudioTask, AudioTask]):
def __post_init__(self):
super().__init__()
self._separator = None
if not (self.keep_waveform_in_task or self.write_to_disk):
msg = "At least one of keep_waveform_in_task or write_to_disk must be True"
raise ValueError(msg)
if self.write_to_disk and not self.separated_audio_dir:
msg = "separated_audio_dir is required when write_to_disk=True"
raise ValueError(msg)

def inputs(self) -> tuple[list[str], list[str]]:
return [], []

def outputs(self) -> tuple[list[str], list[str]]:
return [], ["waveform", "sample_rate", "speaker_id", "num_speakers", "duration"]
outs: list[str] = []
if self.keep_waveform_in_task:
outs.extend([self.waveform_key, self.sample_rate_key])
outs.extend([self.speaker_id_key, self.num_speakers_key, self.duration_key, self.diar_segments_key])
if self.write_to_disk:
outs.append(self.audio_filepath_key)
return [], outs

def describe(self) -> StageContract:
forms = accepts_for_residency(self.input_residency)
reads_one_of = []
if "waveform" in forms:
reads_one_of.append(IOSpec(data_keys=[self.waveform_key, self.sample_rate_key], accepts=["waveform"]))
if "file" in forms:
reads_one_of.append(IOSpec(data_keys=[self.audio_filepath_key], accepts=["file"]))
writes: list[str] = []
produces: list[str] = []
if self.keep_waveform_in_task:
writes.extend([self.waveform_key, self.sample_rate_key])
produces.append("tensor")
writes.extend(
[self.speaker_id_key, self.num_speakers_key, self.duration_key, self.diar_segments_key, "original_file"]
)
if self.write_to_disk:
writes.append(self.audio_filepath_key)
produces.append("disk")
return StageContract(
reads_one_of=reads_one_of,
writes=IOSpec(data_keys=writes, produces=produces),
# children drop the parent's audio_filepath (and blob keys)
preserves_upstream_keys=False,
cardinality="1:N fan-out",
# One child per detected speaker; speaker_id is the per-child key that
# identifies which slice of the iteration a child is (role-resolvable).
iteration_key=self.speaker_id_key,
gates=Gates(
requires_gpu=self.resources.requires_gpu,
requires_internet_first_run=True,
writes_to_disk=self.write_to_disk,
output_path_params=["separated_audio_dir"],
# Diarization runs on one file's audio, and ``num_speakers`` counts the speakers
# found in THAT file. Unlike SplitLongAudioStage, a shared output directory is
# still safe here: ``write_audio_stable`` names each per-speaker WAV after a
# digest of its own samples rather than after the source basename.
per_row_independent=True,
),
)

def setup_on_node(self, _node_info: Any = None, _worker_metadata: Any = None) -> None: # noqa: ANN401
try:
Expand Down Expand Up @@ -160,6 +251,17 @@ def _initialize_separator(self) -> None:
# each child gets its own duration from the diarization result.
_INHERITED_DROP_KEYS = frozenset({"audio", "waveform", "duration", "num_samples"})

def _write_speaker_wav(self, waveform: torch.Tensor, sr: int, original_file: str, speaker_id: str) -> str:
"""Write one per-speaker waveform to ``separated_audio_dir`` and return the path."""
stem = os.path.splitext(os.path.basename(str(original_file)))[0] or "audio"
return write_audio_stable(
waveform,
sr,
output_dir=self.separated_audio_dir,
stem=stem,
tag=str(speaker_id),
)

def _build_speaker_tasks(
self,
speaker_audio_data: dict,
Expand All @@ -174,21 +276,46 @@ def _build_speaker_tasks(
logger.debug(f"Skipping {speaker_id}: duration {result.duration:.2f}s < {self.min_duration}s")
continue
spk_waveform, spk_sr = _pydub_to_waveform_sr(result.audio)
# Drop the parent's file path(s) too: they point at the FULL
# multi-speaker file, so a file-preferring downstream stage would
# process the whole file per speaker instead of this speaker's
# extracted waveform. With the path gone, downstream resolves the
# per-speaker waveform (input_residency="auto") instead.
drop_keys = {
*self._INHERITED_DROP_KEYS,
self.waveform_key,
self.duration_key,
self.audio_filepath_key,
"audio_filepath",
}
speaker_data = {
**{k: v for k, v in item.items() if k not in self._INHERITED_DROP_KEYS},
"waveform": spk_waveform,
"sample_rate": spk_sr,
"speaker_id": speaker_id,
"num_speakers": num_speakers,
"duration": result.duration,
"diar_segments": result.diar_segments,
**{k: v for k, v in item.items() if k not in drop_keys},
self.speaker_id_key: speaker_id,
self.num_speakers_key: num_speakers,
self.duration_key: result.duration,
self.diar_segments_key: result.diar_segments,
# Source identity must survive the audio_filepath drop above —
# TimestampMapper (and any provenance consumer) reads original_file.
"original_file": item.get("original_file")
or item.get(self.audio_filepath_key)
or item.get("audio_filepath")
or "unknown",
}
# Output residency: keep the in-memory waveform (default) and/or persist a
# per-speaker WAV and point audio_filepath at it (opt-in write_to_disk).
if self.keep_waveform_in_task:
speaker_data[self.waveform_key] = spk_waveform
speaker_data[self.sample_rate_key] = spk_sr
if self.write_to_disk:
speaker_data[self.audio_filepath_key] = self._write_speaker_wav(
spk_waveform, spk_sr, speaker_data["original_file"], speaker_id
)
spk_task = AudioTask(
data=speaker_data,
dataset_name=task.dataset_name,
_metadata=dict(task._metadata or {}),
_stage_perf=list(task._stage_perf),
)
if task._metadata:
spk_task._metadata = dict(task._metadata)
results.append(spk_task)
return results

Expand All @@ -208,7 +335,13 @@ def process(self, task: AudioTask) -> list[AudioTask]:
results: list[AudioTask] = []

try:
audio_result = resolve_waveform_from_item(item, task.task_id)
audio_result = resolve_audio(
item,
residency=self.input_residency, # type: ignore[arg-type]
audio_filepath_key=self.audio_filepath_key,
waveform_key=self.waveform_key,
sample_rate_key=self.sample_rate_key,
)
if audio_result is None:
return []
waveform, sample_rate = audio_result
Expand Down
Loading
Loading