tts.py

import hashlib
import io
import os
import re
import logging
import numpy as np
import torch
import librosa
import soundfile as sf
from typing import Tuple, Optional
from http import HTTPStatus

import torchaudio

from model_loader import model_loader, ModelSource
from config.prompts import AUDIO_EDIT_CLONE_SYSTEM_PROMPT_TPL, AUDIO_EDIT_SYSTEM_PROMPT
from stepvocoder.cosyvoice2.cli.cosyvoice import CosyVoice
from transformers.generation.logits_process import LogitsProcessor
from transformers.generation.utils import LogitsProcessorList

# Configure logging
logger = logging.getLogger(__name__)


class HTTPException(Exception):
    """Custom HTTP exception for API errors"""
    def __init__(self, status_code, detail):
        self.status_code = status_code
        self.detail = detail
        super().__init__(detail)


class RepetitionAwareLogitsProcessor(LogitsProcessor):
    """Logits processor to handle repetition in generation"""
    def __call__(
        self, input_ids: torch.LongTensor, scores: torch.FloatTensor
    ) -> torch.FloatTensor:
        window_size = 10
        threshold = 0.1

        window = input_ids[:, -window_size:]
        if window.shape[1] < window_size:
            return scores

        last_tokens = window[:, -1].unsqueeze(-1)
        repeat_counts = (window == last_tokens).sum(dim=1)
        repeat_ratios = repeat_counts.float() / window_size

        mask = repeat_ratios > threshold
        scores[mask, last_tokens[mask].squeeze(-1)] = float("-inf")
        return scores

