models/bicodec_tokenizer/spark_tokenizer.py

# -*- coding: utf-8 -*-
# Time      :2025/3/29 10:30
# Author    :Hui Huang
import os
from typing import Literal, Optional, Tuple, Dict, Any, List, Union

import torch
import torchaudio
import torchaudio.transforms as TT
from transformers import Wav2Vec2Model, Wav2Vec2FeatureExtractor
import numpy as np
from loguru import logger
from pathlib import Path

# ----------------- 假设这些模块位于你的项目路径下 -----------------
from .utils.file import load_config
from .utils.audio import load_audio
from .models.bicodec import BiCodec
from .base_model import SparkBaseModel
from .batch_processor import AsyncBatchEngine
# ---------------------------------------------------------------

__all__ = ["SparkTokenizer"]


class SparkTokenizer:
    def __init__(
            self,
            model_path: str,
            device: Literal["cpu", "cuda", "mps"] | str = "cuda",
            attn_implementation: Optional[Literal["sdpa", "flash_attention_2", "eager"]] = "eager",
            batch_size: int = 32,
            wait_timeout: float = 0.01,
    ):
        self.device = torch.device(device)
        self.model_dir = Path(model_path)

        # 1. 加载配置
        self.config = load_config(self.model_dir / "config.yaml")
        self.device_type = "cuda" if "cuda" in str(device) else "cpu"
        self.dtype = torch.float16 if self.device_type == "cuda" else torch.float32
        self.target_sample_rate = self.config.get("sample_rate", 16000)

        # 2. 加载模型
        wav2vec_path = self.model_dir / "wav2vec2-large-xlsr-53"
        self.processor = Wav2Vec2FeatureExtractor.from_pretrained(wav2vec_path)
        self.feature_extractor = Wav2Vec2Model.from_pretrained(
            wav2vec_path,
            attn_implementation=attn_implementation,
            torch_dtype=self.dtype 
        )
        self.feature_extractor.config.output_hidden_states = True
        self.feature_extractor.to(self.device)
        self.feature_extractor.eval()

        # BiCodec model
        self.model = (
            BiCodec.load_from_checkpoint(str(self.model_dir)).to(self.device).half()
        )
        self.model.eval()

        # 异步处理引擎
        self._batch_processor = AsyncBatchEngine(
            processing_function=self.batch_tokenize_async,
            batch_size=batch_size,
            wait_timeout=wait_timeout
        )

    def _to_ndarray(self, audio_input: Union[str, Path, torch.Tensor]) -> np.ndarray:
        """
        将输入（路径或Tensor）统一转换为指定采样率的 numpy 数组。
        """
        if isinstance(audio_input, (str, Path)):
            # 如果是路径，直接使用原有的 load_audio
            wav = load_audio(
                str(audio_input),
                sampling_rate=self.target_sample_rate,
                volume_normalize=self.config.get("volume_normalize", True),
            )
        elif isinstance(audio_input, torch.Tensor):
            # 如果是 Tensor
            wav = audio_input.detach().cpu().float()

            # 处理通道: [C, T] -> [T]
            if wav.ndim > 1:
                wav = torch.mean(wav, dim=0)

            # 这里默认输入的 Tensor 采样率已经是 self.target_sample_rate
            # 如果需要在这里做重采样，需要额外传入输入采样率参数
            wav = wav.numpy()

            # 可选：音量归一化逻辑（如果 Tensor 没归一化）
            if self.config.get("volume_normalize", True):
                max_val = np.abs(wav).max()
                if max_val > 0:
                    wav = wav / max_val * 0.9
        else:
            raise ValueError(f"Unsupported audio type: {type(audio_input)}")

        return wav

    def get_ref_clip(self, wav: np.ndarray) -> np.ndarray:
        """获取参考音频片段"""
        ref_segment_length = (
            int(self.target_sample_rate * self.config["ref_segment_duration"])
            // self.config["latent_hop_length"]
            * self.config["latent_hop_length"]
        )
        wav_length = len(wav)

        if ref_segment_length > wav_length:
            wav = np.tile(wav, ref_segment_length // wav_length + 1)

        return wav[:ref_segment_length]

    def process_audio(self, audio_input: Union[str, torch.Tensor], ref_audio_input: Union[str, torch.Tensor] = None) -> Tuple[np.ndarray, torch.Tensor]:
        """
        处理音频和参考音频。
        """
        wav = self._to_ndarray(audio_input)

        if ref_audio_input is None:
            wav_ref_np = self.get_ref_clip(wav)
        else:
            ref_wav = self._to_ndarray(ref_audio_input)
            wav_ref_np = self.get_ref_clip(ref_wav)

        wav_ref = torch.from_numpy(wav_ref_np).unsqueeze(0).float()
        return wav, wav_ref

    def extract_wav2vec2_features(self, wavs: torch.Tensor) -> torch.Tensor:
        """提取 wav2vec2 特征"""
        # processor 期望是 list of numpy
        inputs = self.processor(
            [w.cpu().numpy() for w in wavs], 
            sampling_rate=16000,
            return_tensors="pt",
            padding=True,
        ).input_values

        with torch.no_grad():
            with torch.amp.autocast(self.device_type, dtype=self.dtype):
                feat = self.feature_extractor(inputs.to(self.feature_extractor.device))

        feats_mix = (
            feat.hidden_states[11] + feat.hidden_states[14] + feat.hidden_states[16]
        ) / 3

        return feats_mix

    @torch.no_grad()
    def tokenize(self, audios: List[Union[str, torch.Tensor]]):
        """
        支持音频路径列表或 Tensor 列表。
        """
        batch_wavs = []
        batch_ref_wavs = []

        for audio_item in audios:
            wav, wav_ref = self.process_audio(audio_input=audio_item, ref_audio_input=audio_item)
            batch_wavs.append(torch.from_numpy(wav).float())
            batch_ref_wavs.append(wav_ref.squeeze(0))

        # Padding wavs
        wav_lengths = [len(w) for w in batch_wavs]
        max_wav_len = max(wav_lengths)
        padded_wavs = torch.zeros(len(batch_wavs), max_wav_len, dtype=self.dtype).to(self.device)
        for i, w in enumerate(batch_wavs):
            padded_wavs[i, :len(w)] = w.to(self.dtype)

        # Padding ref_wavs
        ref_lengths = [len(w) for w in batch_ref_wavs]
        max_ref_len = max(ref_lengths)
        padded_ref_wavs = torch.zeros(len(batch_ref_wavs), max_ref_len, dtype=self.dtype).to(self.device)
        for i, w in enumerate(batch_ref_wavs):
            padded_ref_wavs[i, :len(w)] = w.to(self.dtype)

        # 提取特征
        feats = self.extract_wav2vec2_features(padded_wavs)

        batch = {
            "wav": padded_wavs,
            "ref_wav": padded_ref_wavs,
            "feat": feats,
        }

        semantic_tokens, global_tokens = self.model.tokenize(batch)

        if self.device.type == "cuda":
            torch.cuda.empty_cache()

        return {"semantic_tokens": semantic_tokens, "global_tokens": global_tokens}

    async def batch_tokenize_async(self, audios: list) -> list[dict[str, torch.Tensor]]:
        tokenized = self.tokenize(audios)
        responses = []
        for i in range(len(audios)):
            responses.append({
                "global_tokens": tokenized["global_tokens"][i],
                "semantic_tokens": tokenized["semantic_tokens"][i]
            })
        return responses

    async def tokenize_async(self, audio: Union[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        output = await self._batch_processor.add_request(
            single_input=audio
        )
        return output

# ------------------------------------------------------------------
# 测试用例
# ------------------------------------------------------------------
if __name__ == "__main__":
    # 配置你的模型路径
    MODEL_DIR = "/data/yumu/model/ark_tts_v1"
    
    # 初始化
    # 注意：在没有真实环境时，这行会因为找不到文件报错，请在有环境的地方运行
    tokenizer = SparkTokenizer(model_path=MODEL_DIR, device="cuda" if torch.cuda.is_available() else "cpu")
    
    # 准备数据：一个是本地存在的 wav 路径，一个是构造的 Tensor
    dummy_wav_path = "/data/yumu/arktts/dufu.wav" 
    # 构造一个 16kHz 的 2 秒音频 Tensor (假设模型要求16k)
    import torchaudio
    dummy_tensor, sr = torchaudio.load(dummy_wav_path)

    # 1. 测试路径输入
    print("Testing path input...")
    if os.path.exists(dummy_wav_path):
        res1 = tokenizer.tokenize([dummy_wav_path])
        print(f"Path results: {res1['semantic_tokens'].shape}")

    # 2. 测试 Tensor 输入
    print("Testing tensor input...")
    res2 = tokenizer.tokenize([dummy_tensor])
    print(f"Tensor results: {res2['semantic_tokens'].shape}")

    # 3. 测试混合输入 (List 包含 str 和 Tensor)
    print("Testing mixed input...")
    # 为了演示，我们传两个相同的 tensor
    res3 = tokenizer.tokenize([dummy_tensor, dummy_tensor])
    print(f"Mixed results: {res3['semantic_tokens'].shape}")
    
    print("All tests passed!")