From 953ebaf5cdc89955bec592a568ea8d07cdfb45fe Mon Sep 17 00:00:00 2001 From: Badalyan Vyacheslav Date: Tue, 13 Jan 2026 04:49:34 +0300 Subject: [PATCH 1/2] Added functionality for transcribing multi-channel audio with automatic diarization (channel/speaker separation). The method correctly processes overlapping speech and sorts the results by time, ensuring the correct alternation of cues between channels. --- README.md | 16 +++ README_ru.md | 16 +++ gigaam/__init__.py | 3 +- gigaam/model.py | 93 ++++++++++++++++- gigaam/preprocess.py | 101 ++++++++++++++++++- gigaam/vad_utils.py | 144 +++++++++++++++++++++++++- tests/test_multichannel.py | 200 +++++++++++++++++++++++++++++++++++++ 7 files changed, 568 insertions(+), 5 deletions(-) create mode 100644 tests/test_multichannel.py diff --git a/README.md b/README.md index c2b29ac..29bd1b4 100644 --- a/README.md +++ b/README.md @@ -116,6 +116,22 @@ for utt in utterances: transcription, (start, end) = utt["transcription"], utt["boundaries"] print(f"[{gigaam.format_time(start)} - {gigaam.format_time(end)}]: {transcription}") +# Multichannel transcription (diarization) +# Supports stereo/multichannel files or a list of separate files +# Results are automatically sorted by time and interleaved between channels +stereo_file = "conversation_stereo.wav" # or list: ["channel_0.wav", "channel_1.wav"] +results = model.transcribe_multichannel( + stereo_file, + batch_size=4, # batch size for processing segments + pause_threshold=2.0, # pause threshold for grouping segments (seconds) + strict_limit_duration=30.0 # maximum segment duration for model (seconds) +) +for seg in results: + channel = seg["channel"] # channel number (0, 1, ...) + transcription = seg["transcription"] + start, end = seg["boundaries"] + print(f"[{start:.2f}s - {end:.2f}s] Channel {channel}: {transcription}") + # Emotion recognition model = gigaam.load_model("emo") emotion2prob = model.get_probs(audio_path) diff --git a/README_ru.md b/README_ru.md index eaae27b..8741afd 100644 --- a/README_ru.md +++ b/README_ru.md @@ -115,6 +115,22 @@ for utt in utterances: transcription, (start, end) = utt["transcription"], utt["boundaries"] print(f"[{gigaam.format_time(start)} - {gigaam.format_time(end)}]: {transcription}") +# Мультиканальная транскрибация (диаризация) +# Поддерживает стерео/многоканальные файлы или список отдельных файлов +# Результаты автоматически сортируются по времени и чередуются между каналами +stereo_file = "conversation_stereo.wav" # или список: ["channel_0.wav", "channel_1.wav"] +results = model.transcribe_multichannel( + stereo_file, + batch_size=4, # размер батча для обработки сегментов + pause_threshold=2.0, # порог паузы для группировки сегментов (секунды) + strict_limit_duration=30.0 # максимальная длительность сегмента для модели (секунды) +) +for seg in results: + channel = seg["channel"] # номер канала (0, 1, ...) + transcription = seg["transcription"] + start, end = seg["boundaries"] + print(f"[{start:.2f}s - {end:.2f}s] Канал {channel}: {transcription}") + # Распознавание эмоций model = gigaam.load_model("emo") emotion2prob = model.get_probs(audio_path) diff --git a/gigaam/__init__.py b/gigaam/__init__.py index d43a8fc..4c15547 100644 --- a/gigaam/__init__.py +++ b/gigaam/__init__.py @@ -9,7 +9,7 @@ from tqdm import tqdm from .model import GigaAM, GigaAMASR, GigaAMEmo -from .preprocess import load_audio +from .preprocess import load_audio, load_multichannel_audio from .utils import format_time __all__ = [ @@ -17,6 +17,7 @@ "GigaAMASR", "GigaAMEmo", "load_audio", + "load_multichannel_audio", "format_time", "load_model", ] diff --git a/gigaam/model.py b/gigaam/model.py index ac38c1d..a49ecfa 100644 --- a/gigaam/model.py +++ b/gigaam/model.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Tuple, Union +from typing import Dict, List, Optional, Tuple, Union import hydra import omegaconf @@ -169,6 +169,97 @@ def transcribe_longform( ) return transcribed_segments + @torch.inference_mode() + def transcribe_multichannel( + self, + audio_input: Union[str, List[str]], + batch_size: int = 4, + **kwargs + ) -> List[Dict[str, Union[int, str, Tuple[float, float]]]]: + """ + Transcribes multichannel audio with synchronized diarization. + + Supports: + - Single stereo/multichannel file (str) + - Multiple separate audio files (List[str]) + + Handles overlapping speech by cutting segments when other channel starts. + Maintains decoder state between segments for better quality. + + Parameters: + ----------- + audio_input : Union[str, List[str]] + Either a single multichannel audio file or list of separate files + batch_size : int + Batch size for processing segments (default: 4) + **kwargs + Additional arguments passed to segment_multichannel_audio + + Returns: + -------- + List of dicts with keys: 'channel', 'transcription', 'boundaries' (start, end) + """ + from .vad_utils import segment_multichannel_audio + + # Segment audio with diarization + segments = segment_multichannel_audio( + audio_input, SAMPLE_RATE, device=self._device, **kwargs + ) + + if not segments: + return [] + + transcribed_segments = [] + + # Process all segments together in batches, regardless of channel + # Channel information is only used in the final output + + # Process all segments in batches - no state preservation needed + for batch_start in range(0, len(segments), batch_size): + batch_segments = segments[batch_start:batch_start + batch_size] + + # Prepare batch - audio is already on GPU from segmentation + batch_audio = [] + batch_lengths = [] + + for seg in batch_segments: + audio = seg["audio"] + # Ensure correct dtype (device should already be correct) + if audio.dtype != self._dtype: + audio = audio.to(self._dtype) + # Ensure audio is 1D: (samples,) + if audio.dim() > 1: + audio = audio.squeeze() + batch_audio.append(audio) + batch_lengths.append(len(audio)) + + # Pad and batch - more efficient: create tensor and fill in one pass + max_len = max(batch_lengths) + batched_audio = torch.zeros( + len(batch_audio), max_len, dtype=self._dtype, device=self._device + ) + for i, audio in enumerate(batch_audio): + batched_audio[i, :len(audio)] = audio + + # Format: (batch, samples) - same as transcribe_longform uses + batched_lengths = torch.tensor(batch_lengths, device=self._device, dtype=torch.long) + + # Forward pass + encoded, encoded_len = self.forward(batched_audio, batched_lengths) + + # Decode + batch_results = self.decoding.decode(self.head, encoded, encoded_len) + + # Store transcribed segments + for idx, seg in enumerate(batch_segments): + transcribed_segments.append({ + "channel": seg["channel"], + "transcription": batch_results[idx], + "boundaries": seg["boundaries"], + }) + + return transcribed_segments + class GigaAMEmo(GigaAM): """ diff --git a/gigaam/preprocess.py b/gigaam/preprocess.py index fb6ebde..dad858d 100644 --- a/gigaam/preprocess.py +++ b/gigaam/preprocess.py @@ -1,6 +1,6 @@ import warnings from subprocess import CalledProcessError, run -from typing import Tuple +from typing import List, Tuple, Union import torch import torchaudio @@ -40,6 +40,105 @@ def load_audio(audio_path: str, sample_rate: int = SAMPLE_RATE) -> Tensor: return torch.frombuffer(audio, dtype=torch.int16).float() / 32768.0 +def load_multichannel_audio( + audio_input: Union[str, List[str]], + sample_rate: int = SAMPLE_RATE +) -> Tuple[List[Tensor], int]: + """ + Load multichannel audio from either: + - A single stereo/multichannel file (str) + - Multiple separate audio files (List[str]) + + Returns: + Tuple of (list of channel tensors, max_length) + """ + if isinstance(audio_input, str): + # Try to load with torchaudio first (more reliable for multichannel) + try: + import torchaudio + waveform, file_sr = torchaudio.load(audio_input) + + # Resample if needed + if file_sr != sample_rate: + resampler = torchaudio.transforms.Resample(file_sr, sample_rate) + waveform = resampler(waveform) + + # Convert to list of channel tensors + num_channels = waveform.shape[0] + channels = [waveform[i] for i in range(num_channels)] + + max_length = max(len(ch) for ch in channels) + return channels, max_length + except Exception: + # Fallback to ffmpeg approach + pass + + # Fallback: Load multichannel file with ffmpeg + cmd = [ + "ffmpeg", + "-nostdin", + "-threads", + "0", + "-i", + audio_input, + "-f", + "s16le", + "-acodec", + "pcm_s16le", + "-ar", + str(sample_rate), + "-", + ] + try: + audio_bytes = run(cmd, capture_output=True, check=True).stdout + except CalledProcessError as exc: + raise RuntimeError(f"Failed to load audio from {audio_input}") from exc + + # Try to determine number of channels from file metadata + # Default to stereo (2 channels) for common cases + num_channels = 2 # Default assumption + + # Try ffprobe if available + cmd_probe = [ + "ffprobe", + "-v", "error", + "-show_entries", "stream=channels", + "-of", "default=noprint_wrappers=1:nokey=1", + audio_input + ] + try: + result = run(cmd_probe, capture_output=True, check=True) + num_channels = int(result.stdout.strip().split()[0]) + except (CalledProcessError, ValueError, IndexError): + # If ffprobe fails, try to infer from data size + # This is a heuristic - may not always work + pass + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=UserWarning) + audio_data = torch.frombuffer(audio_bytes, dtype=torch.int16).float() / 32768.0 + + # Reshape to channels + if num_channels > 1 and len(audio_data) % num_channels == 0: + audio_data = audio_data.view(-1, num_channels).transpose(0, 1) + channels = [audio_data[i] for i in range(num_channels)] + else: + # Single channel or couldn't determine + channels = [audio_data] + + max_length = max(len(ch) for ch in channels) + return channels, max_length + + else: + # Load multiple separate files + channels = [] + for path in audio_input: + channels.append(load_audio(path, sample_rate)) + + max_length = max(len(ch) for ch in channels) + return channels, max_length + + class SpecScaler(nn.Module): """ Module that applies logarithmic scaling to spectrogram values. diff --git a/gigaam/vad_utils.py b/gigaam/vad_utils.py index 5d6219f..4dedf0b 100644 --- a/gigaam/vad_utils.py +++ b/gigaam/vad_utils.py @@ -1,5 +1,6 @@ import os -from typing import List, Tuple +from itertools import chain, groupby +from typing import Dict, List, Tuple, Union import torch from huggingface_hub import snapshot_download @@ -9,7 +10,7 @@ from pyannote.audio.pipelines import VoiceActivityDetection from torch.torch_version import TorchVersion -from .preprocess import load_audio +from .preprocess import load_audio, load_multichannel_audio, SAMPLE_RATE _PIPELINE = None @@ -132,3 +133,142 @@ def _update_segments(curr_start: float, curr_end: float, curr_duration: float): _update_segments(curr_start, curr_end, curr_duration) return segments, boundaries + + +def segment_multichannel_audio( + audio_input: Union[str, List[str]], + sr: int = SAMPLE_RATE, + pause_threshold: float = 2.0, + strict_limit_duration: float = 30.0, + device: torch.device = torch.device("cpu"), +) -> List[Dict[str, Union[int, torch.Tensor, Tuple[float, float]]]]: + """ + Segments multichannel audio with synchronized diarization. + + Simple approach: + 1. Segment each channel with pause_threshold (2 sec by default) + 2. Sort all segments by start_time + 3. Merge segments from same channel (up to strict_limit_duration) + 4. Reduce long pauses to 1 sec to save GPU resources + + Returns: + List of segment dicts with keys: 'channel', 'audio', 'boundaries' (start, end) + boundaries contain REAL start/end times (not affected by pause reduction) + """ + # Load multichannel audio + channels, max_length = load_multichannel_audio(audio_input, sr) + num_channels = len(channels) + + # Move channels to device and pad all channels to same length + # Do this in one pass to minimize CPU-GPU transfers + for i in range(num_channels): + channels[i] = channels[i].to(device) # Move to GPU first + if len(channels[i]) < max_length: + padding = torch.zeros(max_length - len(channels[i]), device=device, dtype=channels[i].dtype) + channels[i] = torch.cat([channels[i], padding]) # Already on GPU, no transfer needed + + pipeline = get_pipeline(device) + + # Step 1: Get ALL small VAD segments for ALL channels (don't merge yet!) + all_segments: List[Dict[str, Union[int, torch.Tensor, Tuple[float, float], float]]] = [] + + for channel_idx, channel_audio in enumerate(channels): + # Track last segment info for THIS channel only (in channel scope!) + prev_end: float = None + prev_global_start: float = None + + # Use pipeline with tensor directly - NO DISK I/O! + channel_audio_tensor = channel_audio.unsqueeze(0) # (1, num_samples) + input_dict = {"waveform": channel_audio_tensor, "sample_rate": sr} + sad_segments = pipeline(input_dict) + + # Get all small VAD segments for this channel (each contains several words) + # Don't merge yet - we need to sort ALL segments from ALL channels first! + for segment in sad_segments.get_timeline().support(): + start = max(0, segment.start) + end = min(max_length / sr, segment.end) + + # Calculate global_start: if pause < pause_threshold from previous segment of THIS channel, + # use previous global_start, otherwise use current start + if prev_end is not None: + pause_from_prev = start - prev_end + if pause_from_prev < pause_threshold: + global_start = prev_global_start + else: + global_start = start + else: + global_start = start + + # Extract audio tensor for this segment (already on GPU) + start_idx = int(start * sr) + end_idx = int(end * sr) + seg_audio = channels[channel_idx][start_idx:end_idx].clone() # Clone to avoid keeping reference to large tensor + + all_segments.append({ + "channel": channel_idx, + "audio": seg_audio, + "boundaries": (start, end), # Real boundaries + "global_start": global_start, # Start of the group this segment belongs to + }) + + # Update last segment info for THIS channel + prev_end = end + prev_global_start = global_start + + # Step 2: Sort ALL segments from ALL channels by global_start, then by start_time + all_segments.sort(key=lambda x: (x["global_start"], x["boundaries"][0])) + + # Step 3: Group by channel, then merge segments in each group (up to strict_limit_duration) + + def merge_channel_segments(channel_segments: List[Dict]) -> List[Dict]: + """Merge segments in a channel group, splitting by strict_limit_duration windows""" + if not channel_segments: + return [] + + # Assign window_idx to each segment based on accumulated audio duration + accumulated_duration = 0.0 + segments_with_window = [] + for seg in channel_segments: + seg_duration = len(seg["audio"]) / sr + window_idx = int(accumulated_duration / strict_limit_duration) + accumulated_duration += seg_duration + segments_with_window.append((window_idx, seg)) + + # Sort by window_idx, then group by window_idx, then map merge function + def merge_window_segments(window_group): + window_idx, window_segs_iter = window_group + # Extract seg from (window_idx, seg) tuples + window_segs = [seg for _, seg in window_segs_iter] + channel_idx = window_segs[0]["channel"] + + # Concatenate all audio in window (no pauses) + all_audio = [seg["audio"] for seg in window_segs] + merged_audio = torch.cat(all_audio) + + # Get min start and max end + all_starts = [seg["boundaries"][0] for seg in window_segs] + all_ends = [seg["boundaries"][1] for seg in window_segs] + merged_start = min(all_starts) + merged_end = max(all_ends) + + return { + "channel": channel_idx, + "audio": merged_audio, + "boundaries": (merged_start, merged_end), + } + + # Sort -> groupby -> map + sorted_segments = sorted(segments_with_window, key=lambda x: x[0]) + grouped_by_window = groupby(sorted_segments, key=lambda x: x[0]) + merged_results = list(map(merge_window_segments, grouped_by_window)) + + return merged_results + + # Step 3: Group by channel, then merge segments in each group (up to strict_limit_duration) + merged_groups = map( + lambda channel_group: merge_channel_segments(list(channel_group[1])), + groupby(all_segments, key=lambda x: x["channel"]) + ) + final_segments = list(chain.from_iterable(merged_groups)) + + return final_segments diff --git a/tests/test_multichannel.py b/tests/test_multichannel.py new file mode 100644 index 0000000..7ff8730 --- /dev/null +++ b/tests/test_multichannel.py @@ -0,0 +1,200 @@ +import logging +import os +import tempfile +import urllib.request +from typing import List, Tuple + +import numpy as np +import pytest +import soundfile as sf +import torch + +import gigaam + + +def download_long_audio(): + """Download test audio file if not exists""" + audio_file = "long_example.wav" + if not os.path.exists(audio_file): + url = "https://cdn.chatwm.opensmodel.sberdevices.ru/GigaAM/long_example.wav" + urllib.request.urlretrieve(url, audio_file) + assert os.path.exists(audio_file), "Long audio file not found" + return audio_file + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +def create_stereo_from_mono(audio: np.ndarray, offset_seconds: float = 0.0, sr: int = 16000) -> np.ndarray: + """ + Create stereo audio from mono by duplicating with optional time offset. + Channel 0: original audio + Channel 1: audio shifted by offset_seconds + """ + offset_samples = int(offset_seconds * sr) + + # Create two channels + channel_0 = audio.copy() + + # Channel 1: shift audio by offset + if offset_samples > 0: + # Pad beginning with zeros + channel_1 = np.pad(audio, (offset_samples, 0), mode='constant') + # Trim to same length + channel_1 = channel_1[:len(audio)] + elif offset_samples < 0: + # Shift left (remove from beginning) + channel_1 = audio[-offset_samples:] + # Pad end with zeros + channel_1 = np.pad(channel_1, (0, -offset_samples), mode='constant') + else: + channel_1 = audio.copy() + + # Ensure same length + min_len = min(len(channel_0), len(channel_1)) + channel_0 = channel_0[:min_len] + channel_1 = channel_1[:min_len] + + # Stack into stereo: (2, samples) + stereo = np.stack([channel_0, channel_1], axis=0) + return stereo + + +@pytest.mark.parametrize("revision", ["v3_e2e_rnnt", "v3_ctc"]) +def test_transcribe_multichannel_stereo(revision): + """Test multichannel transcription with stereo file created from mono""" + # Download test audio + mono_file = download_long_audio() + + # Load mono audio + audio, sr = sf.read(mono_file) + + # Create stereo file with offset (channel 1 starts 5 seconds later) + stereo_audio = create_stereo_from_mono(audio, offset_seconds=5.0, sr=sr) + + # Save to temporary file + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f: + temp_file = f.name + try: + sf.write(temp_file, stereo_audio.T, sr) # soundfile expects (samples, channels) + + # Load model + model = gigaam.load_model(revision) + + # Test with stereo file + results = model.transcribe_multichannel(temp_file, batch_size=4) + + assert isinstance(results, list), "Should return list of segments" + assert len(results) > 0, "Should have at least one segment" + + # Check structure + for seg in results: + assert "channel" in seg, "Missing channel key" + assert "transcription" in seg, "Missing transcription key" + assert "boundaries" in seg, "Missing boundaries key" + assert seg["channel"] in [0, 1], f"Invalid channel: {seg['channel']}" + start, end = seg["boundaries"] + assert start < end, f"Invalid boundaries: {start} >= {end}" + + logger.info(f"Multichannel test: {len(results)} segments for stereo file") + + finally: + if os.path.exists(temp_file): + os.remove(temp_file) + + +@pytest.mark.parametrize("revision", ["v3_e2e_rnnt"]) +def test_transcribe_multichannel_list(revision): + """Test multichannel transcription with list of separate files""" + # Download test audio + mono_file = download_long_audio() + + # Load mono audio + audio, sr = sf.read(mono_file) + + # Split into two parts with offset + split_point = len(audio) // 2 + audio_0 = audio[:split_point] + audio_1 = audio[split_point:] + + # Save to temporary files + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f0, \ + tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f1: + temp_file0 = f0.name + temp_file1 = f1.name + try: + sf.write(temp_file0, audio_0, sr) + sf.write(temp_file1, audio_1, sr) + + # Load model + model = gigaam.load_model(revision) + + # Test with list of files + results = model.transcribe_multichannel([temp_file0, temp_file1], batch_size=4) + + assert isinstance(results, list), "Should return list of segments" + assert len(results) > 0, "Should have at least one segment" + + # Check structure + for seg in results: + assert "channel" in seg, "Missing channel key" + assert "transcription" in seg, "Missing transcription key" + assert "boundaries" in seg, "Missing boundaries key" + assert seg["channel"] in [0, 1], f"Invalid channel: {seg['channel']}" + + logger.info(f"Multichannel test: {len(results)} segments for list of files") + + finally: + for fname in [temp_file0, temp_file1]: + if os.path.exists(fname): + os.remove(fname) + + +def test_multichannel_channel_ordering(): + """Test that channels are correctly identified and ordered""" + # Download test audio + mono_file = download_long_audio() + + # Load mono audio + audio, sr = sf.read(mono_file) + + # Create stereo: channel 0 = first half, channel 1 = second half + split_point = len(audio) // 2 + channel_0 = audio[:split_point] + channel_1 = audio[split_point:] + + # Pad to same length + max_len = max(len(channel_0), len(channel_1)) + channel_0 = np.pad(channel_0, (0, max_len - len(channel_0)), mode='constant') + channel_1 = np.pad(channel_1, (0, max_len - len(channel_1)), mode='constant') + + stereo_audio = np.stack([channel_0, channel_1], axis=0) + + # Save to temporary file + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f: + temp_file = f.name + try: + sf.write(temp_file, stereo_audio.T, sr) + + model = gigaam.load_model("v3_e2e_rnnt") + results = model.transcribe_multichannel(temp_file, batch_size=4) + + # Check that we have segments from both channels + channels_found = set(seg["channel"] for seg in results) + assert 0 in channels_found, "Should have segments from channel 0" + assert 1 in channels_found, "Should have segments from channel 1" + + # Check ordering: segments should be sorted by time + for i in range(len(results) - 1): + assert results[i]["boundaries"][0] <= results[i+1]["boundaries"][0], \ + "Segments should be sorted by start time" + + logger.info(f"Channel ordering test: channels {channels_found}, {len(results)} segments") + + finally: + if os.path.exists(temp_file): + os.remove(temp_file) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 68f1d870362e523e70dc1b5e1f97e6cfb7c665f4 Mon Sep 17 00:00:00 2001 From: Badalian Vyacheslav Date: Wed, 14 Jan 2026 04:23:25 +0300 Subject: [PATCH 2/2] Fix error - stft(torch.cuda.HalfTensor[1, 270], n_fft=320, hop_length=160, win_length=320, window=torch.cuda.FloatTensor{[320]}, normalized=0, onesided=1, return_complex=1, align_to_window=None) : expected 0 < n_fft < 270, but got n_fft=320 --- gigaam/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gigaam/model.py b/gigaam/model.py index a49ecfa..2a18370 100644 --- a/gigaam/model.py +++ b/gigaam/model.py @@ -236,7 +236,7 @@ def transcribe_multichannel( # Pad and batch - more efficient: create tensor and fill in one pass max_len = max(batch_lengths) batched_audio = torch.zeros( - len(batch_audio), max_len, dtype=self._dtype, device=self._device + len(batch_audio), max(max_len, 320), dtype=self._dtype, device=self._device ) for i, audio in enumerate(batch_audio): batched_audio[i, :len(audio)] = audio