class StepAudioTTS:
    """
    Step Audio TTS wrapper for voice cloning and audio editing tasks
    """

    def __init__(
        self,
        model_path,
        audio_tokenizer,
        model_source=ModelSource.AUTO,
        tts_model_id=None,
        quantization_config=None,
        torch_dtype=torch.bfloat16,
        device_map="cuda"
    ):
        """
        Initialize StepAudioTTS

        Args:
            model_path: Model path
            audio_tokenizer: Audio tokenizer for wav2token processing
            model_source: Model source (auto/local/modelscope/huggingface)
            tts_model_id: TTS model ID, if None use model_path
            quantization_config: Quantization configuration ('int4', 'int8', or None)
            torch_dtype: PyTorch data type for model weights (default: torch.bfloat16)
            device_map: Device mapping for model (default: "cuda")
        """
        # Determine model ID or path to load
        if tts_model_id is None:
            tts_model_id = model_path

        logger.info("🔧 StepAudioTTS loading configuration:")
        logger.info(f"   - model_source: {model_source}")
        logger.info(f"   - model_path: {model_path}")
        logger.info(f"   - tts_model_id: {tts_model_id}")

        self.audio_tokenizer = audio_tokenizer

        # Load LLM and tokenizer using model_loader
        try:
            self.llm, self.tokenizer, model_path = model_loader.load_transformers_model(
                tts_model_id,
                source=model_source,
                quantization_config=quantization_config,
                torch_dtype=torch_dtype,
                device_map=device_map
            )
            logger.info(f"✅ Successfully loaded LLM and tokenizer: {tts_model_id}")
        except Exception as e:
            logger.error(f"❌ Failed to load model: {e}")
            raise

        # Load CosyVoice model (usually local path)
        self.cosy_model = CosyVoice(
            os.path.join(model_path, "CosyVoice-300M-25Hz")
        )

        # Print final GPU memory usage after all models are loaded
        logger.info("🎤 CosyVoice model loaded successfully")

        # Use system prompts from config module
        self.edit_clone_sys_prompt_tpl = AUDIO_EDIT_CLONE_SYSTEM_PROMPT_TPL
        self.edit_sys_prompt = AUDIO_EDIT_SYSTEM_PROMPT

    def clone(
        self,
        prompt_wav_path: str,
        prompt_text: str,
        target_text: str
    ) -> Tuple[torch.Tensor, int]:
        """
        Clone voice from reference audio

        Args:
            prompt_wav_path: Path to reference audio file
            prompt_text: Text content of reference audio
            target_text: Text to synthesize with cloned voice

        Returns:
            Tuple[torch.Tensor, int]: Generated audio tensor and sample rate
        """
        try:
            logger.debug(f"Starting voice cloning: {prompt_wav_path}")
            prompt_wav, _ = torchaudio.load(prompt_wav_path)
            vq0206_codes, vq02_codes_ori, vq06_codes_ori, speech_feat, _, speech_embedding = (
                self.preprocess_prompt_wav(prompt_wav_path)
            )
            prompt_speaker = self.generate_clone_voice_id(prompt_text, prompt_wav)
            prompt_wav_tokens = self.audio_tokenizer.merge_vq0206_to_token_str(
                vq02_codes_ori, vq06_codes_ori
            )
            token_ids = self._encode_audio_edit_clone_prompt(
                target_text,
                prompt_text,
                prompt_speaker,
                prompt_wav_tokens,
            )

            output_ids = self.llm.generate(
                torch.tensor([token_ids]).to(torch.long).to("cuda"),
                max_length=8192,
                temperature=0.7,
                do_sample=True,
                logits_processor=LogitsProcessorList([RepetitionAwareLogitsProcessor()]),
            )
            output_ids = output_ids[:, len(token_ids) : -1]  # skip eos token
            logger.debug("Voice cloning generation completed")
            vq0206_codes_vocoder = torch.tensor([vq0206_codes], dtype=torch.long) - 65536
            return (
                self.cosy_model.token2wav_nonstream(
                    output_ids - 65536,
                    vq0206_codes_vocoder,
                    speech_feat.to(torch.bfloat16),
                    speech_embedding.to(torch.bfloat16),
                ),
                24000,
            )
        except Exception as e:
            logger.error(f"Clone failed: {e}")
            raise

    def edit(
        self,
        input_audio_path: str,
        audio_text: str,
        edit_type: str,
        edit_info: Optional[str] = None,
        text: Optional[str] = None
    ) -> Tuple[torch.Tensor, int]:
        """
        Edit audio based on specified edit type

        Args:
            input_audio_path: Path to input audio file
            audio_text: Text content of input audio
            edit_type: Type of edit (emotion, style, speed, etc.)
            edit_info: Specific edit information (happy, sad, etc.)
            text: Target text for para-linguistic editing

        Returns:
            Tuple[torch.Tensor, int]: Edited audio tensor and sample rate
        """
        try:
            logger.debug(f"Starting audio editing: {edit_type} - {edit_info}")            
            vq0206_codes, vq02_codes_ori, vq06_codes_ori, speech_feat, _, speech_embedding = (
                self.preprocess_prompt_wav(input_audio_path)
            )
            audio_tokens = self.audio_tokenizer.merge_vq0206_to_token_str(
                vq02_codes_ori, vq06_codes_ori
            )
            # Build instruction prefix based on edit type
            instruct_prefix = self._build_audio_edit_instruction(audio_text, edit_type, edit_info, text)

            # Encode the complete prompt to token sequence
            prompt_tokens = self._encode_audio_edit_prompt(
                self.edit_sys_prompt, instruct_prefix, audio_tokens
            )

            logger.debug(f"Edit instruction: {instruct_prefix}")
            logger.debug(f"Encoded prompt length: {len(prompt_tokens)}")

            output_ids = self.llm.generate(
                torch.tensor([prompt_tokens]).to(torch.long).to("cuda"),
                max_length=8192,
                temperature=0.7,
                do_sample=True,
                logits_processor=LogitsProcessorList([RepetitionAwareLogitsProcessor()]),
            )
            output_ids = output_ids[:, len(prompt_tokens) : -1]  # skip eos token
            vq0206_codes_vocoder = torch.tensor([vq0206_codes], dtype=torch.long) - 65536
            logger.debug("Audio editing generation completed")
            return (
                self.cosy_model.token2wav_nonstream(
                    output_ids - 65536,
                    vq0206_codes_vocoder,
                    speech_feat.to(torch.bfloat16),
                    speech_embedding.to(torch.bfloat16),
                ),
                24000,
            )
        except Exception as e:
            logger.error(f"Edit failed: {e}")
            raise

    def _build_audio_edit_instruction(
        self,
        audio_text: str,
        edit_type: str,
        edit_info: Optional[str] = None,
        text: Optional[str] = None
        ) -> str:
        """
        Build audio editing instruction based on request

        Args:
            audio_text: Text content of input audio
            edit_type: Type of edit
            edit_info: Specific edit information
            text: Target text for editing

        Returns:
            str: Instruction prefix
        """

        audio_text = audio_text.strip() if audio_text else ""
        if edit_type in {"emotion", "speed"}:
            if edit_info == "remove":
                instruct_prefix = f"Remove any emotion in the following audio and the reference text is: {audio_text}\n"
            else:
                instruct_prefix=f"Make the following audio more {edit_info}. The text corresponding to the audio is: {audio_text}\n"
        elif edit_type == "style":
            if edit_info == "remove":
                instruct_prefix = f"Remove any speaking styles in the following audio and the reference text is: {audio_text}\n"
            else:
                instruct_prefix = f"Make the following audio more {edit_info} style. The text corresponding to the audio is: {audio_text}\n"
        elif edit_type == "denoise":
            instruct_prefix = f"Remove any noise from the given audio while preserving the voice content clearly. Ensure that the speech quality remains intact with minimal distortion, and eliminate all noise from the audio.\n"
        elif edit_type == "vad":
            instruct_prefix = f"Remove any silent portions from the given audio while preserving the voice content clearly. Ensure that the speech quality remains intact with minimal distortion, and eliminate all silence from the audio.\n"
        elif edit_type == "paralinguistic":
            instruct_prefix = f"Add some non-verbal sounds to make the audio more natural, the new text is : {text}\n  The text corresponding to the audio is: {audio_text}\n"
        else:
            raise HTTPException(
                status_code=HTTPStatus.BAD_REQUEST,
                detail=f"Unsupported edit_type: {edit_type}",
            )

        return instruct_prefix

    def _encode_audio_edit_prompt(
        self, sys_prompt: str, instruct_prefix: str, audio_token_str: str
    ) -> list[int]:
        """
        Encode audio edit prompt to token sequence

        Args:
            sys_prompt: System prompt
            instruct_prefix: Instruction prefix
            audio_token_str: Audio tokens as string

        Returns:
            list[int]: Encoded token sequence
        """
        audio_token_str = audio_token_str.strip()
        history = [1]
        sys_tokens = self.tokenizer.encode(f"system\n{sys_prompt}")
        history.extend([4] + sys_tokens + [3])
        qrole_toks = self.tokenizer.encode("human\n")
        arole_toks = self.tokenizer.encode("assistant\n")
        human_turn_toks = self.tokenizer.encode(
            f"{instruct_prefix}\n{audio_token_str}\n"
        )
        history.extend([4] + qrole_toks + human_turn_toks + [3] + [4] + arole_toks)
        return history
    
    def _encode_audio_edit_clone_prompt(
        self, text: str, prompt_text: str, prompt_speaker: str, prompt_wav_tokens: str
    ):
        prompt = self.edit_clone_sys_prompt_tpl.format(
            speaker=prompt_speaker,
            prompt_text=prompt_text,
            prompt_wav_tokens=prompt_wav_tokens
        )
        sys_tokens = self.tokenizer.encode(f"system\n{prompt}")

        history = [1]
        history.extend([4] + sys_tokens + [3])

        _prefix_tokens = self.tokenizer.encode("\n")
        
        target_token_encode = self.tokenizer.encode("\n" + text)
        target_tokens = target_token_encode[len(_prefix_tokens) :]

        qrole_toks = self.tokenizer.encode("human\n")
        arole_toks = self.tokenizer.encode("assistant\n")

        history.extend(
            [4]
            + qrole_toks
            + target_tokens
            + [3]
            + [4]
            + arole_toks
        )
        return history


    def detect_instruction_name(self, text):
        instruction_name = ""
        match_group = re.match(r"^([（\(][^\(\)()]*[）\)]).*$", text, re.DOTALL)
        if match_group is not None:
            instruction = match_group.group(1)
            instruction_name = instruction.strip("()（）")
        return instruction_name

    def process_audio_file(self, audio_path: str) -> Tuple[any, int]:
        """
        Process audio file and return numpy array and sample rate

        Args:
            audio_path: Path to audio file

        Returns:
            Tuple[numpy.ndarray, int]: Audio data and sample rate
        """
        try:
            audio_data, sample_rate = librosa.load(audio_path)
            logger.debug(f"Audio file processed successfully: {audio_path}")
            return audio_data, sample_rate
        except Exception as e:
            logger.error(f"Failed to process audio file: {e}")
            raise

    def preprocess_prompt_wav(self, prompt_wav_path : str):
        prompt_wav, prompt_wav_sr = torchaudio.load(prompt_wav_path)
        if prompt_wav.shape[0] > 1:
            prompt_wav = prompt_wav.mean(dim=0, keepdim=True)  # 将多通道音频转换为单通道

        # volume-normalize avoid clipping
        norm = torch.max(torch.abs(prompt_wav), dim=1, keepdim=True)[0]
        if norm > 0.6: # hard code;  max absolute value is 0.6
            prompt_wav = prompt_wav / norm * 0.6 

        speech_feat, speech_feat_len = self.cosy_model.frontend.extract_speech_feat(
            prompt_wav, prompt_wav_sr
        )
        speech_embedding = self.cosy_model.frontend.extract_spk_embedding(
            prompt_wav, prompt_wav_sr
        )
        vq0206_codes, vq02_codes_ori, vq06_codes_ori = self.audio_tokenizer.wav2token(prompt_wav, prompt_wav_sr)
        return (
            vq0206_codes,
            vq02_codes_ori,
            vq06_codes_ori,
            speech_feat,
            speech_feat_len,
            speech_embedding,
        )
        
    def generate_clone_voice_id(self, prompt_text, prompt_wav):
        hasher = hashlib.sha256()
        hasher.update(prompt_text.encode('utf-8'))
        wav_data = prompt_wav.cpu().numpy()
        if wav_data.size > 2000:
            audio_sample = np.concatenate([wav_data.flatten()[:1000], wav_data.flatten()[-1000:]])
        else:
            audio_sample = wav_data.flatten()
        hasher.update(audio_sample.tobytes())
        voice_hash = hasher.hexdigest()[:16]
        return f"clone_{voice_hash}"