rwkvtts-respark-webrwkv/ref_audio_utilities.py

import onnxruntime as ort
import numpy as np
import librosa
import soundfile as sf
import soxr
from pathlib import Path
from typing import Tuple, Union, Optional
import soundfile as sf


class RefAudioUtilities:
    """音频处理工具类，使用ONNX模型生成tokens"""
    
    def __init__(self, onnx_model_path: str, wav2vec2_path, 
                 ref_segment_duration: float = 6.0, latent_hop_length: int = 320):
        """
        初始化ONNX模型
        
        Args:
            onnx_model_path: ONNX模型文件路径
            wav2vec2_path: wav2vec2 ONNX模型文件路径，如果为None则不加载wav2vec2模型
            ref_segment_duration: 参考音频时长（秒）
            latent_hop_length: 潜在特征跳长度
        """
        self.ort_session = ort.InferenceSession(onnx_model_path, 
                                                providers=['CUDAExecutionProvider','CPUExecutionProvider'])
        print(f"🖥️ONNX Session actual providers: {self.ort_session.get_providers()}")
        self.sample_rate = 16000
        self.ref_segment_duration = ref_segment_duration
        self.latent_hop_length = latent_hop_length
        
        # 获取模型输入输出信息
        self.input_names = [input_info.name for input_info in self.ort_session.get_inputs()]
        self.output_names = [output_info.name for output_info in self.ort_session.get_outputs()]
        
        print(f"模型输入: {self.input_names}")
        print(f"模型输出: {self.output_names}")
        
        # 初始化wav2vec2模型
        self.wav2vec2_session = ort.InferenceSession(wav2vec2_path, 
                                                providers=['CUDAExecutionProvider','CPUExecutionProvider'])
        print(f"🖥️Wav2Vec2 Session actual providers: {self.wav2vec2_session.get_providers()}")
    def load_audio(self, audio_path: Union[str, Path], target_sr: int = 16000, 
                   volume_normalize: bool = False) -> np.ndarray:
        """
        加载音频文件，与BiCodecTokenizer保持一致
        
        Args:
            audio_path: 音频文件路径
            target_sr: 目标采样率
            volume_normalize: 是否进行音量归一化
            
        Returns:
            音频数据数组
        """
        if isinstance(audio_path, str):
            audio_path = Path(audio_path)
        
        # 使用soundfile加载音频，与BiCodecTokenizer保持一致
        audio, sr = sf.read(audio_path)
        if len(audio.shape) > 1:
            audio = audio[:, 0]  # 如果是立体声，取第一个通道
        
        # 重采样到目标采样率
        if sr != target_sr:
            audio = soxr.resample(audio, sr, target_sr, quality="VHQ")
            sr = target_sr
        
        # 音量归一化
        if volume_normalize:
            audio = self._audio_volume_normalize(audio)
        
        return audio
    
    def _audio_volume_normalize(self, audio: np.ndarray, coeff: float = 0.2) -> np.ndarray:
        """音频音量归一化"""
        # Sort the absolute values of the audio signal
        temp = np.sort(np.abs(audio))

        # If the maximum value is less than 0.1, scale the array to have a maximum of 0.1
        if temp[-1] < 0.1:
            scaling_factor = max(
                temp[-1], 1e-3
            )  # Prevent division by zero with a small constant
            audio = audio / scaling_factor * 0.1

        # Filter out values less than 0.01 from temp
        temp = temp[temp > 0.01]
        L = temp.shape[0]  # Length of the filtered array

        # If there are fewer than or equal to 10 significant values, return the audio without further processing
        if L <= 10:
            return audio

        # Compute the average of the top 10% to 1% of values in temp
        volume = np.mean(temp[int(0.9 * L) : int(0.99 * L)])

        # Normalize the audio to the target coefficient level, clamping the scale factor between 0.1 and 10
        audio = audio * np.clip(coeff / volume, a_min=0.1, a_max=10)

        # Ensure the maximum absolute value in the audio does not exceed 1
        max_value = np.max(np.abs(audio))
        if max_value > 1:
            audio = audio / max_value

        return audio
    
    def extract_mel_spectrogram(self, wav: np.ndarray, n_mels: int = 128, 
                               n_fft: int = 1024, hop_length: int = 320, 
                               win_length: int = 640) -> np.ndarray:
        """
        提取梅尔频谱图
        
        Args:
            wav: 音频数据
            n_mels: 梅尔滤波器组数量
            n_fft: FFT窗口大小
            hop_length: 帧移
            win_length: 窗口长度
            
        Returns:
            梅尔频谱图
        """
        mel_spec = librosa.feature.melspectrogram(
            y=wav,
            sr=self.sample_rate,
            n_mels=n_mels,
            n_fft=n_fft,
            hop_length=hop_length,
            win_length=win_length,
            power=1,
            norm="slaney",
            fmin=10,
        )
        
        return mel_spec
    
    def extract_wav2vec2_features(self, wav: np.ndarray) -> np.ndarray:
        """
        使用ONNX wav2vec2模型提取特征，模拟BiCodecTokenizer的行为
        
        Args:
            wav: 音频数据
            
        Returns:
            特征向量
        """
        # 检查wav2vec2模型是否已加载
        if self.wav2vec2_session is None:
            raise RuntimeError("wav2vec2模型未加载，请在初始化时提供wav2vec2_path参数")
        
        # 添加batch维度
        input_data = wav[np.newaxis, :].astype(np.float32)  # [1, sequence_length]
        
        # 运行wav2vec2推理
        # 注意：这个ONNX模型已经包含了特征提取器的预处理和多个隐藏层的组合
        inputs = {'input': input_data}
        outputs = self.wav2vec2_session.run(None, inputs)
        
        # 输出形状应该是 [1, time_steps, 1024]
        # 这个输出已经是通过选择隐藏层11, 14, 16并计算平均值得到的
        print(f'outputs: {outputs}')
        print(f'outputs: {outputs[0].shape}')
        features = outputs[0][0]  # 移除batch维度，得到 [time_steps, 1024]
        
        return features.astype(np.float32)
            
    
    
    def get_ref_clip(self, wav: np.ndarray) -> np.ndarray:
        """
        获取参考音频片段，与BiCodecTokenizer保持一致
        
        Args:
            wav: 原始音频数据
            
        Returns:
            参考音频片段
        """
        # 使用与BiCodecTokenizer相同的计算方式
        ref_segment_length = (
            int(self.sample_rate * self.ref_segment_duration)
            // self.latent_hop_length
            * self.latent_hop_length
        )
        wav_length = len(wav)
        
        if ref_segment_length > wav_length:
            # 如果音频不足指定长度，重复音频直到达到要求
            repeat_times = ref_segment_length // wav_length + 1
            wav = np.tile(wav, repeat_times)
        
        # 截取指定长度
        return wav[:ref_segment_length]
    
    def process_audio(self, audio_path: Union[str, Path], volume_normalize: bool = False) -> Tuple[np.ndarray, np.ndarray]:
        """
        处理音频文件，返回原始音频和参考音频，与BiCodecTokenizer保持一致
        
        Args:
            audio_path: 音频文件路径
            volume_normalize: 是否进行音量归一化
            
        Returns:
            (原始音频, 参考音频)
        """
        wav = self.load_audio(audio_path, volume_normalize=volume_normalize)
        ref_wav = self.get_ref_clip(wav)
        
        return wav, ref_wav
    
    def tokenize(self, audio_path: Union[str, Path]) -> Tuple[np.ndarray, np.ndarray]:
        """
        使用ONNX模型生成tokens
        
        Args:
            audio_path: 音频文件路径
            
        Returns:
            (global_tokens, semantic_tokens)
        """
        # 处理音频
        wav, ref_wav = self.process_audio(audio_path)
        
        # 提取特征
        feat = self.extract_wav2vec2_features(wav)
        ref_mel = self.extract_mel_spectrogram(ref_wav)
        
   
        # 添加batch维度
        ref_mel_input = ref_mel[np.newaxis, :, :].astype(np.float32)  # [1, 128, 301]
        feat_input = feat[np.newaxis, :, :].astype(np.float32)  # [1, feat_len, 1024]
        
        # 运行ONNX模型
        inputs = {
            'ref_wav_mel': ref_mel_input,
            'feat': feat_input
        }
        
        outputs = self.ort_session.run(self.output_names, inputs)
        
        # 解析输出
        semantic_tokens = outputs[0]  # 第一个输出
        global_tokens = outputs[1]    # 第二个输出
        
        return global_tokens, semantic_tokens
    
    def tokenize_batch(self, audio_paths: list) -> Tuple[list, list]:
        """
        批量处理音频文件
        
        Args:
            audio_paths: 音频文件路径列表
            
        Returns:
            (global_tokens_list, semantic_tokens_list)
        """
        global_tokens_list = []
        semantic_tokens_list = []
        
        for audio_path in audio_paths:
            global_tokens, semantic_tokens = self.tokenize(audio_path)
            global_tokens_list.append(global_tokens)
            semantic_tokens_list.append(semantic_tokens)
        
        return global_tokens_list, semantic_tokens_list


