lenML/ChatTTS-Forge

1import logging2from typing import Generator3 4import numpy as np5 6from modules.api.impl.handler.AudioHandler import AudioHandler7from modules.api.impl.model.audio_model import AdjustConfig8from modules.api.impl.model.chattts_model import ChatTTSConfig, InferConfig9from modules.api.impl.model.enhancer_model import EnhancerConfig10from modules.Enhancer.ResembleEnhance import apply_audio_enhance_full11from modules.normalization import text_normalize12from modules.speaker import Speaker13from modules.synthesize_audio import synthesize_audio14from modules.synthesize_stream import synthesize_stream15from modules.utils.audio import apply_normalize, apply_prosody_to_audio_data16 17logger = logging.getLogger(__name__)18 19 20class TTSHandler(AudioHandler):21    def __init__(22        self,23        text_content: str,24        spk: Speaker,25        tts_config: ChatTTSConfig,26        infer_config: InferConfig,27        adjust_config: AdjustConfig,28        enhancer_config: EnhancerConfig,29    ):30        assert isinstance(text_content, str), "text_content should be str"31        assert isinstance(spk, Speaker), "spk should be Speaker"32        assert isinstance(33            tts_config, ChatTTSConfig34        ), "tts_config should be ChatTTSConfig"35        assert isinstance(36            infer_config, InferConfig37        ), "infer_config should be InferConfig"38        assert isinstance(39            adjust_config, AdjustConfig40        ), "adjest_config should be AdjustConfig"41        assert isinstance(42            enhancer_config, EnhancerConfig43        ), "enhancer_config should be EnhancerConfig"44 45        self.text_content = text_content46        self.spk = spk47        self.tts_config = tts_config48        self.infer_config = infer_config49        self.adjest_config = adjust_config50        self.enhancer_config = enhancer_config51 52        self.validate()53 54    def validate(self):55        # TODO params checker56        pass57 58    def enqueue(self) -> tuple[np.ndarray, int]:59        text = text_normalize(self.text_content)60        tts_config = self.tts_config61        infer_config = self.infer_config62        adjust_config = self.adjest_config63        enhancer_config = self.enhancer_config64 65        sample_rate, audio_data = synthesize_audio(66            text,67            spk=self.spk,68            temperature=tts_config.temperature,69            top_P=tts_config.top_p,70            top_K=tts_config.top_k,71            prompt1=tts_config.prompt1,72            prompt2=tts_config.prompt2,73            prefix=tts_config.prefix,74            infer_seed=infer_config.seed,75            batch_size=infer_config.batch_size,76            spliter_threshold=infer_config.spliter_threshold,77            end_of_sentence=infer_config.eos,78        )79 80        if enhancer_config.enabled:81            nfe = enhancer_config.nfe82            solver = enhancer_config.solver83            lambd = enhancer_config.lambd84            tau = enhancer_config.tau85 86            audio_data, sample_rate = apply_audio_enhance_full(87                audio_data=audio_data,88                sr=sample_rate,89                nfe=nfe,90                solver=solver,91                lambd=lambd,92                tau=tau,93            )94 95        audio_data = apply_prosody_to_audio_data(96            audio_data=audio_data,97            rate=adjust_config.speed_rate,98            pitch=adjust_config.pitch,99            volume=adjust_config.volume_gain_db,100            sr=sample_rate,101        )102 103        if adjust_config.normalize:104            sample_rate, audio_data = apply_normalize(105                audio_data=audio_data,106                headroom=adjust_config.headroom,107                sr=sample_rate,108            )109 110        return audio_data, sample_rate111 112    def enqueue_stream(self) -> Generator[tuple[np.ndarray, int], None, None]:113        text = text_normalize(self.text_content)114        tts_config = self.tts_config115        infer_config = self.infer_config116        adjust_config = self.adjest_config117        enhancer_config = self.enhancer_config118 119        if enhancer_config.enabled:120            logger.warning(121                "enhancer_config is enabled, but it is not supported in stream mode"122            )123 124        gen = synthesize_stream(125            text,126            spk=self.spk,127            temperature=tts_config.temperature,128            top_P=tts_config.top_p,129            top_K=tts_config.top_k,130            prompt1=tts_config.prompt1,131            prompt2=tts_config.prompt2,132            prefix=tts_config.prefix,133            infer_seed=infer_config.seed,134            spliter_threshold=infer_config.spliter_threshold,135            end_of_sentence=infer_config.eos,136        )137 138        # FIXME: 很奇怪，合并出来的音频每个 chunk 之前会有一段异常，暂时没有查出来是哪里的问题，可能是解码时候切割漏了？或者多了？139        for sr, wav in gen:140 141            wav = apply_prosody_to_audio_data(142                audio_data=wav,143                rate=adjust_config.speed_rate,144                pitch=adjust_config.pitch,145                volume=adjust_config.volume_gain_db,146                sr=sr,147            )148 149            if adjust_config.normalize:150                sr, wav = apply_normalize(151                    audio_data=wav,152                    headroom=adjust_config.headroom,153                    sr=sr,154                )155 156            yield wav, sr157