minicpmo-utils 0.1.0__py3-none-any.whl

stepaudio2/token2wav.py

@@ -0,0 +1,248 @@
+import io
+import sys
+import types
+
+import torch
+import torchaudio
+import s3tokenizer
+import onnxruntime
+import numpy as np
+from copy import deepcopy
+
+import torchaudio.compliance.kaldi as kaldi
+from stepaudio2.flashcosyvoice.modules.hifigan import HiFTGenerator
+from stepaudio2.flashcosyvoice.utils.audio import mel_spectrogram
+from hyperpyyaml import load_hyperpyyaml
+
+
+def _setup_cosyvoice2_alias():
+    """给 hyperpyyaml 提供 cosyvoice2.* 的兼容别名，不改 flow.yaml。"""
+    if 'cosyvoice2.flow.flow' in sys.modules:
+        # 已经设置过，直接复用
+        return
+
+    # 导入 stepaudio2 里真实的实现
+    import stepaudio2.cosyvoice2.flow.flow as _step_flow
+    import stepaudio2.cosyvoice2.flow.flow_matching as _step_flow_matching
+    import stepaudio2.cosyvoice2.flow.decoder_dit as _step_decoder_dit
+    import stepaudio2.cosyvoice2.transformer.upsample_encoder_v2 as _step_upsample
+
+    # 创建顶层 cosyvoice2 包和子包
+    cosyvoice2_pkg = types.ModuleType('cosyvoice2')
+    cosyvoice2_flow_pkg = types.ModuleType('cosyvoice2.flow')
+    cosyvoice2_transformer_pkg = types.ModuleType('cosyvoice2.transformer')
+
+    # 挂载子模块
+    cosyvoice2_flow_pkg.flow = _step_flow
+    cosyvoice2_flow_pkg.flow_matching = _step_flow_matching
+    cosyvoice2_flow_pkg.decoder_dit = _step_decoder_dit
+
+    cosyvoice2_transformer_pkg.upsample_encoder_v2 = _step_upsample
+
+    cosyvoice2_pkg.flow = cosyvoice2_flow_pkg
+    cosyvoice2_pkg.transformer = cosyvoice2_transformer_pkg
+
+    # 注册到 sys.modules，让 `cosyvoice2.flow.flow.*` 这类路径可被 import
+    sys.modules['cosyvoice2'] = cosyvoice2_pkg
+    sys.modules['cosyvoice2.flow'] = cosyvoice2_flow_pkg
+    sys.modules['cosyvoice2.flow.flow'] = _step_flow
+    sys.modules['cosyvoice2.flow.flow_matching'] = _step_flow_matching
+    sys.modules['cosyvoice2.flow.decoder_dit'] = _step_decoder_dit
+    sys.modules['cosyvoice2.transformer'] = cosyvoice2_transformer_pkg
+    sys.modules['cosyvoice2.transformer.upsample_encoder_v2'] = _step_upsample
+
+def fade_in_out(fade_in_mel:torch.Tensor, fade_out_mel:torch.Tensor, window:torch.Tensor):
+    """perform fade_in_out in tensor style
+    """
+    mel_overlap_len = int(window.shape[0] / 2)
+    fade_in_mel = fade_in_mel.clone()
+    fade_in_mel[..., :mel_overlap_len] = \
+        fade_in_mel[..., :mel_overlap_len] * window[:mel_overlap_len] + \
+        fade_out_mel[..., -mel_overlap_len:] * window[mel_overlap_len:]
+    return fade_in_mel
+
+
+class Token2wav():
+    def __init__(self, model_path, float16=False, n_timesteps=10):
+        self.float16 = float16
+        self.n_timesteps = n_timesteps
+
+        # 在加载 flow.yaml 之前，先把 cosyvoice2 的别名注册好
+        _setup_cosyvoice2_alias()
+
+        self.audio_tokenizer = s3tokenizer.load_model(f"{model_path}/speech_tokenizer_v2_25hz.onnx").cuda().eval()
+
+        option = onnxruntime.SessionOptions()
+        option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+        option.intra_op_num_threads = 1
+        self.spk_model = onnxruntime.InferenceSession(f"{model_path}/campplus.onnx", sess_options=option, providers=["CPUExecutionProvider"])
+
+        with open(f"{model_path}/flow.yaml", "r") as f:
+            configs = load_hyperpyyaml(f)
+            self.flow = configs['flow']
+        if float16:
+            self.flow.half()
+        self.flow.load_state_dict(torch.load(f"{model_path}/flow.pt", map_location="cpu", weights_only=True), strict=True)
+        self.flow.cuda().eval()
+
+        self.hift = HiFTGenerator()
+        hift_state_dict = {k.replace('generator.', ''): v for k, v in torch.load(f"{model_path}/hift.pt", map_location="cpu", weights_only=True).items()}
+        self.hift.load_state_dict(hift_state_dict, strict=True)
+        self.hift.cuda().eval()
+
+        self.cache = None
+
+        # stream conf
+        self.mel_cache_len = 8  # hard-coded, 160ms
+        self.source_cache_len = int(self.mel_cache_len * 480)   # 50hz mel -> 24kHz wave
+        self.speech_window = torch.from_numpy(np.hamming(2 * self.source_cache_len)).cuda()
+
+        # hifigan cache
+        self.hift_cache_dict = {}
+
+    def _prepare_prompt(self, prompt_wav):
+        audio = s3tokenizer.load_audio(prompt_wav, sr=16000)  # [T]
+        # TODO 在audio 后面 pad 3个token长度 = 0.12s的音频
+        mels = s3tokenizer.log_mel_spectrogram(audio)
+        mels, mels_lens = s3tokenizer.padding([mels])
+        prompt_speech_tokens, prompt_speech_tokens_lens = self.audio_tokenizer.quantize(mels.cuda(), mels_lens.cuda())
+        # TODO 手动赋值最后3个token为 4218 静音token
+
+        spk_feat = kaldi.fbank(audio.unsqueeze(0), num_mel_bins=80, dither=0, sample_frequency=16000)
+        spk_feat = spk_feat - spk_feat.mean(dim=0, keepdim=True)
+        spk_emb = torch.tensor(self.spk_model.run(
+            None, {self.spk_model.get_inputs()[0].name: spk_feat.unsqueeze(dim=0).cpu().numpy()}
+        )[0], device='cuda')
+
+        audio, sample_rate = torchaudio.load(prompt_wav, backend='soundfile')
+        audio = audio.mean(dim=0, keepdim=True)  # [1, T]
+        if sample_rate != 24000:
+            audio = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=24000)(audio)
+        prompt_mel = mel_spectrogram(audio).transpose(1, 2).squeeze(0)  # [T, num_mels]
+        prompt_mels = prompt_mel.unsqueeze(0).cuda()
+        prompt_mels_lens = torch.tensor([prompt_mels.shape[1]], dtype=torch.int32, device='cuda')
+        prompt_mels = torch.nn.functional.pad(prompt_mels, (0, 0, 0, prompt_speech_tokens.shape[1] * self.flow.up_rate - prompt_mels.shape[1]), mode='replicate')
+        return prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens
+
+    def __call__(self, generated_speech_tokens, prompt_wav):
+        if self.cache is None:
+            self.cache = self._prepare_prompt(prompt_wav)
+        prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens = self.cache
+
+        generated_speech_tokens = torch.tensor([generated_speech_tokens], dtype=torch.int32, device='cuda')
+        generated_speech_tokens_lens = torch.tensor([generated_speech_tokens.shape[1]], dtype=torch.int32, device='cuda')
+
+        with torch.amp.autocast("cuda", dtype=torch.float16 if self.float16 else torch.float32):
+            mel = self.flow.inference(generated_speech_tokens, generated_speech_tokens_lens,
+                prompt_speech_tokens, prompt_speech_tokens_lens,
+                prompt_mels, prompt_mels_lens, spk_emb, 10)
+
+        wav, _ = self.hift(speech_feat=mel)
+        output = io.BytesIO()
+        torchaudio.save(output, wav.cpu(), sample_rate=24000, format='wav')
+
+        return output.getvalue()
+
+    def set_stream_cache(self, prompt_wav):
+        self.cache = self._prepare_prompt(prompt_wav)
+        prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens = self.cache
+
+        print("prompt_speech_tokens_lens", prompt_speech_tokens_lens)
+        print("prompt_mels.shape", prompt_mels.shape)
+        print("prompt_mels_lens", prompt_mels_lens)
+        
+        right_pad_speech_tokens = torch.ones(1, 3, device=prompt_speech_tokens.device, dtype=prompt_speech_tokens.dtype) * 4218
+        
+        # self.stream_cache = self.flow.setup_cache(
+        #     torch.cat([prompt_speech_tokens, prompt_speech_tokens[:, :3]], dim=1),
+        #     prompt_mels, spk_emb, n_timesteps=self.n_timesteps)
+        
+        stream_cache = self.flow.setup_cache(
+            torch.cat([prompt_speech_tokens, right_pad_speech_tokens], dim=1),
+            prompt_mels, spk_emb, n_timesteps=self.n_timesteps)
+
+        # hift cache
+        hift_cache_dict = dict(
+            mel = torch.zeros(1, prompt_mels.shape[2], 0, device='cuda'), 
+            source = torch.zeros(1, 1, 0, device='cuda'),
+            speech = torch.zeros(1, 0, device='cuda'),
+        )
+
+        return stream_cache, hift_cache_dict
+
+
+    def stream(self, generated_speech_tokens, prompt_wav, last_chunk=False, return_waveform=False):
+        if self.cache is None:
+            self.cache = self._prepare_prompt(prompt_wav)
+        prompt_speech_tokens, prompt_speech_tokens_lens, spk_emb, prompt_mels, prompt_mels_lens = self.cache
+
+        generated_speech_tokens = torch.tensor([generated_speech_tokens], dtype=torch.int32, device='cuda')
+        generated_speech_tokens_lens = torch.tensor([generated_speech_tokens.shape[1]], dtype=torch.int32, device='cuda')
+
+        if self.stream_cache is None:
+            raise ValueError("stream_cache is not set")
+
+        with torch.amp.autocast("cuda", dtype=torch.float16 if self.float16 else torch.float32):
+            chunk_mel, self.stream_cache = self.flow.inference_chunk(
+                token=generated_speech_tokens,
+                spk=spk_emb,
+                cache=self.stream_cache,
+                last_chunk=last_chunk,
+                n_timesteps=self.n_timesteps,
+            )
+        if self.stream_cache['estimator_att_cache'].shape[4] > (prompt_mels.shape[1] + 100):
+            self.stream_cache['estimator_att_cache'] = torch.cat([
+                self.stream_cache['estimator_att_cache'][:, :, :, :, :prompt_mels.shape[1]],
+                self.stream_cache['estimator_att_cache'][:, :, :, :, -100:],
+            ], dim=4)
+
+        # bug fix - 20260107
+        # 同样截断 conformer_att_cache，防止无限增长导致 position embedding 不匹配
+        if self.stream_cache['conformer_att_cache'].shape[3] > (prompt_mels.shape[1] + 100):
+            self.stream_cache['conformer_att_cache'] = torch.cat([
+                self.stream_cache['conformer_att_cache'][:, :, :, :prompt_mels.shape[1], :],
+                self.stream_cache['conformer_att_cache'][:, :, :, -100:, :],
+            ], dim=3)
+
+        # vocoder cache
+        hift_cache_mel = self.hift_cache_dict['mel']
+        hift_cache_source = self.hift_cache_dict['source']
+        hift_cache_speech = self.hift_cache_dict['speech']
+        mel = torch.concat([hift_cache_mel, chunk_mel], dim=2)
+
+        speech, source = self.hift(mel, hift_cache_source)
+
+        # overlap speech smooth
+        if hift_cache_speech.shape[-1] > 0:
+            speech = fade_in_out(speech, hift_cache_speech, self.speech_window)
+
+        # 检测是否是第一个 chunk（没有有效的 speech cache）
+        is_first_chunk = hift_cache_speech.shape[-1] == 0
+
+        # update vocoder cache
+        self.hift_cache_dict = dict(
+            mel = mel[..., -self.mel_cache_len:].clone().detach(),
+            source = source[:, :, -self.source_cache_len:].clone().detach(),
+            speech = speech[:, -self.source_cache_len:].clone().detach(),
+        )
+
+        if not last_chunk:
+            if is_first_chunk:
+                # 第一个 chunk：截断尾部，在开头添加静音来补偿长度
+                # 不做淡入，保持原始音频内容不变
+                silence_padding = torch.zeros(1, self.source_cache_len, device=speech.device)
+                speech = torch.cat([silence_padding, speech[:, :-self.source_cache_len]], dim=1)
+            else:
+                # 后续 chunk：正常截断尾部（由下一个 chunk 的 fade_in 补偿）
+                speech = speech[:, :-self.source_cache_len]
+
+        wav_np = speech.cpu().numpy()
+        if return_waveform:
+            return wav_np
+
+        # Clip to [-1, 1] to avoid overflow, then scale to int16
+        wav_np = np.clip(wav_np, -1.0, 1.0)
+        wav_int16 = (wav_np * 32767.0).astype('<i2')  # 16-bit little-endian PCM
+        pcm_bytes = wav_int16.tobytes()
+        return pcm_bytes
+