# 测试函数
def test_ref_audio_utilities():
    """测试RefAudioUtilities类"""
    # 初始化工具类
    onnx_model_path = '/Volumes/bigdata/models/RWKVTTS_WebRWKV/BiCodecTokenize.onnx'
    wav2vec2_path = "/Volumes/bigdata/models/RWKVTTS_WebRWKV/wav2vec2-large-xlsr-53.onnx"
    # 使用与BiCodecTokenizer相同的参数
    utilities = RefAudioUtilities(
        onnx_model_path, 
        wav2vec2_path,
        ref_segment_duration=6.0,  # 6秒参考音频
        latent_hop_length=320       # 潜在特征跳长度
    )
    
    # 测试音频文件（使用项目中的示例音频）
    test_audio_path = "demos/刘德华/dehua_zh.wav"
    
    if Path(test_audio_path).exists():
        print(f"测试音频文件: {test_audio_path}")
        
        try:
            # 生成tokens
            global_tokens, semantic_tokens = utilities.tokenize(test_audio_path)
            
            print(f"Global tokens shape: {global_tokens.shape}")
            print(f"Semantic tokens shape: {semantic_tokens.shape}")
            print(f"Global tokens: {global_tokens.flatten().tolist()}")
            print(f"Semantic tokens : {semantic_tokens.flatten().tolist()}")
            
        except Exception as e:
            print(f"处理音频时出错: {e}")
    else:
        print(f"测试音频文件不存在: {test_audio_path}")
        print("请确保测试音频文件存在")


if __name__ == "__main__":
    test_ref_audio_utilities()