✨ webui in tts_pipeline

- add infer_config.sync_gen 以支持 track_tqdm
- remove 老代码

refs #91 #90

modules/core/handler/AudioHandler.py

@@ -1,12 +1,10 @@
 import base64
 import io
-import struct
 import wave
 from typing import AsyncGenerator, Generator
 
 import numpy as np
 from fastapi import Request
-from pydub import AudioSegment
 
 from modules.core.handler.encoder.StreamEncoder import StreamEncoder
 from modules.core.handler.encoder.WavFile import WAVFileBytes

modules/core/handler/datacls/chattts_model.py

@@ -26,3 +26,6 @@ class InferConfig(BaseModel):
     stream_chunk_size: int = 96
 
     no_cache: bool = False
+
+    # 开启同步生成 （主要是给gradio用）
+    sync_gen: bool = False

modules/core/models/tts/ChatTtsModel.py

@@ -1,11 +1,11 @@
-from typing import Any, Generator
+from typing import Any, Generator, Union
 
 import numpy as np
 
 from modules.core.models.TTSModel import TTSModel
 from modules.core.models.zoo.ChatTTS import ChatTTS, load_chat_tts, unload_chat_tts
 from modules.core.models.zoo.ChatTTSInfer import ChatTTSInfer
-from modules.core.models.zoo.InerCache import InferCache
+from modules.core.models.tts.InerCache import InferCache
 from modules.core.pipeline.dcls import TTSPipelineContext
 from modules.core.pipeline.pipeline import TTSSegment
 from modules.core.pipeline.processor import NP_AUDIO
@@ -16,7 +16,7 @@ class ChatTTSModel(TTSModel):
     model_id = "chat-tts"
 
     def __init__(self) -> None:
-        super().__init__("chat-tts-4w")
+        super().__init__("chat-tts")
         self.chat: ChatTTS = None
 
     def load(self, context: TTSPipelineContext) -> ChatTTS:
@@ -75,11 +75,15 @@ def get_cache_kwargs(self, segments: list[TTSSegment], context: TTSPipelineConte
 
     def get_cache(
         self, segments: list[TTSSegment], context: TTSPipelineContext
-    ) -> list[NP_AUDIO]:
+    ) -> Union[list[NP_AUDIO], None]:
         no_cache = context.infer_config.no_cache
         if no_cache:
             return None
 
+        is_random_generate = context.infer_config.seed == -1
+        if is_random_generate:
+            return None
+
         kwargs = self.get_cache_kwargs(segments=segments, context=context)
 
         if InferCache.get_cache_val(model_id=self.model_id, **kwargs):
@@ -102,7 +106,7 @@ def set_cache(
 
     def generate_batch_base(
         self, segments: list[TTSSegment], context: TTSPipelineContext, stream=False
-    ) -> list[NP_AUDIO] | Generator[list[NP_AUDIO], Any, None]:
+    ) -> Union[list[NP_AUDIO], Generator[list[NP_AUDIO], Any, None]]:
         cached = self.get_cache(segments=segments, context=context)
         if cached is not None:
             if not stream:

modules/core/models/zoo/ChatTTSInfer.py

@@ -9,6 +9,7 @@
 from modules import config
 from modules.ChatTTS.ChatTTS.core import Chat
 from modules.ChatTTS.ChatTTS.model import GPT
+from modules.core.models import zoo
 from modules.utils.monkey_tqdm import disable_tqdm
 
 
@@ -48,6 +49,9 @@ def __init__(self, instance: Chat) -> None:
         self.instance = instance
         ChatTTSInfer.current_infer = self
 
+        if zoo.zoo_config.debug_generate:
+            self.logger.setLevel(logging.DEBUG)
+
     def get_tokenizer(self) -> LlamaTokenizer:
         return self.instance.pretrain_models["tokenizer"]
 
@@ -102,6 +106,14 @@ def _infer(
         # smooth_decoding = stream
         smooth_decoding = False
 
+        self.logger.debug(
+            f"Start infer: stream={stream}, skip_refine_text={skip_refine_text}, refine_text_only={refine_text_only}, use_decoder={use_decoder}, smooth_decoding={smooth_decoding}"
+        )
+        self.logger.debug(
+            f"params_refine_text={params_refine_text}, params_infer_code={params_infer_code}"
+        )
+        self.logger.debug(f"Text: {text}")
+
         with torch.no_grad():
 
             if not skip_refine_text:
@@ -131,6 +143,7 @@ def _infer(
                     wavs = self._decode_to_wavs(result, length, use_decoder)
                     yield wavs
             else:
+                # NOTE: 貌似没什么用...?
                 # smooth_decoding 即使用了滑动窗口的解码，每次都保留上一段的隐藏状态一起解码，并且保留上一段的音频长度用于截取
                 @dataclass(repr=False, eq=False)
                 class WavWindow:

modules/core/models/zoo/__init__.py

@@ -1 +1,2 @@
 from . import ChatTTS
+from . import zoo_config

modules/core/models/zoo/zoo_config.py

@@ -0,0 +1 @@
+debug_generate = False

modules/core/pipeline/dcls.py

@@ -31,7 +31,7 @@ class TTSPipelineContext:
     text: Optional[str] = None
     ssml: Optional[str] = None
 
-    spk: Speaker = None
+    spk: Optional[Speaker] = None
     tts_config: ChatTTSConfig = ChatTTSConfig()
     infer_config: InferConfig = InferConfig()
     adjust_config: AdjustConfig = AdjustConfig()

modules/core/pipeline/generate/BatchSynth.py

@@ -24,6 +24,7 @@ def __init__(
         self.generator = BatchGenerate(
             buckets=self.buckets, context=context, model=model
         )
+        self.context = context
 
         self.thread1 = None
 
@@ -40,9 +41,19 @@ def read(self):
         return self.streamer.read()
 
     def start_generate(self):
+        sync_gen = self.context.infer_config.sync_gen
+        if sync_gen:
+            self.start_generate_sync()
+        else:
+            self.start_generate_async()
+
+    def start_generate_async(self):
         if self.thread1 is not None:
             return
         gen_t1 = threading.Thread(target=self.generator.generate, args=(), daemon=True)
         gen_t1.start()
         self.thread1 = gen_t1
         return gen_t1
+
+    def start_generate_sync(self):
+        self.generator.generate()