Merge pull request #35 from xingchensong/longaudio

[feat] support long audio

Author: xingchensong

.flake8

@@ -0,0 +1,28 @@
+[flake8]
+# Suggested config from pytorch that we can adapt
+select = B,C,E,F,N,P,T4,W,B9,TOR0,TOR1,TOR2
+max-line-length = 120
+# C408 ignored because we like the dict keyword argument syntax
+# E501 is not flexible enough, we're using B950 instead
+# N812 ignored because import torch.nn.functional as F is PyTorch convention
+# N817 ignored because importing using acronyms is convention (DistributedDataParallel as DDP)
+# E731 allow usage of assigning lambda expressions
+# N803,N806 allow caps and mixed case in function params. This is to work with Triton kernel coding style.
+ignore =
+    E203,E305,E402,E501,E721,E741,F405,F821,F841,F999,W503,W504,C408,E302,W291,E303,N812,N817,E731,N803,N806
+    # shebang has extra meaning in fbcode lints, so I think it's not worth trying
+    # to line this up with executable bit
+    EXE001,
+    # these ignores are from flake8-bugbear; please fix!
+    B007,B008,
+optional-ascii-coding = True
+exclude =
+    ./.git,
+    ./docs
+    ./build
+    ./scripts,
+    ./venv,
+    *.pyi
+    .pre-commit-config.yaml
+    *.md
+    .flake8

.github/workflows/unit_test_cpu.yaml