stepaudio2/utils.py

@@ -0,0 +1,91 @@
+import librosa
+import torch
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
+import torchaudio
+from typing import List
+
+
+def _mel_filters(n_mels: int) -> torch.Tensor:
+    """Load the mel filterbank matrix for projecting STFT into a Mel spectrogram."""
+    assert n_mels in {80, 128}, f"Unsupported n_mels: {n_mels}"
+    if n_mels == 128:
+        return torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=128))
+    else:
+        return torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=80))
+
+def load_audio(file_path, target_rate=16000, max_length=None):
+    """
+    Open an audio file and read as mono waveform, resampling as necessary
+    If max_length is provided, truncate the audio to that length
+    """
+    waveform, sample_rate = torchaudio.load(file_path)
+    if sample_rate != target_rate:
+        waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_rate)(waveform)
+    audio = waveform[0]  # get the first channel
+
+    # Truncate audio if it exceeds max_length
+    if max_length is not None and audio.shape[0] > max_length:
+        audio = audio[:max_length]
+
+    return audio
+
+def log_mel_spectrogram(audio, n_mels=128, padding=479, device=None):
+    """
+    Compute the log-Mel spectrogram with specific padding for StepAudio
+    """
+    if not torch.is_tensor(audio):
+        if isinstance(audio, str):
+            audio = load_audio(audio)
+        audio = torch.from_numpy(audio)
+    if device is not None:
+        audio = audio.to(device)
+    if padding > 0:
+        audio = F.pad(audio, (0, padding))
+    window = torch.hann_window(400).to(audio.device)
+    stft = torch.stft(audio, 400, 160, window=window, return_complex=True)
+    magnitudes = stft[..., :-1].abs() ** 2
+    filters = _mel_filters(n_mels)
+    mel_spec = filters @ magnitudes
+
+    log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+    log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
+    log_spec = (log_spec + 4.0) / 4.0
+    return log_spec
+
+def compute_token_num(max_feature_len):
+    # First, audio goes through encoder:
+    # 1. conv1: kernel=3, stride=1, padding=1 -> size unchanged
+    # 2. conv2: kernel=3, stride=2, padding=1 -> size/2
+    # 3. avg_pooler: kernel=2, stride=2 -> size/2
+    max_feature_len = max_feature_len - 2  # remove padding
+    encoder_output_dim = (max_feature_len + 1) // 2 // 2  # after conv2 and avg_pooler
+
+    # Then through adaptor (parameters from config file):
+    padding = 1
+    kernel_size = 3  # from config: audio_encoder_config.kernel_size
+    stride = 2      # from config: audio_encoder_config.adapter_stride
+    adapter_output_dim = (encoder_output_dim + 2 * padding - kernel_size) // stride + 1
+    return adapter_output_dim
+
+def padding_mels(data: List[torch.Tensor]):
+    """ Padding the data into batch data
+
+    Parameters
+    ----------
+        data: List[Tensor], shape of Tensor (128, T)
+
+    Returns:
+    -------
+        feats, feats lengths
+    """
+    sample = data
+    assert isinstance(sample, list)
+    feats_lengths = torch.tensor([s.size(1)-2 for s in sample],
+                                dtype=torch.int32)
+    feats = [s.t() for s in sample]
+    padded_feats = pad_sequence(feats,
+                                batch_first=True,
+                                padding_value=0)
+
+    return padded_feats.transpose(1, 2), feats_lengths