fix: encoder args #225 #224

- 修复 ffmpeg 参数覆盖问题
- 增加单测
- 增加 AudioHandler.get_sample_rate()

Author: zhzLuke96

modules/core/handler/AudioHandler.py

@@ -62,6 +62,11 @@ def enqueue_stream(self) -> AsyncGenerator[NP_AUDIO, None]:
             "Method 'enqueue_stream' must be implemented by subclass"
         )
 
+    def get_sample_rate(self) -> int:
+        raise NotImplementedError(
+            "Method 'get_sample_rate' must be implemented by subclass"
+        )
+
     def set_current_request(self, request: Request):
         assert self.current_request is None, "current_request has been set"
         assert isinstance(
@@ -93,7 +98,9 @@ def get_encoder(self) -> StreamEncoder:
         else:
             raise ValueError(f"Unsupported audio format: {format}")
 
+        encoder.set_header(sample_rate=self.get_sample_rate())
         encoder.open(bitrate=bitrate, acodec=acodec)
+        encoder.write_header_data()
 
         return encoder
 
@@ -104,9 +111,6 @@ async def enqueue_to_stream(self) -> AsyncGenerator[bytes, None]:
 
             chunk_data = bytes()
             async for sample_rate, audio_data in self.enqueue_stream():
-                encoder.set_header(
-                    sample_rate=sample_rate, sample_width=audio_data.dtype.itemsize
-                )
                 audio_bytes = covert_to_s16le(audio_data=audio_data)
 
                 logger.debug(f"write audio_bytes len: {len(audio_bytes)}")
@@ -153,9 +157,6 @@ async def enqueue_to_stream_join(self) -> AsyncGenerator[bytes, None]:
         encoder = self.get_encoder()
         chunk_data = bytes()
         async for sample_rate, audio_data in self.enqueue_stream():
-            encoder.set_header(
-                sample_rate=sample_rate, sample_width=audio_data.dtype.itemsize
-            )
             audio_bytes = covert_to_s16le(audio_data=audio_data)
             encoder.write(audio_bytes)
 
@@ -166,26 +167,30 @@ async def enqueue_to_stream_join(self) -> AsyncGenerator[bytes, None]:
 
         encoder.terminate()
 
+    async def _enqueue_to_bytes(self) -> bytes:
+        """
+        为了测试拆分的函数
+        这个函数不依赖 current_request 状态
+        """
+        encoder = self.get_encoder()
+        buffer = bytes()
+        try:
+            sample_rate, audio_data = await self.enqueue()
+            audio_bytes = covert_to_s16le(audio_data=audio_data)
+            encoder.write(audio_bytes)
+            encoder.close()
+            buffer = encoder.read_all()
+        finally:
+            encoder.terminate()
+        return buffer
+
     async def enqueue_to_bytes(self) -> bytes:
         if self.current_request is None:
             raise ValueError("current_request is not set")
 
-        encoder = self.get_encoder()
-
         # NOTE: 这里的逻辑类似 goto
         async with cancel_on_disconnect(self.current_request):
-            try:
-                sample_rate, audio_data = await self.enqueue()
-                audio_bytes = covert_to_s16le(audio_data=audio_data)
-                encoder.set_header(
-                    sample_rate=sample_rate, sample_width=audio_data.dtype.itemsize
-                )
-                encoder.write(audio_bytes)
-                encoder.close()
-                buffer = encoder.read_all()
-            finally:
-                encoder.terminate()
-            return buffer
+            return self._enqueue_to_bytes()
 
         logger.debug(f"disconnected")
         self.interrupt()

modules/core/handler/TTSHandler.py

@@ -100,6 +100,9 @@ def interrupt(self):
         self.ctx.stop = True
         self.pipeline.model.interrupt()
 
+    def get_sample_rate(self):
+        return self.pipeline.model.get_sample_rate()
+
     async def enqueue(self) -> NP_AUDIO:
         timeout = self.ctx.infer_config.timeout
         return await self.pipeline.generate(timeout=timeout)

modules/core/handler/VCHandler.py

@@ -32,7 +32,7 @@ def __init__(
         if self.model is None:
             raise Exception(f"Model {self.vc_config.mid} is not supported")
 
-    def get_model(self):
+    def get_model(self) -> VCModel:
         model_id = (
             self.vc_config.mid.lower()
             .replace(" ", "")
@@ -55,3 +55,6 @@ async def enqueue_stream(self) -> AsyncGenerator[NP_AUDIO, None]:
         raise NotImplementedError(
             "Method 'enqueue_stream' not implemented in VCHandler"
         )
+
+    def get_sample_rate(self):
+        return self.model.get_sample_rate()

modules/core/handler/encoder/StreamEncoder.py

@@ -43,21 +43,47 @@ def __init__(self) -> None:
     def set_header(
         self, *, frame_input=b"", channels=1, sample_width=2, sample_rate=24000
     ):
+        """
+        基本上只需要改 sample_rate 因为我们输入的都是 pcm s16le (int16)
+        """
+        self.channels = channels
+        self.sample_width = sample_width
+        self.sample_rate = sample_rate
+
+        logger.info(
+            f"StreamEncoder header set, channels: {channels}, sample_width: {sample_width}, sample_rate: {sample_rate}"
+        )
+
+    def write_header_data(self):
         if self.header:
             return
         header_bytes = wave_header_chunk(
-            frame_input, channels, sample_width, sample_rate
+            channels=self.channels,
+            sample_width=self.sample_width,
+            sample_rate=self.sample_rate,
         )
         self.header = header_bytes
         self.write(header_bytes)
 
         logger.info(
-            f"StreamEncoder header set, channels: {channels}, sample_width: {sample_width}, sample_rate: {sample_rate}"
+            f"StreamEncoder header written, channels: {self.channels}, sample_width: {self.sample_width}, sample_rate: {self.sample_rate}"
         )
 
     def open(
-        self, format: str = "mp3", acodec: str = "libmp3lame", bitrate: str = "320k"
+        self,
+        format: str = "mp3",
+        acodec: str = "libmp3lame",
+        bitrate: str = "320k",
+        input_dtype: str = "s16le",  # s16le or s32le
     ):
+        """
+        打开编码器
+
+        :param format: 输出格式
+        :param acodec: 输出编码器
+        :param bitrate: 输出比特率
+        :param input_dtype: 输入数据类型 s16le or s32le
+        """
         encoder = self.encoder
         self.p = subprocess.Popen(
             [
@@ -66,14 +92,12 @@ def open(
                 "-threads",
                 str(os.cpu_count() or 4),
                 # NOTE: 指定输入格式为 16 位 PCM
-                # NOTE: 其实文件头里面有写，但是没有文件名，所以需要手动指定
                 "-f",
-                "s16le",
-                # NOTE: 不要在这里传递 ar/ac ，我们写在wav文件头上，这里会覆盖掉文件头读取的数据
-                # "-ar",
-                # str(self.sample_rate),  # 输入采样率
-                # "-ac",
-                # "1",  # 输入单声道
+                input_dtype,
+                "-ar",
+                str(self.sample_rate),  # 输入采样率
+                "-ac",
+                str(self.channels),  # 输入单声道
                 "-i",
                 "pipe:0",
                 "-f",

modules/core/models/vc/OpenVoice.py

@@ -59,6 +59,9 @@ def sampling_rate(self) -> int:
         hps = self.model.hps
         return hps.data.sampling_rate
 
+    def get_sample_rate(self):
+        return self.sampling_rate
+
     def audio_to_se(self, audio: NP_AUDIO) -> torch.Tensor:
         hps = self.model.hps
         device = self.device

modules/core/models/vc/VCModel.py

@@ -10,3 +10,6 @@ def convert(
         self, src_audio: NP_AUDIO, ref_spk: TTSSpeaker, config: VCConfig
     ) -> NP_AUDIO:
         raise NotImplementedError
+
+    def get_sample_rate(self) -> int:
+        raise NotImplementedError

tests/encders/test_encoders.py

@@ -0,0 +1,49 @@
+import pytest
+
+from modules.core.handler.TTSHandler import TTSHandler
+from modules.core.handler.datacls.audio_model import EncoderConfig
+from modules.core.handler.datacls.tts_model import TTSConfig
+from modules.core.spk.SpkMgr import spk_mgr
+
+# 这里测试 encoder 是否可以正常工作
+# 正常工作的定义是输出音频正常无噪音无变速
+
+need_test_models = ["chat-tts", "fish-speech", "cosy-voice"]
+
+
+@pytest.mark.parametrize(
+    "model_id, format",
+    [
+        # raw 格式就是 wav，只是直接输出pcm
+        ("chat-tts", "raw"),
+        ("fish-speech", "raw"),
+        ("cosy-voice", "raw"),
+        ("chat-tts", "mp3"),
+        ("fish-speech", "mp3"),
+        ("cosy-voice", "mp3"),
+        ("chat-tts", "wav"),
+        ("fish-speech", "wav"),
+        ("cosy-voice", "wav"),
+    ],
+)
+@pytest.mark.encoders
+@pytest.mark.asyncio
+async def test_encoders(model_id, format):
+    spk_mona = spk_mgr.get_speaker("mona")
+    handler = TTSHandler(
+        text_content="云想衣裳花想容，春风拂槛露华浓。 若非群玉山头见，会向瑶台月下逢。",
+        spk=spk_mona,
+        tts_config=TTSConfig(mid=model_id),
+        encoder_config=EncoderConfig(format=format),
+    )
+    file_bytes = await handler._enqueue_to_bytes()
+
+    ext = format
+    if format == "raw":
+        ext = "wav"
+
+    # 1. 不为空
+    assert len(file_bytes) > 0
+    # 2. 保存到 tests/test_outputs 之下，然后人工检查
+    with open(f"tests/test_outputs/test_encder_{model_id}_{format}.{ext}", "wb") as f:
+        f.write(file_bytes)