@@ -0,0 +1,47 @@
+name: CPU Unit Test
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+
+concurrency:
+  group: unit-test${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  unit-test:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      max-parallel: 20
+      matrix:
+        os: [ubuntu-22.04]
+        python-version: [3.10.16]
+    steps:
+      - name: Cache Python Packages
+        uses: actions/cache@v4
+        with:
+          path: ~/.cache/pip
+          key: ${{ runner.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('**/pyproject.toml') }}
+      - name: Setup Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+          architecture: x64
+      - name: Fetch S3Tokenizer
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          ref: ${{ github.event.pull_request.head.ref || github.ref }}
+      - name: Install S3Tokenizer Dependencies
+        run: |
+          set -eux
+          sudo apt update && sudo apt install -y ffmpeg  libsox-dev libsndfile1
+          pip install -e .
+      - name: Run Pytest
+        run: |
+          set -eux
+          pip install pytest onnxruntime
+          pytest --version
+          PYTHONPATH="${PYTHONPATH:-}:$(pwd)" pytest test/ -q
+          if [ $? != 0 ]; then exit 1; fi

README.md

@@ -14,10 +14,11 @@ This repository undertakes a reverse engineering of the S3Tokenizer, offering:
 2. High-throughput (distributed) batch inference, achieving a ~790x speedup compared to the original inference pipeline in [[cosyvoice/tools/extract_speech_token.py]](https://github.com/FunAudioLLM/CosyVoice/blob/main/tools/extract_speech_token.py).
 3. The capability to perform online speech code extraction during SpeechLLM training.
 
-## Supported Models 🔥
-- [x] [S3Tokenizer V1 50hz](https://modelscope.cn/models/iic/CosyVoice-300M)
-- [x] [S3Tokenizer V1 25hz](https://modelscope.cn/models/iic/CosyVoice-300M-25Hz)
-- [x] [S3Tokenizer V2 25hz](https://modelscope.cn/models/iic/CosyVoice2-0.5B)
+## Supported Models 🔥 && New Features 🎉
+- [x] Model: [S3Tokenizer V1 50hz](https://modelscope.cn/models/iic/CosyVoice-300M)
+- [x] Model: [S3Tokenizer V1 25hz](https://modelscope.cn/models/iic/CosyVoice-300M-25Hz)
+- [x] Model: [S3Tokenizer V2 25hz](https://modelscope.cn/models/iic/CosyVoice2-0.5B)
+- [x] Feature: S3Tokenizer now has built-in **long audio processing** capabilities, requiring no additional operations from users!
 
 
 # Setup
@@ -39,7 +40,7 @@ for wav_path in wav_paths:
     audio = s3tokenizer.load_audio(wav_path)
     mels.append(s3tokenizer.log_mel_spectrogram(audio))
 mels, mels_lens = s3tokenizer.padding(mels)
-codes, codes_lens = tokenizer.quantize(mels.cuda(), mels_lens.cuda())
+codes, codes_lens = tokenizer.quantize(mels.cuda(), mels_lens.cuda())  # Automatically handles long audio internally!
 
 for i in range(len(wav_paths)):
     print(codes[i, :codes_lens[i].item()])
@@ -139,9 +140,9 @@ class SpeechLLM(nn.Module):
 </tr>
 </table>
 
+# Usage-4: Long Audio Processing (Built-in Automatic Processing)
 
-# TODO
-
-- [x] Usage-1: Offline batch inference
-- [x] Usage-2: Distributed offline batch inference via command-line tools
-- [x] Usage-3: Online speech code extraction
+- **Automatic Detection**: Model automatically detects audio length (>30 seconds triggers long audio processing)
+- **Sliding Window**: 30-second window with 4-second overlap, automatically segments long audio
+- **Batch Processing**: Internal batch processing of multiple segments for improved efficiency
+- **Complete Transparency**: User calling method is identical to short audio

s3tokenizer/__init__.py

@@ -28,11 +28,11 @@
 
 from .model import S3Tokenizer
 from .utils import (load_audio, log_mel_spectrogram, make_non_pad_mask,
-                    mask_to_bias, onnx2torch, padding)
+                    mask_to_bias, onnx2torch, padding, merge_tokenized_segments)
 
 __all__ = [
     'load_audio', 'log_mel_spectrogram', 'make_non_pad_mask', 'mask_to_bias',
-    'onnx2torch', 'padding'
+    'onnx2torch', 'padding', 'merge_tokenized_segments'
 ]
 _MODELS = {
     "speech_tokenizer_v1":

s3tokenizer/cli.py

@@ -61,13 +61,7 @@ def __getitem__(self, idx):
         file_path = self.data[idx]
         key = self.keys[idx]
         audio = s3tokenizer.load_audio(file_path)
-        if audio.shape[0] / 16000 > 30:
-            print(
-                f'do not support extract speech token for audio longer than 30s, file_path: {file_path}'  # noqa
-            )
-            mel = torch.zeros(128, 0)
-        else:
-            mel = s3tokenizer.log_mel_spectrogram(audio)
+        mel = s3tokenizer.log_mel_spectrogram(audio)
         return key, mel
 
 

s3tokenizer/model.py

@@ -25,7 +25,7 @@
 from einops import rearrange
 from torch import Tensor, nn
 
-from .utils import make_non_pad_mask, mask_to_bias, onnx2torch
+from .utils import make_non_pad_mask, mask_to_bias, onnx2torch, merge_tokenized_segments
 
 
 @dataclass
@@ -236,7 +236,7 @@ def preprocess(self, x: Tensor) -> Tensor:
 
     @torch.inference_mode()
     def quantize(self, x: Tensor) -> Tensor:
-        embed = self.embed.t()
+        embed = self.embed.t().to(x.dtype)
         dist = -(x.pow(2).sum(1, keepdim=True) - 2 * x @ embed +
                  embed.pow(2).sum(0, keepdim=True))
         embed_ind = dist.max(dim=-1).indices
@@ -287,7 +287,7 @@ def codebook(self):
 
     @torch.inference_mode()
     def encode(self, x: Tensor) -> Tensor:
-        x = F.normalize(x, p=2, dim=-1)
+        x = F.normalize(x.float(), p=2, dim=-1)
         embed_in = self._codebook.encode(x)
         return embed_in
 
@@ -306,6 +306,7 @@ class S3Tokenizer(nn.Module):
 
     def __init__(self, name: str, config: ModelConfig = ModelConfig()):
         super().__init__()
+        self.name = name  # Store model name for token_rate determination
         self.config = config
         self.encoder = AudioEncoder(
             self.config.n_mels,
@@ -324,9 +325,209 @@ def forward(self, mel: Tensor, mel_len: Tensor) -> Tuple[Tensor, Tensor]:
 
     @torch.inference_mode()
     def quantize(self, mel: Tensor, mel_len: Tensor) -> Tuple[Tensor, Tensor]:
-        hidden, code_len = self.encoder(mel, mel_len)
-        code = self.quantizer.encode(hidden)
-        return code, code_len
+        """
+        Quantize mel spectrogram to tokens, with automatic long audio handling.
+
+        Args:
+            mel: mel spectrogram tensor, shape (batch_size, n_mels, T)
+            mel_len: mel length tensor, shape (batch_size,)
+
+        Returns:
+            code: quantized tokens, shape (batch_size, T')
+            code_len: token length, shape (batch_size,)
+        """
+        # Check if any audio in the batch exceeds 30 seconds
+        # Assuming 16kHz sample rate and hop_length=160, 30s = 30*16000/160 = 3000 frames
+        max_frames = 3000
+
+        # Check which samples are long audio
+        long_audio_mask = mel_len > max_frames
+
+        if long_audio_mask.any():
+            # Has long audio - need special processing
+            return self._quantize_mixed_batch(mel, mel_len, long_audio_mask,
+                                              max_frames)
+        else:
+            # All short audio - use original method
+            hidden, code_len = self.encoder(mel, mel_len)
+            code = self.quantizer.encode(hidden)
+            return code, code_len
+
+    @torch.inference_mode()
+    def _quantize_mixed_batch(self, mel: Tensor, mel_len: Tensor,
+                              long_audio_mask: Tensor,
+                              max_frames: int) -> Tuple[Tensor, Tensor]:
+        """
+        Handle mixed batch with both short and long audio using unified batch processing.
+
+        Args:
+            mel: mel spectrogram tensor, shape (batch_size, n_mels, T)
+            mel_len: mel length tensor, shape (batch_size,)
+            long_audio_mask: boolean mask for long audio, shape (batch_size,)
+            max_frames: maximum frames for short audio
+
+        Returns:
+            code: quantized tokens, shape (batch_size, T')
+            code_len: token length, shape (batch_size,)
+        """
+        batch_size = mel.size(0)
+
+        # Parameters for sliding window
+        sample_rate = 16000
+        hop_length = 160  # Default hop length for mel spectrogram
+        window_size = 30  # seconds
+        overlap = 4  # seconds
+
+        # Calculate frame-based parameters
+        frames_per_window = window_size * sample_rate // hop_length  # 3000 frames
+        frames_per_overlap = overlap * sample_rate // hop_length  # 400 frames
+        frames_per_stride = frames_per_window - frames_per_overlap  # 2600 frames
+
+        # Collect all segments to process (including short and long audio segments)
+        all_segments = []
+        all_segments_len = []
+        segment_info = [
+        ]  # Record which audio each segment belongs to and whether it's long audio
+
+        # Process all audio in the batch
+        for batch_idx in range(batch_size):
+            audio_mel = mel[batch_idx]
+            audio_mel_len = mel_len[batch_idx]
+            is_long_audio = long_audio_mask[batch_idx].item()
+
+            if not is_long_audio:
+                # Short audio: process directly as a single segment
+                segment = audio_mel[:, :audio_mel_len]
+                seg_len = audio_mel_len.item()
+
+                # Pad to max_frames if necessary
+                if seg_len < frames_per_window:
+                    pad_size = frames_per_window - seg_len
+                    segment = F.pad(segment, (0, pad_size))
+
+                all_segments.append(segment)
+                all_segments_len.append(
+                    torch.tensor(seg_len, device=mel.device))
+                segment_info.append({
+                    'batch_idx': batch_idx,
+                    'is_long_audio': False,
+                    'segment_idx': 0,
+                    'total_segments': 1
+                })
+            else:
+                # Long audio: split into multiple segments
+                start = 0
+                segment_idx = 0
+                while start < audio_mel_len:
+                    end = min(start + frames_per_window, audio_mel_len)
+                    segment = audio_mel[:, start:end]
+
+                    seg_len = segment.size(1)
+                    # Pad if necessary
+                    if seg_len < frames_per_window:
+                        pad_size = frames_per_window - seg_len
+                        segment = F.pad(segment, (0, pad_size))
+
+                    all_segments.append(segment)
+                    all_segments_len.append(
+                        torch.tensor(seg_len, device=mel.device))
+                    segment_info.append({
+                        'batch_idx': batch_idx,
+                        'is_long_audio': True,
+                        'segment_idx': segment_idx,
+                        'total_segments': None  # Will be filled later
+                    })
+
+                    segment_idx += 1
+                    start += frames_per_stride
+
+                # Update total_segments info
+                total_segments = segment_idx
+                for info in segment_info:
+                    if info['batch_idx'] == batch_idx and info['is_long_audio']:
+                        info['total_segments'] = total_segments
+
+        if not all_segments:
+            # Fallback if no segments
+            return torch.zeros(batch_size,
+                               0,
+                               dtype=torch.long,
+                               device=mel.device), torch.zeros(
+                                   batch_size,
+                                   dtype=torch.long,
+                                   device=mel.device)
+
+        # Unified batch processing for all segments
+        unified_batch_mel = torch.stack(all_segments)
+        unified_batch_lens = torch.stack(all_segments_len)
+
+        # Process all segments at once
+        hidden, code_len = self.encoder(unified_batch_mel, unified_batch_lens)
+        codes = self.quantizer.encode(hidden)
+
+        # Reorganize results based on segment_info
+        results = {}  # batch_idx -> (code_tensor, code_len)
+
+        for seg_idx, info in enumerate(segment_info):
+            batch_idx = info['batch_idx']
+            is_long_audio = info['is_long_audio']
+            segment_idx = info['segment_idx']
+
+            # Get codes for current segment
+            segment_code = codes[
+                seg_idx, :code_len[seg_idx].item()].cpu().numpy().tolist()
+
+            if not is_long_audio:
+                # Short audio: use directly
+                code_tensor = torch.tensor(segment_code,
+                                           dtype=torch.long,
+                                           device=mel.device)
+                results[batch_idx] = (code_tensor, len(segment_code))
+            else:
+                # Long audio: collect all segments
+                if batch_idx not in results:
+                    results[batch_idx] = []
+                results[batch_idx].append(segment_code)
+
+        # Process long audio segment merging
+        for batch_idx in range(batch_size):
+            if long_audio_mask[batch_idx].item():
+                # Merge long audio segments
+                audio_codes = results[batch_idx]
+
+                # Determine token rate based on model name
+                if hasattr(self,
+                           'name') and self.name == "speech_tokenizer_v1":
+                    token_rate = 50
+                else:
+                    token_rate = 25
+
+                merged_codes = merge_tokenized_segments(audio_codes,
+                                                        overlap=overlap,
+                                                        token_rate=token_rate)
+
+                # Convert to tensor
+                merged_codes_tensor = torch.tensor(merged_codes,
+                                                   dtype=torch.long,
+                                                   device=mel.device)
+                results[batch_idx] = (merged_codes_tensor, len(merged_codes))
+
+        # Construct final output
+        max_code_len = max(code_info[1] for code_info in results.values())
+
+        output_codes = torch.zeros(batch_size,
+                                   max_code_len,
+                                   dtype=torch.long,
+                                   device=mel.device)
+        output_codes_len = torch.zeros(batch_size,
+                                       dtype=torch.long,
+                                       device=mel.device)
+
+        for batch_idx, (code_tensor, code_len) in results.items():
+            output_codes[batch_idx, :code_len] = code_tensor
+            output_codes_len[batch_idx] = code_len
+
+        return output_codes, output_codes_len
 
     @property
     def device(self):