update

.env.webui

@@ -17,5 +17,5 @@ TTS_MAX_LEN=1000
 SSML_MAX_LEN=3000
 MAX_BATCH_SIZE=12
 
-V_GIT_TAG="🤗hf(0.5.6-rc)"
+V_GIT_TAG="🤗hf(0.6.1-rc)"
 V_GIT_COMMIT=main

language/zh-CN.json

@@ -80,6 +80,9 @@
   "readme": "readme",
   "changelog": "changelog",
   "💼Speaker file": "💼音色文件",
+  "🎛️Spliter": "🎛️分割器配置",
+  "eos": "句尾词",
+  "Spliter Threshold": "分割器阈值",
   "TTS_STYLE_GUIDE": ["后缀为 _p 表示带prompt，效果更强但是影响质量"],
   "SSML_SPLITER_GUIDE": [
     "- 如果尾字吞字不读，可以试试结尾加上 `[lbreak]`",

modules/ChatTTS/ChatTTS/core.py

@@ -17,7 +17,7 @@ from .infer.api import refine_text, infer_code
 
 from huggingface_hub import snapshot_download
 
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.ERROR)
 
 
 class Chat:

modules/SynthesizeSegments.py

@@ -1,8 +1,10 @@
+import copy
 from box import Box
 from pydub import AudioSegment
 from typing import List, Union
 from scipy.io.wavfile import write
 import io
+from modules.SentenceSplitter import SentenceSplitter
 from modules.api.utils import calc_spk_style
 from modules.ssml_parser.SSMLParser import SSMLSegment, SSMLBreak, SSMLContext
 from modules.utils import rng
@@ -56,27 +58,27 @@ def to_number(value, t, default=0):
 
 
 class TTSAudioSegment(Box):
-    text: str
-    temperature: float
-    top_P: float
-    top_K: int
-    spk: int
-    infer_seed: int
-    prompt1: str
-    prompt2: str
-    prefix: str
-
-    _type: str
-
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self._type = kwargs.get("_type", "voice")
+        self.text = kwargs.get("text", "")
+        self.temperature = kwargs.get("temperature", 0.3)
+        self.top_P = kwargs.get("top_P", 0.5)
+        self.top_K = kwargs.get("top_K", 20)
+        self.spk = kwargs.get("spk", -1)
+        self.infer_seed = kwargs.get("infer_seed", -1)
+        self.prompt1 = kwargs.get("prompt1", "")
+        self.prompt2 = kwargs.get("prompt2", "")
+        self.prefix = kwargs.get("prefix", "")
 
 
 class SynthesizeSegments:
-    def __init__(self, batch_size: int = 8):
+    def __init__(self, batch_size: int = 8, eos="", spliter_thr=100):
         self.batch_size = batch_size
         self.batch_default_spk_seed = rng.np_rng()
         self.batch_default_infer_seed = rng.np_rng()
+        self.eos = eos
+        self.spliter_thr = spliter_thr
 
     def segment_to_generate_params(
         self, segment: Union[SSMLSegment, SSMLBreak]
@@ -85,9 +87,11 @@ class SynthesizeSegments:
             return TTSAudioSegment(_type="break")
 
         if segment.get("params", None) is not None:
-            return TTSAudioSegment(**segment.get("params"))
+            params = segment.get("params")
+            text = segment.get("text", None) or segment.text or ""
+            return TTSAudioSegment(**params, text=text)
 
-        text = segment.get("text", "")
+        text = segment.get("text", None) or segment.text or ""
         is_end = segment.get("is_end", False)
 
         text = str(text).strip()
@@ -156,7 +160,7 @@ class SynthesizeSegments:
         for i in range(0, len(bucket), self.batch_size):
             batch = bucket[i : i + self.batch_size]
             param_arr = [self.segment_to_generate_params(segment) for segment in batch]
-            texts = [params.text for params in param_arr]
+            texts = [params.text + self.eos for params in param_arr]
 
             params = param_arr[0]
             audio_datas = generate_audio.generate_audio_batch(
@@ -204,9 +208,38 @@ class SynthesizeSegments:
 
         return buckets
 
+    def split_segments(self, segments: List[Union[SSMLSegment, SSMLBreak]]):
+        """
+        将 segments 中的 text 经过 spliter 处理成多个 segments
+        """
+        spliter = SentenceSplitter(threshold=self.spliter_thr)
+        ret_segments: List[Union[SSMLSegment, SSMLBreak]] = []
+
+        for segment in segments:
+            if isinstance(segment, SSMLBreak):
+                ret_segments.append(segment)
+                continue
+
+            text = segment.text
+            if not text:
+                continue
+
+            sentences = spliter.parse(text)
+            for sentence in sentences:
+                ret_segments.append(
+                    SSMLSegment(
+                        text=sentence,
+                        attrs=segment.attrs.copy(),
+                        params=copy.copy(segment.params),
+                    )
+                )
+
+        return ret_segments
+
     def synthesize_segments(
         self, segments: List[Union[SSMLSegment, SSMLBreak]]
     ) -> List[AudioSegment]:
+        segments = self.split_segments(segments)
         audio_segments = [None] * len(segments)
         buckets = self.bucket_segments(segments)
 

modules/api/api_setup.py

@@ -18,6 +18,7 @@ from modules.api.impl import (
     speaker_api,
     ping_api,
     models_api,
+    xtts_v2_api,
 )
 
 logger = logging.getLogger(__name__)
@@ -35,6 +36,7 @@ def create_api(app, exclude=[]):
     google_api.setup(app_mgr)
     openai_api.setup(app_mgr)
     refiner_api.setup(app_mgr)
+    xtts_v2_api.setup(app_mgr)
 
     return app_mgr
 
@@ -42,9 +44,9 @@ def create_api(app, exclude=[]):
 def setup_model_args(parser: argparse.ArgumentParser):
     parser.add_argument("--compile", action="store_true", help="Enable model compile")
     parser.add_argument(
-        "--half",
+        "--no_half",
         action="store_true",
-        help="Enable half precision for model inference",
+        help="Disalbe half precision for model inference",
     )
     parser.add_argument(
         "--off_tqdm",
@@ -82,7 +84,7 @@ def process_model_args(args):
     compile = env.get_and_update_env(args, "compile", False, bool)
     device_id = env.get_and_update_env(args, "device_id", None, str)
     use_cpu = env.get_and_update_env(args, "use_cpu", [], list)
-    half = env.get_and_update_env(args, "half", False, bool)
+    no_half = env.get_and_update_env(args, "no_half", False, bool)
     off_tqdm = env.get_and_update_env(args, "off_tqdm", False, bool)
     debug_generate = env.get_and_update_env(args, "debug_generate", False, bool)
 

modules/api/impl/google_api.py

@@ -13,6 +13,7 @@ from modules.Enhancer.ResembleEnhance import (
 )
 from modules.api.Api import APIManager
 from modules.synthesize_audio import synthesize_audio
+from modules.utils import audio
 from modules.utils.audio import apply_prosody_to_audio_data
 from modules.normalization import text_normalize
 
@@ -44,6 +45,9 @@ class VoiceSelectionParams(BaseModel):
     topK: int = 20
     seed: int = 42
 
+    # end_of_sentence
+    eos: str = "[uv_break]"
+
 
 class AudioConfig(BaseModel):
     audioEncoding: api_utils.AudioFormat = "mp3"
@@ -87,6 +91,7 @@ async def google_text_synthesize(request: GoogleTextSynthesizeRequest):
     language_code = voice.languageCode
     voice_name = voice.name
     infer_seed = voice.seed or 42
+    eos = voice.eos or "[uv_break]"
     audio_format = audioConfig.audioEncoding or "mp3"
     speaking_rate = audioConfig.speakingRate or 1
     pitch = audioConfig.pitch or 0
@@ -94,11 +99,9 @@ async def google_text_synthesize(request: GoogleTextSynthesizeRequest):
 
     batch_size = audioConfig.batchSize or 1
 
-    # TODO spliter_threshold
     spliter_threshold = audioConfig.spliterThreshold or 100
 
-    # TODO sample_rate
-    sample_rate_hertz = audioConfig.sampleRateHertz or 24000
+    sample_rate = audioConfig.sampleRateHertz or 24000
 
     params = api_utils.calc_spk_style(spk=voice.name, style=voice.style)
 
@@ -137,10 +140,10 @@ async def google_text_synthesize(request: GoogleTextSynthesizeRequest):
                 prefix=params.get("prefix", ""),
                 batch_size=batch_size,
                 spliter_threshold=spliter_threshold,
+                end_of_sentence=eos,
             )
 
         elif input.ssml:
-            # 处理SSML合成逻辑
             parser = create_ssml_parser()
             segments = parser.parse(input.ssml)
             for seg in segments:
@@ -151,17 +154,13 @@ async def google_text_synthesize(request: GoogleTextSynthesizeRequest):
                     status_code=422, detail="The SSML text is empty or parsing failed."
                 )
 
-            synthesize = SynthesizeSegments(batch_size=batch_size)
+            synthesize = SynthesizeSegments(
+                batch_size=batch_size, eos=eos, spliter_thr=spliter_threshold
+            )
             audio_segments = synthesize.synthesize_segments(segments)
             combined_audio = combine_audio_segments(audio_segments)
 
-            buffer = io.BytesIO()
-            combined_audio.export(buffer, format="wav")
-
-            buffer.seek(0)
-
-            audio_data = buffer.read()
-
+            sample_rate, audio_data = audio.pydub_to_np(combined_audio)
         else:
             raise HTTPException(
                 status_code=422, detail="Either text or SSML input must be provided."

modules/api/impl/openai_api.py

@@ -41,6 +41,8 @@ class AudioSpeechRequest(BaseModel):
     spliter_threshold: float = Field(
         100, ge=10, le=1024, description="Threshold for sentence spliter"
     )
+    # end of sentence
+    eos: str = "[uv_break]"
 
 
 async def openai_speech_api(
@@ -52,6 +54,7 @@ async def openai_speech_api(
     input_text = request.input
     voice = request.voice
     style = request.style
+    eos = request.eos
     response_format = request.response_format
     batch_size = request.batch_size
     spliter_threshold = request.spliter_threshold
@@ -95,6 +98,7 @@ async def openai_speech_api(
             prompt1=prompt1,
             prompt2=prompt2,
             prefix=prefix,
+            end_of_sentence=eos,
         )
 
         if speed != 1:

modules/api/impl/ssml_api.py

@@ -26,8 +26,13 @@ class SSMLRequest(BaseModel):
     # NOTE: 🤔 也许这个值应该配置成系统变量？ 传进来有点奇怪
     batch_size: int = 4
 
+    # end of sentence
+    eos: str = "[uv_break]"
 
-async def synthesize_ssml(
+    spliter_thr: int = 100
+
+
+async def synthesize_ssml_api(
     request: SSMLRequest = Body(
         ..., description="JSON body with SSML string and format"
     )
@@ -36,12 +41,19 @@ async def synthesize_ssml(
         ssml = request.ssml
         format = request.format.lower()
         batch_size = request.batch_size
+        eos = request.eos
+        spliter_thr = request.spliter_thr
 
         if batch_size < 1:
             raise HTTPException(
                 status_code=400, detail="Batch size must be greater than 0."
             )
 
+        if spliter_thr < 50:
+            raise HTTPException(
+                status_code=400, detail="Spliter threshold must be greater than 50."
+            )
+
         if not ssml or ssml == "":
             raise HTTPException(status_code=400, detail="SSML content is required.")
 
@@ -55,7 +67,9 @@ async def synthesize_ssml(
         for seg in segments:
             seg["text"] = text_normalize(seg["text"], is_end=True)
 
-        synthesize = SynthesizeSegments(batch_size)
+        synthesize = SynthesizeSegments(
+            batch_size=batch_size, eos=eos, spliter_thr=spliter_thr
+        )
         audio_segments = synthesize.synthesize_segments(segments)
         combined_audio = combine_audio_segments(audio_segments)
         buffer = io.BytesIO()
@@ -77,4 +91,4 @@ async def synthesize_ssml(
 
 
 def setup(api_manager: APIManager):
-    api_manager.post("/v1/ssml", response_class=FileResponse)(synthesize_ssml)
+    api_manager.post("/v1/ssml", response_class=FileResponse)(synthesize_ssml_api)

modules/api/impl/tts_api.py

@@ -38,6 +38,7 @@ class TTSParams(BaseModel):
     prefix: str = Query("", description="Text prefix for inference")
     bs: str = Query("8", description="Batch size for inference")
     thr: str = Query("100", description="Threshold for sentence spliter")
+    eos: str = Query("", description="End of sentence str")
 
 
 async def synthesize_tts(params: TTSParams = Depends()):
@@ -87,6 +88,7 @@ async def synthesize_tts(params: TTSParams = Depends()):
         prefix = params.prefix or calc_params.get("prefix", params.prefix)
         prompt1 = params.prompt1 or calc_params.get("prompt1", params.prompt1)
         prompt2 = params.prompt2 or calc_params.get("prompt2", params.prompt2)
+        eos = params.eos or ""
 
         batch_size = int(params.bs)
         threshold = int(params.thr)
@@ -103,6 +105,7 @@ async def synthesize_tts(params: TTSParams = Depends()):
             prefix=prefix,
             batch_size=batch_size,
             spliter_threshold=threshold,
+            end_of_sentence=eos,
         )
 
         buffer = io.BytesIO()

modules/api/impl/xtts_v2_api.py

@@ -0,0 +1,160 @@
+import io
+from fastapi import HTTPException
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+from modules.api import utils as api_utils
+from modules.api.Api import APIManager
+
+import soundfile as sf
+
+from modules import config
+from modules.normalization import text_normalize
+from modules.speaker import speaker_mgr
+from modules.synthesize_audio import synthesize_audio
+
+import logging
+
+from modules.utils.audio import apply_prosody_to_audio_data
+
+logger = logging.getLogger(__name__)
+
+
+class XTTS_V2_Settings:
+    def __init__(self):
+        self.stream_chunk_size = 100
+        self.temperature = 0.3
+        self.speed = 1
+        self.length_penalty = 0.5
+        self.repetition_penalty = 1.0
+        self.top_p = 0.7
+        self.top_k = 20
+        self.enable_text_splitting = True
+
+
+class TTSSettingsRequest(BaseModel):
+    stream_chunk_size: int
+    temperature: float
+    speed: float
+    length_penalty: float
+    repetition_penalty: float
+    top_p: float
+    top_k: int
+    enable_text_splitting: bool
+
+
+class SynthesisRequest(BaseModel):
+    text: str
+    speaker_wav: str
+    language: str
+
+
+def setup(app: APIManager):
+    XTTSV2 = XTTS_V2_Settings()
+
+    @app.get("/v1/xtts_v2/speakers")
+    async def speakers():
+        spks = speaker_mgr.list_speakers()
+        return [
+            {
+                "name": spk.name,
+                "voice_id": spk.id,
+                # TODO: 也许可以放一个 "/v1/tts" 接口地址在这里
+                "preview_url": "",
+            }
+            for spk in spks
+        ]
+
+    @app.post("/v1/xtts_v2/tts_to_audio", response_class=StreamingResponse)
+    async def tts_to_audio(request: SynthesisRequest):
+        text = request.text
+        # speaker_wav 就是 speaker id 。。。
+        voice_id = request.speaker_wav
+        language = request.language
+
+        spk = speaker_mgr.get_speaker_by_id(voice_id) or speaker_mgr.get_speaker(
+            voice_id
+        )
+        if spk is None:
+            raise HTTPException(status_code=400, detail="Invalid speaker id")
+
+        text = text_normalize(text, is_end=True)
+        sample_rate, audio_data = synthesize_audio(
+            text=text,
+            temperature=XTTSV2.temperature,
+            # length_penalty=XTTSV2.length_penalty,
+            # repetition_penalty=XTTSV2.repetition_penalty,
+            top_P=XTTSV2.top_p,
+            top_K=XTTSV2.top_k,
+            spk=spk,
+            spliter_threshold=XTTSV2.stream_chunk_size,
+            # TODO 支持设置 batch_size
+            batch_size=4,
+            end_of_sentence="[uv_break]",
+        )
+
+        if XTTSV2.speed:
+            audio_data = apply_prosody_to_audio_data(
+                audio_data,
+                rate=XTTSV2.speed,
+                sr=sample_rate,
+            )
+
+        # to mp3
+        buffer = io.BytesIO()
+        sf.write(buffer, audio_data, sample_rate, format="wav")
+        buffer.seek(0)
+
+        buffer = api_utils.wav_to_mp3(buffer)
+
+        return StreamingResponse(buffer, media_type="audio/mpeg")
+
+    @app.get("/v1/xtts_v2/tts_stream")
+    async def tts_stream():
+        raise HTTPException(status_code=501, detail="Not implemented")
+
+    @app.post("/v1/xtts_v2/set_tts_settings")
+    async def set_tts_settings(request: TTSSettingsRequest):
+        try:
+            if request.stream_chunk_size < 50:
+                raise HTTPException(
+                    status_code=400, detail="stream_chunk_size must be greater than 0"
+                )
+            if request.temperature < 0:
+                raise HTTPException(
+                    status_code=400, detail="temperature must be greater than 0"
+                )
+            if request.speed < 0:
+                raise HTTPException(
+                    status_code=400, detail="speed must be greater than 0"
+                )
+            if request.length_penalty < 0:
+                raise HTTPException(
+                    status_code=400, detail="length_penalty must be greater than 0"
+                )
+            if request.repetition_penalty < 0:
+                raise HTTPException(
+                    status_code=400, detail="repetition_penalty must be greater than 0"
+                )
+            if request.top_p < 0:
+                raise HTTPException(
+                    status_code=400, detail="top_p must be greater than 0"
+                )
+            if request.top_k < 0:
+                raise HTTPException(
+                    status_code=400, detail="top_k must be greater than 0"
+                )
+
+            XTTSV2.stream_chunk_size = request.stream_chunk_size
+            XTTSV2.temperature = request.temperature
+            XTTSV2.speed = request.speed
+            XTTSV2.length_penalty = request.length_penalty
+            XTTSV2.repetition_penalty = request.repetition_penalty
+            XTTSV2.top_p = request.top_p
+            XTTSV2.top_k = request.top_k
+            XTTSV2.enable_text_splitting = request.enable_text_splitting
+            return {"message": "Settings successfully applied"}
+        except Exception as e:
+            if isinstance(e, HTTPException):
+                raise e
+            logger.error(e)
+            raise HTTPException(status_code=500, detail=str(e))

modules/devices/devices.py

@@ -127,7 +127,7 @@ def reset_device():
     global dtype_gpt
     global dtype_decoder
 
-    if config.runtime_env_vars.half:
+    if not config.runtime_env_vars.no_half:
         dtype = torch.float16
         dtype_dvae = torch.float16
         dtype_vocos = torch.float16

modules/finetune/model/encoder.py

@@ -0,0 +1,87 @@
+import torch
+import torch.nn as nn
+
+from modules.ChatTTS.ChatTTS.model.dvae import ConvNeXtBlock, DVAEDecoder
+
+from .wavenet import WaveNet
+
+
+def get_encoder_config(decoder: DVAEDecoder) -> dict[str, int | bool]:
+    return {
+        "idim": decoder.conv_out.out_channels,
+        "odim": decoder.conv_in[0].in_channels,
+        "n_layer": len(decoder.decoder_block),
+        "bn_dim": decoder.conv_in[0].out_channels,
+        "hidden": decoder.conv_in[2].out_channels,
+        "kernel": decoder.decoder_block[0].dwconv.kernel_size[0],
+        "dilation": decoder.decoder_block[0].dwconv.dilation[0],
+        "down": decoder.up,
+    }
+
+
+class DVAEEncoder(nn.Module):
+    def __init__(
+        self,
+        idim: int,
+        odim: int,
+        n_layer: int = 12,
+        bn_dim: int = 64,
+        hidden: int = 256,
+        kernel: int = 7,
+        dilation: int = 2,
+        down: bool = False,
+    ) -> None:
+        super().__init__()
+        self.wavenet = WaveNet(
+            input_channels=100,
+            residual_channels=idim,
+            residual_layers=20,
+            dilation_cycle=4,
+        )
+        self.conv_in_transpose = nn.ConvTranspose1d(
+            idim, hidden, kernel_size=1, bias=False
+        )
+        # nn.Sequential(
+        #     nn.ConvTranspose1d(100, idim, 3, 1, 1, bias=False),
+        #     nn.ConvTranspose1d(idim, hidden, kernel_size=1, bias=False)
+        # )
+        self.encoder_block = nn.ModuleList(
+            [
+                ConvNeXtBlock(
+                    hidden,
+                    hidden * 4,
+                    kernel,
+                    dilation,
+                )
+                for _ in range(n_layer)
+            ]
+        )
+        self.conv_out_transpose = nn.Sequential(
+            nn.Conv1d(hidden, bn_dim, 3, 1, 1),
+            nn.GELU(),
+            nn.Conv1d(bn_dim, odim, 3, 1, 1),
+        )
+
+    def forward(
+        self,
+        audio_mel_specs: torch.Tensor,  # (batch_size, audio_len*2, 100)
+        audio_attention_mask: torch.Tensor,  # (batch_size, audio_len)
+        conditioning=None,
+    ) -> torch.Tensor:
+        mel_attention_mask = (
+            audio_attention_mask.unsqueeze(-1).repeat(1, 1, 2).flatten(1)
+        )
+        x: torch.Tensor = self.wavenet(
+            audio_mel_specs.transpose(1, 2)
+        )  # (batch_size, idim, audio_len*2)
+        x = x * mel_attention_mask.unsqueeze(1)
+        x = self.conv_in_transpose(x)  # (batch_size, hidden, audio_len*2)
+        for f in self.encoder_block:
+            x = f(x, conditioning)
+        x = self.conv_out_transpose(x)  # (batch_size, odim, audio_len*2)
+        x = (
+            x.view(x.size(0), x.size(1), 2, x.size(2) // 2)
+            .permute(0, 3, 1, 2)
+            .flatten(2)
+        )
+        return x  # (batch_size, audio_len, audio_dim=odim*2)

modules/finetune/model/wavenet.py

@@ -0,0 +1,227 @@
+"""https://github.com/fishaudio/fish-speech/blob/main/fish_speech/models/vqgan/modules/wavenet.py"""
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Mish(nn.Module):
+    def forward(self, x):
+        return x * torch.tanh(F.softplus(x))
+
+
+class DiffusionEmbedding(nn.Module):
+    """Diffusion Step Embedding"""
+
+    def __init__(self, d_denoiser):
+        super(DiffusionEmbedding, self).__init__()
+        self.dim = d_denoiser
+
+    def forward(self, x):
+        device = x.device
+        half_dim = self.dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
+        emb = x[:, None] * emb[None, :]
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb
+
+
+class LinearNorm(nn.Module):
+    """LinearNorm Projection"""
+
+    def __init__(self, in_features, out_features, bias=False):
+        super(LinearNorm, self).__init__()
+        self.linear = nn.Linear(in_features, out_features, bias)
+
+        nn.init.xavier_uniform_(self.linear.weight)
+        if bias:
+            nn.init.constant_(self.linear.bias, 0.0)
+
+    def forward(self, x):
+        x = self.linear(x)
+        return x
+
+
+class ConvNorm(nn.Module):
+    """1D Convolution"""
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=1,
+        stride=1,
+        padding=None,
+        dilation=1,
+        bias=True,
+        w_init_gain="linear",
+    ):
+        super(ConvNorm, self).__init__()
+
+        if padding is None:
+            assert kernel_size % 2 == 1
+            padding = int(dilation * (kernel_size - 1) / 2)
+
+        self.conv = nn.Conv1d(
+            in_channels,
+            out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            bias=bias,
+        )
+        nn.init.kaiming_normal_(self.conv.weight)
+
+    def forward(self, signal):
+        conv_signal = self.conv(signal)
+
+        return conv_signal
+
+
+class ResidualBlock(nn.Module):
+    """Residual Block"""
+
+    def __init__(
+        self,
+        residual_channels,
+        use_linear_bias=False,
+        dilation=1,
+        condition_channels=None,
+    ):
+        super(ResidualBlock, self).__init__()
+        self.conv_layer = ConvNorm(
+            residual_channels,
+            2 * residual_channels,
+            kernel_size=3,
+            stride=1,
+            padding=dilation,
+            dilation=dilation,
+        )
+
+        if condition_channels is not None:
+            self.diffusion_projection = LinearNorm(
+                residual_channels, residual_channels, use_linear_bias
+            )
+            self.condition_projection = ConvNorm(
+                condition_channels, 2 * residual_channels, kernel_size=1
+            )
+
+        self.output_projection = ConvNorm(
+            residual_channels, 2 * residual_channels, kernel_size=1
+        )
+
+    def forward(self, x, condition=None, diffusion_step=None):
+        y = x
+
+        if diffusion_step is not None:
+            diffusion_step = self.diffusion_projection(diffusion_step).unsqueeze(-1)
+            y = y + diffusion_step
+
+        y = self.conv_layer(y)
+
+        if condition is not None:
+            condition = self.condition_projection(condition)
+            y = y + condition
+
+        gate, filter = torch.chunk(y, 2, dim=1)
+        y = torch.sigmoid(gate) * torch.tanh(filter)
+
+        y = self.output_projection(y)
+        residual, skip = torch.chunk(y, 2, dim=1)
+
+        return (x + residual) / math.sqrt(2.0), skip
+
+
+class WaveNet(nn.Module):
+    def __init__(
+        self,
+        input_channels: Optional[int] = None,
+        output_channels: Optional[int] = None,
+        residual_channels: int = 512,
+        residual_layers: int = 20,
+        dilation_cycle: Optional[int] = 4,
+        is_diffusion: bool = False,
+        condition_channels: Optional[int] = None,
+    ):
+        super().__init__()
+
+        # Input projection
+        self.input_projection = None
+        if input_channels is not None and input_channels != residual_channels:
+            self.input_projection = ConvNorm(
+                input_channels, residual_channels, kernel_size=1
+            )
+
+        if input_channels is None:
+            input_channels = residual_channels
+
+        self.input_channels = input_channels
+
+        # Residual layers
+        self.residual_layers = nn.ModuleList(
+            [
+                ResidualBlock(
+                    residual_channels=residual_channels,
+                    use_linear_bias=False,
+                    dilation=2 ** (i % dilation_cycle) if dilation_cycle else 1,
+                    condition_channels=condition_channels,
+                )
+                for i in range(residual_layers)
+            ]
+        )
+
+        # Skip projection
+        self.skip_projection = ConvNorm(
+            residual_channels, residual_channels, kernel_size=1
+        )
+
+        # Output projection
+        self.output_projection = None
+        if output_channels is not None and output_channels != residual_channels:
+            self.output_projection = ConvNorm(
+                residual_channels, output_channels, kernel_size=1
+            )
+
+        if is_diffusion:
+            self.diffusion_embedding = DiffusionEmbedding(residual_channels)
+            self.mlp = nn.Sequential(
+                LinearNorm(residual_channels, residual_channels * 4, False),
+                Mish(),
+                LinearNorm(residual_channels * 4, residual_channels, False),
+            )
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv1d, nn.Linear)):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if getattr(m, "bias", None) is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x, t=None, condition=None):
+        if self.input_projection is not None:
+            x = self.input_projection(x)
+            x = F.silu(x)
+
+        if t is not None:
+            t = self.diffusion_embedding(t)
+            t = self.mlp(t)
+
+        skip = []
+        for layer in self.residual_layers:
+            x, skip_connection = layer(x, condition, t)
+            skip.append(skip_connection)
+
+        x = torch.sum(torch.stack(skip), dim=0) / math.sqrt(len(self.residual_layers))
+        x = self.skip_projection(x)
+
+        if self.output_projection is not None:
+            x = F.silu(x)
+            x = self.output_projection(x)
+
+        return x

modules/finetune/train_gpt.py

@@ -0,0 +1,246 @@
+import functools
+import torch
+import transformers
+import peft
+from transformers.trainer_pt_utils import LabelSmoother
+from utils.dataset import AudioCollator
+from utils.logger import MetricLogger
+from utils.output import ansi, get_ansi_len, output_iter
+
+IGNORE_TOKEN_ID = LabelSmoother.ignore_index
+
+
+def train_gpt_lora(
+    chat,
+    dataset,
+    decoder_encoder,
+    dvae_encoder,
+    batch_size=16,
+    epochs=10,
+    train_text=True,
+    speaker_embeds=None,
+    lora_r=8,
+    lora_alpha=16,
+):
+    if speaker_embeds is None:
+        speaker_embeds = {}
+
+    tokenizer = chat.pretrain_models["tokenizer"]
+    decoder_decoder = chat.pretrain_models["decoder"]
+    decoder_decoder.eval().requires_grad_(False)
+    decoder_encoder.to(device=dataset.device).eval().requires_grad_(False)
+    dvae_decoder = chat.pretrain_models["dvae"]
+    dvae_decoder.eval().requires_grad_(False)
+    dvae_encoder.to(device=dataset.device).eval().requires_grad_(False)
+
+    gpt = chat.pretrain_models["gpt"]
+    gpt.train().requires_grad_()
+
+    # Add LoRA to GPT model
+    lora_config = peft.LoraConfig(r=lora_r, lora_alpha=lora_alpha)
+    gpt.gpt = peft.get_peft_model(gpt.gpt, lora_config)
+
+    speaker_embeds = {
+        speaker: torch.randn(768, device=dataset.device, requires_grad=True)
+        for speaker in dataset.speakers
+    } | speaker_embeds
+
+    for speaker_embed in speaker_embeds.values():
+        std, mean = chat.pretrain_models["spk_stat"].chunk(2)
+        speaker_embed.data = speaker_embed.data * std + mean
+
+    SPEAKER_TOKEN_ID = tokenizer.convert_tokens_to_ids("[spk_emb]")
+    AUDIO_EOS_TOKEN_ID = 0
+    AUDIO_PAD_TOKEN_ID = AUDIO_EOS_TOKEN_ID
+
+    train_params = list(gpt.parameters()) + list(speaker_embeds.values())
+    optimizer = torch.optim.Adam(
+        gpt.parameters(), lr=1e-3, weight_decay=0, betas=[0.9, 0.95], eps=1e-5
+    )
+    optimizer.add_param_group({"params": speaker_embeds.values(), "lr": 1e-1})
+
+    loss_fn = torch.nn.CrossEntropyLoss()
+    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs, 1e-7)
+
+    loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        collate_fn=AudioCollator(text_pad=tokenizer.pad_token_id),
+    )
+    logger = MetricLogger()
+    logger.create_meters(loss=None, mse_loss=None, audio_loss=None, text_loss=None)
+
+    for _epoch in range(epochs):
+        _epoch += 1
+        logger.reset()
+        header = "{blue_light}{0}: {1}{reset}".format(
+            "Epoch", output_iter(_epoch, epochs), **ansi
+        )
+        header = header.ljust(max(len("Epoch"), 30) + get_ansi_len(header))
+        iterator = logger.log_every(loader, header=header, tqdm_header="Batch")
+
+        for batch in iterator:
+            speakers = batch["speaker"]
+            text_input_ids = batch["text_input_ids"]
+            text_attention_mask = batch["text_attention_mask"]
+            audio_mel_specs = batch["audio_mel_specs"]
+            audio_attention_mask = batch["audio_attention_mask"]
+
+            batch_size, text_len = text_attention_mask.size()
+
+            dvae_audio_latents = dvae_encoder(audio_mel_specs, audio_attention_mask)
+            _, dvae_audio_input_ids = quantize(
+                dvae_decoder.vq_layer.quantizer, dvae_audio_latents
+            )
+            dvae_audio_input_ids[~audio_attention_mask.bool()] = AUDIO_PAD_TOKEN_ID
+
+            extended_audio_attention_mask = torch.cat(
+                [
+                    audio_attention_mask,
+                    torch.zeros(
+                        (batch_size, 1),
+                        dtype=audio_attention_mask.dtype,
+                        device=audio_attention_mask.device,
+                    ),
+                ],
+                dim=1,
+            )
+            extended_audio_input_ids = torch.cat(
+                [
+                    dvae_audio_input_ids,
+                    AUDIO_PAD_TOKEN_ID
+                    * torch.ones(
+                        (batch_size, 1, gpt.num_vq),
+                        dtype=dvae_audio_input_ids.dtype,
+                        device=dvae_audio_input_ids.device,
+                    ),
+                ],
+                dim=1,
+            )
+
+            indices = audio_attention_mask.int().sum(dim=1)
+            for i in range(batch_size):
+                extended_audio_attention_mask[i, indices[i]] = 1
+                extended_audio_input_ids[i, indices[i]] = AUDIO_EOS_TOKEN_ID
+
+            input_ids = torch.cat(
+                [
+                    text_input_ids.unsqueeze(-1).repeat(1, 1, gpt.num_vq),
+                    extended_audio_input_ids,
+                ],
+                dim=1,
+            )
+            attention_mask = torch.cat(
+                [text_attention_mask, extended_audio_attention_mask], dim=1
+            )
+            text_mask = torch.cat(
+                [
+                    torch.ones_like(text_attention_mask, dtype=bool),
+                    torch.zeros_like(extended_audio_attention_mask, dtype=bool),
+                ],
+                dim=1,
+            )
+            labels = input_ids.clone()
+            labels[~attention_mask.bool()] = IGNORE_TOKEN_ID
+
+            inputs_embeds = gpt.get_emb(input_ids=input_ids, text_mask=text_mask)
+
+            indices = torch.all(input_ids == SPEAKER_TOKEN_ID, dim=-1)
+            for i, speaker in enumerate(speakers):
+                inputs_embeds[i, indices[i]] = torch.nn.functional.normalize(
+                    speaker_embeds[speaker].to(dtype=inputs_embeds.dtype),
+                    p=2.0,
+                    dim=-1,
+                    eps=1e-12,
+                ).unsqueeze(0)
+
+            outputs = gpt.gpt.forward(
+                inputs_embeds=inputs_embeds, attention_mask=attention_mask
+            )
+            hidden_states = outputs.last_hidden_state
+            text_hidden_states = hidden_states[:, : text_len - 1]
+            audio_hidden_states = hidden_states[:, text_len - 1 : -1]
+
+            audio_logits = torch.stack(
+                [gpt.head_code[i](audio_hidden_states) for i in range(gpt.num_vq)],
+                dim=2,
+            )
+            audio_loss = loss_fn(
+                audio_logits.flatten(0, 2), labels[:, text_len:].flatten(0, 2)
+            )
+            loss = audio_loss
+
+            if train_text:
+                text_logits = gpt.head_text(text_hidden_states)
+                text_loss = loss_fn(
+                    text_logits.flatten(0, 1), labels[:, 1:text_len, 0].flatten(0, 1)
+                )
+                loss += text_loss
+                logger.meters["text_loss"].update(text_loss.item(), n=batch_size)
+
+            gpt_gen_mel_specs = decoder_decoder(
+                audio_hidden_states[:, :-1].transpose(1, 2)
+            ).transpose(1, 2)
+            mse_loss = torch.nn.functional.mse_loss(gpt_gen_mel_specs, audio_mel_specs)
+            loss += 0.01 * mse_loss
+
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(train_params, 1.0)
+            optimizer.step()
+
+            logger.meters["loss"].update(loss.item(), n=batch_size)
+            logger.meters["mse_loss"].update(mse_loss.item(), n=batch_size)
+            logger.meters["audio_loss"].update(audio_loss.item(), n=batch_size)
+
+        lr_scheduler.step()
+    optimizer.zero_grad()
+    return speaker_embeds
+
+
+# Example usage
+def main():
+    # Load necessary models and data paths
+    chat = ChatTTS.Chat()
+    chat.load_models()
+    dataset = XzListTar(
+        root="data/all.list",
+        tokenizer=chat.pretrain_models["tokenizer"],
+        vocos_model=chat.pretrain_models["vocos"],
+        tar_path="data/Xz.tar",
+        tar_in_memory=True,
+        process_ahead=True,
+    )
+
+    decoder_encoder = DVAEEncoder(
+        **get_encoder_config(chat.pretrain_models["decoder"].decoder)
+    )
+    dvae_encoder = DVAEEncoder(
+        **get_encoder_config(chat.pretrain_models["dvae"].decoder)
+    )
+
+    # Train GPT with LoRA
+    speaker_embeds = train_gpt_lora(
+        chat=chat,
+        dataset=dataset,
+        decoder_encoder=decoder_encoder,
+        dvae_encoder=dvae_encoder,
+        batch_size=32,
+        epochs=10,
+        train_text=True,
+        lora_r=8,
+        lora_alpha=16,
+    )
+
+    # Save LoRA parameters and embeddings
+    lora_save_path = "./saved_models/gpt_lora.pth"
+    peft.save_pretrained(gpt.gpt, lora_save_path)
+    np.savez(
+        "./saved_models/speaker_embeds.npz",
+        **{k: v.cpu().numpy() for k, v in speaker_embeds.items()}
+    )
+
+
+if __name__ == "__main__":
+    main()

modules/finetune/train_speaker.py

@@ -0,0 +1,296 @@
+import torch
+import torch.nn.functional as F
+import transformers
+
+from modules.finetune.model.encoder import DVAEEncoder, get_encoder_config
+from modules.finetune.utils.output import get_ansi_len, output_iter, ansi
+from .utils.logger import MetricLogger
+from .utils.dataset import AudioCollator, XzListTar
+from .utils.model import quantize
+
+IGNORE_TOKEN_ID = transformers.trainer_pt_utils.LabelSmoother.ignore_index
+
+
+def train_speaker_embeddings(
+    chat,
+    dataset,
+    gpt,
+    batch_size=16,
+    epochs=10,
+    train_text=True,
+    speaker_embeds=None,
+):
+    tokenizer = chat.pretrain_models["tokenizer"]
+
+    decoder_decoder = chat.pretrain_models["decoder"]
+    decoder_decoder.eval().requires_grad_(False)
+    decoder_encoder = DVAEEncoder(**get_encoder_config(decoder_decoder.decoder)).to(
+        device=dataset.device
+    )
+    decoder_encoder.eval().requires_grad_(False)
+
+    dvae_decoder = chat.pretrain_models["dvae"]
+    dvae_decoder.eval().requires_grad_(False)
+    dvae_encoder = DVAEEncoder(**get_encoder_config(dvae_decoder.decoder)).to(
+        device=dataset.device
+    )
+    dvae_encoder.eval().requires_grad_(False)
+
+    if speaker_embeds is None:
+        speaker_embeds = {
+            speaker: torch.randn(
+                768,
+                device=dataset.device,
+                requires_grad=True,
+            )
+            for speaker in dataset.speakers
+        }
+        for speaker_embed in speaker_embeds.values():
+            std, mean = chat.pretrain_models["spk_stat"].chunk(2)
+            speaker_embed.data = speaker_embed.data * std + mean
+
+    SPEAKER_TOKEN_ID = tokenizer.convert_tokens_to_ids("[spk_emb]")
+    AUDIO_EOS_TOKEN_ID = 0
+    AUDIO_PAD_TOKEN_ID = AUDIO_EOS_TOKEN_ID
+
+    optimizer = torch.optim.Adam(
+        speaker_embeds.values(), lr=1e-2, weight_decay=0, betas=[0.9, 0.95], eps=1e-5
+    )
+    loss_fn = torch.nn.CrossEntropyLoss()
+    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, epochs, 1e-7)
+
+    loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        collate_fn=AudioCollator(text_pad=tokenizer.pad_token_id),
+    )
+    logger = MetricLogger()
+    logger.create_meters(loss=None, mse_loss=None, audio_loss=None, text_loss=None)
+
+    for _epoch in range(epochs):
+        _epoch += 1
+        logger.reset()
+        header = "{blue_light}{0}: {1}{reset}".format(
+            "Epoch", output_iter(_epoch, epochs), **ansi
+        )
+        header = header.ljust(max(len("Epoch"), 30) + get_ansi_len(header))
+        iterator = logger.log_every(loader, header=header, tqdm_header="Batch")
+
+        for batch in iterator:
+            speakers = batch["speaker"]
+            text_input_ids = batch["text_input_ids"]
+            text_attention_mask = batch["text_attention_mask"]
+            audio_mel_specs = batch["audio_mel_specs"]
+            audio_attention_mask = batch["audio_attention_mask"]
+
+            batch_size, text_len = text_attention_mask.size()
+
+            dvae_audio_latents = dvae_encoder(audio_mel_specs, audio_attention_mask)
+            _, dvae_audio_input_ids = quantize(
+                dvae_decoder.vq_layer.quantizer, dvae_audio_latents
+            )
+            dvae_audio_input_ids[~audio_attention_mask.bool()] = AUDIO_PAD_TOKEN_ID
+
+            extended_audio_attention_mask = torch.cat(
+                [
+                    audio_attention_mask,
+                    torch.zeros(
+                        (batch_size, 1),
+                        dtype=audio_attention_mask.dtype,
+                        device=audio_attention_mask.device,
+                    ),
+                ],
+                dim=1,
+            )
+            extended_audio_input_ids = torch.cat(
+                [
+                    dvae_audio_input_ids,
+                    AUDIO_PAD_TOKEN_ID
+                    * torch.ones(
+                        (batch_size, 1, gpt.num_vq),
+                        dtype=dvae_audio_input_ids.dtype,
+                        device=dvae_audio_input_ids.device,
+                    ),
+                ],
+                dim=1,
+            )
+            indices = audio_attention_mask.int().sum(dim=1)
+            for i in range(batch_size):
+                extended_audio_attention_mask[i, indices[i]] = 1
+                extended_audio_input_ids[i, indices[i]] = AUDIO_EOS_TOKEN_ID
+
+            input_ids = torch.cat(
+                [
+                    text_input_ids.unsqueeze(-1).repeat(1, 1, gpt.num_vq),
+                    extended_audio_input_ids,
+                ],
+                dim=1,
+            )
+            attention_mask = torch.cat(
+                [text_attention_mask, extended_audio_attention_mask], dim=1
+            )
+            text_mask = torch.cat(
+                [
+                    torch.ones_like(text_attention_mask, dtype=bool),
+                    torch.zeros_like(extended_audio_attention_mask, dtype=bool),
+                ],
+                dim=1,
+            )
+
+            labels = input_ids.clone()
+            labels[~attention_mask.bool()] = IGNORE_TOKEN_ID
+
+            inputs_embeds = gpt.get_emb(input_ids=input_ids, text_mask=text_mask)
+
+            indices = torch.all(input_ids == SPEAKER_TOKEN_ID, dim=-1)
+            for i, speaker in enumerate(speakers):
+                inputs_embeds[i, indices[i]] = F.normalize(
+                    speaker_embeds[speaker].to(dtype=inputs_embeds.dtype),
+                    p=2.0,
+                    dim=-1,
+                    eps=1e-12,
+                ).unsqueeze(0)
+            outputs = gpt.gpt.forward(
+                inputs_embeds=inputs_embeds, attention_mask=attention_mask
+            )
+            hidden_states = outputs.last_hidden_state
+            text_hidden_states = hidden_states[:, : text_len - 1]
+            audio_hidden_states = hidden_states[:, text_len - 1 : -1]
+
+            audio_logits = torch.stack(
+                [gpt.head_code[i](audio_hidden_states) for i in range(gpt.num_vq)],
+                dim=2,
+            )
+            audio_loss = loss_fn(
+                audio_logits.flatten(0, 2), labels[:, text_len:].flatten(0, 2)
+            )
+            loss = audio_loss
+            if train_text:
+                text_logits = gpt.head_text(text_hidden_states)
+                text_loss = loss_fn(
+                    text_logits.flatten(0, 1), labels[:, 1:text_len, 0].flatten(0, 1)
+                )
+                loss += text_loss
+                logger.meters["text_loss"].update(text_loss.item(), n=batch_size)
+
+            gpt_gen_mel_specs = decoder_decoder(
+                audio_hidden_states[:, :-1].transpose(1, 2)
+            ).transpose(1, 2)
+            mse_loss = torch.nn.functional.mse_loss(gpt_gen_mel_specs, audio_mel_specs)
+            loss += 0.01 * mse_loss
+
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(speaker_embeds.values(), 1.0)
+            optimizer.step()
+            logger.meters["loss"].update(loss.item(), n=batch_size)
+            logger.meters["mse_loss"].update(mse_loss.item(), n=batch_size)
+            logger.meters["audio_loss"].update(audio_loss.item(), n=batch_size)
+        lr_scheduler.step()
+    optimizer.zero_grad()
+    return speaker_embeds
+
+
+if __name__ == "__main__":
+    import argparse
+    import os
+    import numpy as np
+    import pathlib
+    from modules.models import load_chat_tts
+    from modules.devices import devices
+    from modules import config
+    from modules.speaker import Speaker
+
+    config.runtime_env_vars.no_half = True
+    devices.reset_device()
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--save_folder", type=str, default="./")
+    parser.add_argument("--batch_size", type=int, default=16)
+    parser.add_argument("--epochs", type=int, default=100)
+    parser.add_argument("--train_text", action="store_true", help="train text loss")
+    # 初始化 speaker
+    parser.add_argument("--init_speaker", type=str)
+    parser.add_argument(
+        "--data_path",
+        type=str,
+        default="datasets/data_speaker_a/speaker_a.list",
+        help="the data_path to json/list file",
+    )
+    parser.add_argument("--tar_path", type=str, help="the tarball path with wavs")
+    parser.add_argument(
+        "--tar_in_memory", action="store_true", help="load tarball in memory"
+    )
+
+    args = parser.parse_args()
+
+    data_path: str = args.data_path
+    tar_path: str | None = args.tar_path
+    tar_in_memory: bool = args.tar_in_memory
+    train_text: bool = args.train_text
+    # gpt_lora: bool = args.gpt_lora
+    # gpt_kbit: int = args.gpt_kbit
+    save_folder: str = args.save_folder
+    batch_size: int = args.batch_size
+    epochs: int = args.epochs
+    init_speaker: str = args.init_speaker
+
+    speaker_embeds_save_path = os.path.join(save_folder, "speaker_embeds.npz")
+
+    chat = load_chat_tts()
+    dataset = XzListTar(
+        root=data_path,
+        tokenizer=chat.pretrain_models["tokenizer"],
+        vocos_model=chat.pretrain_models["vocos"],
+        tar_path=tar_path,
+        tar_in_memory=tar_in_memory,
+        device=devices.device,
+        # speakers=None,  # set(['speaker_A', 'speaker_B'])
+    )
+
+    print("len(dataset)", len(dataset))
+
+    speaker_embeds = None
+    if init_speaker:
+        spk: Speaker = Speaker.from_file(init_speaker)
+        speaker_embeds = {
+            speaker: torch.tensor(
+                spk.emb,
+                device=devices.device,
+                requires_grad=True,
+            )
+            for speaker in dataset.speakers
+        }
+
+    speaker_embeds = train_speaker_embeddings(
+        chat,
+        dataset,
+        chat.pretrain_models["gpt"],
+        batch_size=batch_size,
+        epochs=epochs,
+        train_text=train_text,
+        speaker_embeds=speaker_embeds,
+    )
+    speaker_outs = {
+        speaker: Speaker(speaker_embed.detach().cpu(), f"ep{epochs}_{speaker}")
+        for speaker, speaker_embed in speaker_embeds.items()
+    }
+    time_str = np.datetime_as_string(np.datetime64("now", "s"))
+    time_str = time_str.replace(":", "_").replace(" ", "_").replace("-", "_")
+    for speaker, speaker_out in speaker_outs.items():
+        torch.save(
+            speaker_out,
+            pathlib.Path(save_folder) / f"spk_{speaker}_{time_str}_ep{epochs}.pt",
+        )
+
+# example
+"""
+python -m modules.finetune.train_speaker \
+    --data_path datasets/data_speaker_a/speaker_a.list \
+    --save_folder ./data \
+    --init_speaker ./data/speakers/Bob.pt \
+    --epochs 100 \
+    --batch_size 6 \
+    --train_text
+"""

modules/finetune/utils/dataset.py

@@ -0,0 +1,487 @@
+import os
+import functools
+import json
+import tarfile
+import io
+import logging
+import abc
+import typing
+
+import torch.utils.data
+import torchaudio
+from torchvision.datasets.utils import download_url
+import transformers
+import vocos
+
+from modules.ChatTTS.ChatTTS.utils.infer_utils import (
+    count_invalid_characters,
+    apply_character_map,
+)
+
+
+class LazyDataType(typing.TypedDict):
+    filepath: str
+    speaker: str
+    lang: str
+    text: str
+
+
+class DataType(LazyDataType):
+    text_input_ids: torch.Tensor  # (batch_size, text_len)
+    text_attention_mask: torch.Tensor  # (batch_size, text_len)
+    audio_mel_specs: torch.Tensor  # (batch_size, audio_len*2, 100)
+    audio_attention_mask: torch.Tensor  # (batch_size, audio_len)
+
+
+class XzListTarKwargsType(typing.TypedDict):
+    tokenizer: typing.Union[transformers.PreTrainedTokenizer, None]
+    vocos_model: typing.Union[vocos.Vocos, None]
+    device: typing.Union[str, torch.device, None]
+    speakers: typing.Union[typing.Iterable[str], None]
+    sample_rate: typing.Union[int]
+    default_speaker: typing.Union[str, None]
+    default_lang: typing.Union[str, None]
+    tar_in_memory: typing.Union[bool, None]
+    process_ahead: typing.Union[bool, None]
+
+
+class AudioFolder(torch.utils.data.Dataset, abc.ABC):
+    def __init__(
+        self,
+        root: str | io.BytesIO,
+        tokenizer: transformers.PreTrainedTokenizer | None = None,
+        vocos_model: vocos.Vocos | None = None,
+        device: str | torch.device | None = None,
+        speakers: typing.Iterable[str] | None = None,
+        sample_rate: int = 24_000,
+        default_speaker: str | None = None,
+        default_lang: str | None = None,
+        tar_path: str | None = None,
+        tar_in_memory: bool = False,
+        process_ahead: bool = False,
+    ) -> None:
+        self.root = root
+        self.sample_rate = sample_rate
+        self.default_speaker = default_speaker
+        self.default_lang = default_lang
+
+        self.logger = logging.getLogger(__name__)
+        self.normalizer = {}
+
+        self.tokenizer = tokenizer
+        self.vocos = vocos_model
+        self.vocos_device = (
+            None if self.vocos is None else next(self.vocos.parameters()).device
+        )
+        self.device = device or self.vocos_device
+
+        # tar -cvf ../Xz.tar *
+        # tar -xf Xz.tar -C ./Xz
+        self.tar_path = tar_path
+        self.tar_file = None
+        self.tar_io = None
+        if tar_path is not None:
+            if tar_in_memory:
+                with open(tar_path, "rb") as f:
+                    self.tar_io = io.BytesIO(f.read())
+                self.tar_file = tarfile.open(fileobj=self.tar_io)
+            else:
+                self.tar_file = tarfile.open(tar_path)
+
+        self.lazy_data, self.speakers = self.get_lazy_data(root, speakers)
+
+        self.text_input_ids: dict[int, torch.Tensor] = {}
+        self.audio_mel_specs: dict[int, torch.Tensor] = {}
+        if process_ahead:
+            for n, item in enumerate(self.lazy_data):
+                self.audio_mel_specs[n] = self.preprocess_audio(item["filepath"])
+                self.text_input_ids[n] = self.preprocess_text(
+                    item["text"], item["lang"]
+                )
+            if self.tar_file is not None:
+                self.tar_file.close()
+            if self.tar_io is not None:
+                self.tar_io.close()
+
+    @abc.abstractmethod
+    def get_raw_data(self, root: str | io.BytesIO) -> list[dict[str, str]]: ...
+
+    @staticmethod
+    @abc.abstractmethod
+    def save_config(
+        save_path: str, lazy_data: list[LazyDataType], rel_path: str = "./"
+    ) -> None: ...
+
+    def __len__(self):
+        return len(self.lazy_data)
+
+    def __getitem__(self, n: int) -> DataType:
+        lazy_data = self.lazy_data[n]
+        if n in self.audio_mel_specs:
+            audio_mel_specs = self.audio_mel_specs[n]
+            text_input_ids = self.text_input_ids[n]
+        else:
+            audio_mel_specs = self.preprocess_audio(lazy_data["filepath"])
+            text_input_ids = self.preprocess_text(lazy_data["text"], lazy_data["lang"])
+            self.audio_mel_specs[n] = audio_mel_specs
+            self.text_input_ids[n] = text_input_ids
+            if len(self.audio_mel_specs) == len(self.lazy_data):
+                if self.tar_file is not None:
+                    self.tar_file.close()
+                if self.tar_io is not None:
+                    self.tar_io.close()
+        text_attention_mask = torch.ones(
+            len(text_input_ids), device=text_input_ids.device
+        )
+        audio_attention_mask = torch.ones(
+            (len(audio_mel_specs) + 1) // 2,
+            device=audio_mel_specs.device,
+        )
+        return {
+            "filepath": lazy_data["filepath"],
+            "speaker": lazy_data["speaker"],
+            "lang": lazy_data["lang"],
+            "text": lazy_data["text"],
+            "text_input_ids": text_input_ids,
+            "text_attention_mask": text_attention_mask,
+            "audio_mel_specs": audio_mel_specs,
+            "audio_attention_mask": audio_attention_mask,
+        }
+
+    def get_lazy_data(
+        self,
+        root: str | io.BytesIO,
+        speakers: typing.Iterable[str] | None = None,
+    ) -> tuple[list[LazyDataType], set[str]]:
+        if speakers is not None:
+            new_speakers = set(speakers)
+        else:
+            new_speakers = set()
+        lazy_data = []
+
+        raw_data = self.get_raw_data(root)
+        folder_path = os.path.dirname(root) if isinstance(root, str) else ""
+        for item in raw_data:
+            if "speaker" not in item:
+                item["speaker"] = self.default_speaker
+            if "lang" not in item:
+                item["lang"] = self.default_lang
+
+            if speakers is not None and item["speaker"] not in speakers:
+                continue
+            if speakers is None and item["speaker"] not in new_speakers:
+                new_speakers.add(item["speaker"])
+            if self.tar_file is None and isinstance(root, str):
+                filepath = os.path.join(folder_path, item["filepath"])
+            else:
+                filepath = item["filepath"]
+            lazy_data.append(
+                {
+                    "filepath": filepath,
+                    "speaker": item["speaker"],
+                    "lang": item["lang"].lower(),
+                    "text": item["text"],
+                }
+            )
+        return lazy_data, new_speakers
+
+    def preprocess_text(
+        self,
+        text: str,
+        lang: str,
+    ) -> torch.Tensor:
+        invalid_characters = count_invalid_characters(text)
+        if len(invalid_characters):
+            # self.logger.log(logging.WARNING, f'Invalid characters found! : {invalid_characters}')
+            text = apply_character_map(text)
+
+        # if not skip_refine_text:
+        #     text_tokens = refine_text(self.pretrain_models, text, **params_refine_text)['ids']
+        #     text_tokens = [i[i < self.pretrain_models['tokenizer'].convert_tokens_to_ids('[break_0]')] for i in text_tokens]
+        #     text = self.pretrain_models['tokenizer'].batch_decode(text_tokens)
+        #     if refine_text_only:
+        #         return text
+
+        text = f"[Stts][spk_emb]{text}[Ptts]"
+        # text = f'[Stts][empty_spk]{text}[Ptts]'
+
+        text_token = self.tokenizer(
+            text, return_tensors="pt", add_special_tokens=False
+        ).to(device=self.device)
+        return text_token["input_ids"].squeeze(0)
+
+    def preprocess_audio(self, filepath: str) -> torch.Tensor:
+        if self.tar_file is not None:
+            file = self.tar_file.extractfile(filepath)
+            waveform, sample_rate = torchaudio.load(file)
+        else:
+            waveform, sample_rate = torchaudio.load(filepath)
+        waveform = waveform.to(device=self.vocos_device)
+        if sample_rate != self.sample_rate:
+            waveform = torchaudio.functional.resample(
+                waveform,
+                orig_freq=sample_rate,
+                new_freq=self.sample_rate,
+            )
+        mel_spec: torch.Tensor = self.vocos.feature_extractor(waveform)
+        return (
+            mel_spec.to(device=self.device).squeeze(0).transpose(0, 1)
+        )  # (audio_len*2, 100)
+
+
+class JsonFolder(AudioFolder):
+    """
+    In json file, each item is formatted as following example:
+    `{"filepath": "path/to/file.wav", "speaker": "John", "lang": "ZH", "text": "Hello"}`.
+
+    filepath is relative to the dirname of root json file.
+    """
+
+    def get_raw_data(self, root: str | io.BytesIO) -> list[dict[str, str]]:
+        with open(root, "r", encoding="utf-8") as f:
+            raw_data = json.load(f)
+        return raw_data
+
+    @staticmethod
+    def save_config(
+        save_path: str, lazy_data: list[LazyDataType], rel_path: str = "./"
+    ) -> None:
+        save_data = [item.copy() for item in lazy_data]
+        for item in save_data:
+            item["filepath"] = os.path.relpath(item["filepath"], rel_path)
+        with open(save_path, "w", encoding="utf-8") as f:
+            json.dump(save_data, f, ensure_ascii=False, indent=4)
+
+
+class ListFolder(AudioFolder):
+    """
+    In list file, each row is formatted as `filepath|speaker|lang|text` with `|` as separator.
+    `path/to/file.wav|John|ZH|Hello`.
+
+    filepath is relative to the dirname of root list file.
+    """
+
+    def get_raw_data(self, root: str | io.BytesIO) -> list[dict[str, str]]:
+        raw_data = []
+        with open(root, "r", encoding="utf-8") as f:
+            for line in f.readlines():
+                line = line.strip().removesuffix("\n")
+                if len(line) == 0:
+                    continue
+                filepath, speaker, lang, text = line.split(sep="|", maxsplit=3)
+                raw_data.append(
+                    {
+                        "text": text,
+                        "filepath": filepath,
+                        "speaker": speaker,
+                        "lang": lang,
+                    }
+                )
+        return raw_data
+
+    @staticmethod
+    def save_config(
+        save_path: str, lazy_data: list[LazyDataType], rel_path: str = "./"
+    ) -> None:
+        save_data = [item.copy() for item in lazy_data]
+        for item in save_data:
+            item["filepath"] = os.path.relpath(item["filepath"], rel_path)
+        with open(save_path, "w", encoding="utf-8") as f:
+            for item in save_data:
+                f.write(
+                    f"{item['filepath']}|{item['speaker']}|{item['lang']}|{item['text']}\n"
+                )
+
+
+class XzListTar(ListFolder):
+    def __init__(
+        self,
+        *args,
+        root: str | io.BytesIO,
+        tar_path: str | None = None,
+        **kwargs,
+    ):
+        if isinstance(root, io.BytesIO):
+            assert tar_path is not None
+        else:
+            # make sure root is a list file
+            if not root.endswith(".list"):  # folder case
+                if os.path.isfile(root):
+                    raise FileExistsError(f"{root} is a file!")
+                elif not os.path.exists(root):
+                    os.makedirs(root)
+                root = os.path.join(root, "all.list")
+        if isinstance(root, str) and not os.path.isfile(root):
+            # prepare all.list
+            self.concat_dataset(
+                save_folder=os.path.dirname(root),
+                langs=kwargs.get("langs", ["zh", "en"]),
+            )
+
+        super().__init__(root, *args, tar_path=tar_path, **kwargs)
+
+    def concat_dataset(
+        self, save_folder: str | None = None, langs: list[str] = ["zh", "en"]
+    ) -> None:
+        if save_folder is None:
+            save_folder = os.path.dirname(self.root)
+        if os.path.isfile(save_folder):
+            raise FileExistsError(f"{save_folder} already exists as a file!")
+        elif not os.path.exists(save_folder):
+            os.makedirs(save_folder)
+        lazy_data = []
+
+        for member in self.tar_file.getmembers():
+            if not member.isfile():
+                continue
+            if member.name.endswith(".list"):
+                print(member.name)
+                root_io = self.tar_file.extractfile(member)
+                lazy_data += ListFolder(root_io).lazy_data
+            if member.name.endswith(".json"):
+                print(member.name)
+                root_io = self.tar_file.extractfile(member)
+                lazy_data += JsonFolder(root_io).lazy_data
+        if langs is not None:
+            lazy_data = [item for item in lazy_data if item["lang"] in langs]
+        ListFolder.save_config(os.path.join(save_folder, "all.list"), lazy_data)
+        JsonFolder.save_config(os.path.join(save_folder, "all.json"), lazy_data)
+        print(f"all.list and all.json are saved to {save_folder}")
+
+
+class XzListFolder(ListFolder):
+    """
+    [Xz乔希](https://space.bilibili.com/5859321)
+
+    Only look at the basename of filepath in list file. Previous folder paths are ignored.
+    Files are organized as `[list basename]/[file basename]`
+
+    Example tree structure:
+
+    [folder]
+    ├── speaker_A
+    │   ├── 1.wav
+    │   └── 2.wav
+    ├── speaker_A.list
+    ├── speaker_B
+    │   ├── 1.wav
+    │   └── 2.wav
+    └── speaker_B.list
+    """
+
+    def get_raw_data(self, root: str | io.BytesIO) -> list[dict[str, str]]:
+        raw_data = super().get_raw_data(root)
+        for item in raw_data:
+            item["filepath"] = os.path.join(
+                os.path.basename(root).removesuffix(".list"),
+                os.path.basename(item["filepath"]),
+            )
+        return raw_data
+
+
+class AudioCollator:
+    def __init__(self, text_pad: int = 0, audio_pad: int = 0):
+        self.text_pad = text_pad
+        self.audio_pad = audio_pad
+
+    def __call__(self, batch: list[DataType]):
+        batch = [x for x in batch if x is not None]
+
+        audio_maxlen = max(len(item["audio_attention_mask"]) for item in batch)
+        text_maxlen = max(len(item["text_attention_mask"]) for item in batch)
+
+        filepath = []
+        speaker = []
+        lang = []
+        text = []
+        text_input_ids = []
+        text_attention_mask = []
+        audio_mel_specs = []
+        audio_attention_mask = []
+
+        for x in batch:
+            filepath.append(x["filepath"])
+            speaker.append(x["speaker"])
+            lang.append(x["lang"])
+            text.append(x["text"])
+            text_input_ids.append(
+                torch.nn.functional.pad(
+                    x["text_input_ids"],
+                    (text_maxlen - len(x["text_input_ids"]), 0),
+                    value=self.text_pad,
+                )
+            )
+            text_attention_mask.append(
+                torch.nn.functional.pad(
+                    x["text_attention_mask"],
+                    (text_maxlen - len(x["text_attention_mask"]), 0),
+                    value=0,
+                )
+            )
+            audio_mel_specs.append(
+                torch.nn.functional.pad(
+                    x["audio_mel_specs"],
+                    (0, 0, 0, audio_maxlen * 2 - len(x["audio_mel_specs"])),
+                    value=self.audio_pad,
+                )
+            )
+            audio_attention_mask.append(
+                torch.nn.functional.pad(
+                    x["audio_attention_mask"],
+                    (0, audio_maxlen - len(x["audio_attention_mask"])),
+                    value=0,
+                )
+            )
+        return {
+            "filepath": filepath,
+            "speaker": speaker,
+            "lang": lang,
+            "text": text,
+            "text_input_ids": torch.stack(text_input_ids),
+            "text_attention_mask": torch.stack(text_attention_mask),
+            "audio_mel_specs": torch.stack(audio_mel_specs),
+            "audio_attention_mask": torch.stack(audio_attention_mask),
+        }
+
+
+def formalize_xz_list(src_folder: str):
+    for root, _, files in os.walk(src_folder):
+        for file in files:
+            if file.endswith(".list"):
+                filepath = os.path.join(root, file)
+                print(filepath)
+                lazy_data = XzListFolder(filepath).lazy_data
+                XzListFolder.save_config(filepath, lazy_data, rel_path=src_folder)
+
+
+def concat_dataset(
+    src_folder: str, save_folder: str | None = None, langs: list[str] = ["zh", "en"]
+) -> None:
+    if save_folder is None:
+        save_folder = src_folder
+    if os.path.isfile(save_folder):
+        raise FileExistsError(f"{save_folder} already exists as a file!")
+    elif not os.path.exists(save_folder):
+        os.makedirs(save_folder)
+    lazy_data = []
+    same_folder = os.path.samefile(src_folder, save_folder)
+    for root, _, files in os.walk(src_folder):
+        for file in files:
+            filepath = os.path.join(root, file)
+            if same_folder and file in ("all.list", "all.json"):
+                continue
+            if file.endswith(".list"):
+                print(filepath)
+                lazy_data += ListFolder(filepath).lazy_data
+            if file.endswith(".json"):
+                print(filepath)
+                lazy_data += JsonFolder(filepath).lazy_data
+    if langs is not None:
+        lazy_data = [item for item in lazy_data if item["lang"] in langs]
+    ListFolder.save_config(
+        os.path.join(save_folder, "all.list"), lazy_data, rel_path=save_folder
+    )
+    JsonFolder.save_config(
+        os.path.join(save_folder, "all.json"), lazy_data, rel_path=save_folder
+    )
+    print(f"all.list and all.json are saved to {save_folder}")

modules/finetune/utils/logger.py

@@ -0,0 +1,409 @@
+#!/usr/bin/env python3
+
+import statistics
+import time
+from collections import defaultdict, deque
+from tqdm import tqdm as tqdm_class
+
+from typing import Generator, Iterable, TypeVar
+from typing_extensions import Self
+
+import torch
+import torch.distributed as dist
+
+from .output import ansi, prints, get_ansi_len
+
+__all__ = ["SmoothedValue", "MetricLogger"]
+
+MB = 1 << 20
+T = TypeVar("T")
+
+
+class SmoothedValue:
+    r"""Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+
+    See Also:
+        https://github.com/pytorch/vision/blob/main/references/classification/utils.py
+
+    Args:
+        name (str): Name string.
+        window_size (int): The :attr:`maxlen` of :class:`~collections.deque`.
+        fmt (str): The format pattern of ``str(self)``.
+
+    Attributes:
+        name (str): Name string.
+        fmt (str): The string pattern.
+        deque (~collections.deque): The unique data series.
+        count (int): The amount of data.
+        total (float): The sum of all data.
+
+        median (float): The median of :attr:`deque`.
+        avg (float): The avg of :attr:`deque`.
+        global_avg (float): :math:`\frac{\text{total}}{\text{count}}`
+        max (float): The max of :attr:`deque`.
+        min (float): The min of :attr:`deque`.
+        last_value (float): The last value of :attr:`deque`.
+    """
+
+    def __init__(
+        self, name: str = "", window_size: int = None, fmt: str = "{global_avg:.3f}"
+    ):
+        self.name = name
+        self.deque: deque[float] = deque(maxlen=window_size)
+        self.count: int = 0
+        self.total: float = 0.0
+        self.fmt = fmt
+
+    def update(self, value: float, n: int = 1) -> Self:
+        r"""Update :attr:`n` pieces of data with same :attr:`value`.
+
+        .. code-block:: python
+
+            self.deque.append(value)
+            self.total += value * n
+            self.count += n
+
+        Args:
+            value (float): the value to update.
+            n (int): the number of data with same :attr:`value`.
+
+        Returns:
+            SmoothedValue: return ``self`` for stream usage.
+        """
+        self.deque.append(value)
+        self.total += value * n
+        self.count += n
+        return self
+
+    def update_list(self, value_list: list[float]) -> Self:
+        r"""Update :attr:`value_list`.
+
+        .. code-block:: python
+
+            for value in value_list:
+                self.deque.append(value)
+                self.total += value
+            self.count += len(value_list)
+
+        Args:
+            value_list (list[float]): the value list to update.
+
+        Returns:
+            SmoothedValue: return ``self`` for stream usage.
+        """
+        for value in value_list:
+            self.deque.append(value)
+            self.total += value
+        self.count += len(value_list)
+        return self
+
+    def reset(self) -> Self:
+        r"""Reset ``deque``, ``count`` and ``total`` to be empty.
+
+        Returns:
+            SmoothedValue: return ``self`` for stream usage.
+        """
+        self.deque = deque(maxlen=self.deque.maxlen)
+        self.count = 0
+        self.total = 0.0
+        return self
+
+    def synchronize_between_processes(self):
+        r"""
+        Warning:
+            Does NOT synchronize the deque!
+        """
+        if not (dist.is_available() and dist.is_initialized()):
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = float(t[1])
+
+    @property
+    def median(self) -> float:
+        try:
+            return statistics.median(self.deque)
+        except Exception:
+            return 0.0
+
+    @property
+    def avg(self) -> float:
+        try:
+            return statistics.mean(self.deque)
+        except Exception:
+            return 0.0
+
+    @property
+    def global_avg(self) -> float:
+        try:
+            return self.total / self.count
+        except Exception:
+            return 0.0
+
+    @property
+    def max(self) -> float:
+        try:
+            return max(self.deque)
+        except Exception:
+            return 0.0
+
+    @property
+    def min(self) -> float:
+        try:
+            return min(self.deque)
+        except Exception:
+            return 0.0
+
+    @property
+    def last_value(self) -> float:
+        try:
+            return self.deque[-1]
+        except Exception:
+            return 0.0
+
+    def __str__(self):
+        return self.fmt.format(
+            name=self.name,
+            count=self.count,
+            total=self.total,
+            median=self.median,
+            avg=self.avg,
+            global_avg=self.global_avg,
+            min=self.min,
+            max=self.max,
+            last_value=self.last_value,
+        )
+
+    def __format__(self, format_spec: str) -> str:
+        return self.__str__()
+
+
+class MetricLogger:
+    r"""
+    See Also:
+        https://github.com/pytorch/vision/blob/main/references/classification/utils.py
+
+    Args:
+        delimiter (str): The delimiter to join different meter strings.
+            Defaults to ``''``.
+        meter_length (int): The minimum length for each meter.
+            Defaults to ``20``.
+        tqdm (bool): Whether to use tqdm to show iteration information.
+            Defaults to ``env['tqdm']``.
+        indent (int): The space indent for the entire string.
+            Defaults to ``0``.
+
+    Attributes:
+        meters (dict[str, SmoothedValue]): The meter dict.
+        iter_time (SmoothedValue): Iteration time meter.
+        data_time (SmoothedValue): Data loading time meter.
+        memory (SmoothedValue): Memory usage meter.
+    """
+
+    def __init__(
+        self,
+        delimiter: str = "",
+        meter_length: int = 20,
+        tqdm: bool = True,
+        indent: int = 0,
+        **kwargs,
+    ):
+        self.meters: defaultdict[str, SmoothedValue] = defaultdict(SmoothedValue)
+        self.create_meters(**kwargs)
+        self.delimiter = delimiter
+        self.meter_length = meter_length
+        self.tqdm = tqdm
+        self.indent = indent
+
+        self.iter_time = SmoothedValue()
+        self.data_time = SmoothedValue()
+        self.memory = SmoothedValue(fmt="{max:.0f}")
+
+    def create_meters(self, **kwargs: str) -> Self:
+        r"""Create meters with specific ``fmt`` in :attr:`self.meters`.
+
+        ``self.meters[meter_name] = SmoothedValue(fmt=fmt)``
+
+        Args:
+            **kwargs: ``(meter_name: fmt)``
+
+        Returns:
+            MetricLogger: return ``self`` for stream usage.
+        """
+        for k, v in kwargs.items():
+            self.meters[k] = SmoothedValue(fmt="{global_avg:.3f}" if v is None else v)
+        return self
+
+    def update(self, n: int = 1, **kwargs: float) -> Self:
+        r"""Update values to :attr:`self.meters` by calling :meth:`SmoothedValue.update()`.
+
+        ``self.meters[meter_name].update(float(value), n=n)``
+
+        Args:
+            n (int): the number of data with same value.
+            **kwargs: ``{meter_name: value}``.
+
+        Returns:
+            MetricLogger: return ``self`` for stream usage.
+        """
+        for k, v in kwargs.items():
+            if k not in self.meters:
+                self.meters[k] = SmoothedValue()
+            self.meters[k].update(float(v), n=n)
+        return self
+
+    def update_list(self, **kwargs: list) -> Self:
+        r"""Update values to :attr:`self.meters` by calling :meth:`SmoothedValue.update_list()`.
+
+        ``self.meters[meter_name].update_list(value_list)``
+
+        Args:
+            **kwargs: ``{meter_name: value_list}``.
+
+        Returns:
+            MetricLogger: return ``self`` for stream usage.
+        """
+        for k, v in kwargs.items():
+            self.meters[k].update_list(v)
+        return self
+
+    def reset(self) -> Self:
+        r"""Reset meter in :attr:`self.meters` by calling :meth:`SmoothedValue.reset()`.
+
+        Returns:
+            MetricLogger: return ``self`` for stream usage.
+        """
+        for meter in self.meters.values():
+            meter.reset()
+        return self
+
+    def get_str(self, cut_too_long: bool = True, strip: bool = True, **kwargs) -> str:
+        r"""Generate formatted string based on keyword arguments.
+
+        ``key: value`` with max length to be :attr:`self.meter_length`.
+
+        Args:
+            cut_too_long (bool): Whether to cut too long values to first 5 characters.
+                Defaults to ``True``.
+            strip (bool): Whether to strip trailing whitespaces.
+                Defaults to ``True``.
+            **kwargs: Keyword arguments to generate string.
+        """
+        str_list: list[str] = []
+        for k, v in kwargs.items():
+            v_str = str(v)
+            _str: str = "{green}{k}{reset}: {v}".format(k=k, v=v_str, **ansi)
+            max_length = self.meter_length + get_ansi_len(_str)
+            if cut_too_long:
+                _str = _str[:max_length]
+            str_list.append(_str.ljust(max_length))
+        _str = self.delimiter.join(str_list)
+        if strip:
+            _str = _str.rstrip()
+        return _str
+
+    def __getattr__(self, attr: str) -> float:
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in vars(self):  # TODO: use hasattr
+            return vars(self)[attr]
+        raise AttributeError(
+            "'{}' object has no attribute '{}'".format(type(self).__name__, attr)
+        )
+
+    def __str__(self) -> str:
+        return self.get_str(**self.meters)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def log_every(
+        self,
+        iterable: Iterable[T],
+        header: str = "",
+        tqdm: bool = None,
+        tqdm_header: str = "Iter",
+        indent: int = None,
+        verbose: int = 1,
+    ) -> Generator[T, None, None]:
+        r"""Wrap an :class:`collections.abc.Iterable` with formatted outputs.
+
+        * Middle Output:
+          ``{tqdm_header}: [ current / total ] str(self) {memory} {iter_time} {data_time} {time}<{remaining}``
+        * Final Output
+          ``{header} str(self) {memory} {iter_time} {data_time} {total_time}``
+
+        Args:
+            iterable (~collections.abc.Iterable): The raw iterator.
+            header (str): The header string for final output.
+                Defaults to ``''``.
+            tqdm (bool): Whether to use tqdm to show iteration information.
+                Defaults to ``self.tqdm``.
+            tqdm_header (str): The header string for middle output.
+                Defaults to ``'Iter'``.
+            indent (int): The space indent for the entire string.
+                if ``None``, use ``self.indent``.
+                Defaults to ``None``.
+            verbose (int): The verbose level of output information.
+        """
+        tqdm = tqdm if tqdm is not None else self.tqdm
+        indent = indent if indent is not None else self.indent
+        iterator = iterable
+        if len(header) != 0:
+            header = header.ljust(30 + get_ansi_len(header))
+        if tqdm:
+            length = len(str(len(iterable)))
+            pattern: str = (
+                "{tqdm_header}: {blue_light}"
+                "[ {red}{{n_fmt:>{length}}}{blue_light} "
+                "/ {red}{{total_fmt}}{blue_light} ]{reset}"
+            ).format(tqdm_header=tqdm_header, length=length, **ansi)
+            offset = len(f"{{n_fmt:>{length}}}{{total_fmt}}") - 2 * length
+            pattern = pattern.ljust(30 + offset + get_ansi_len(pattern))
+            time_str = self.get_str(time="{elapsed}<{remaining}", cut_too_long=False)
+            bar_format = f"{pattern}{{desc}}{time_str}"
+            iterator = tqdm_class(iterable, leave=False, bar_format=bar_format)
+
+        self.iter_time.reset()
+        self.data_time.reset()
+        self.memory.reset()
+
+        end = time.time()
+        start_time = time.time()
+        for obj in iterator:
+            cur_data_time = time.time() - end
+            self.data_time.update(cur_data_time)
+            yield obj
+            cur_iter_time = time.time() - end
+            self.iter_time.update(cur_iter_time)
+            if torch.cuda.is_available():
+                cur_memory = torch.cuda.max_memory_allocated() / MB
+                self.memory.update(cur_memory)
+            if tqdm:
+                _dict = {k: v for k, v in self.meters.items()}
+                if verbose > 2 and torch.cuda.is_available():
+                    _dict.update(memory=f"{cur_memory:.0f} MB")
+                if verbose > 1:
+                    _dict.update(
+                        iter=f"{cur_iter_time:.3f} s", data=f"{cur_data_time:.3f} s"
+                    )
+                iterator.set_description_str(self.get_str(**_dict, strip=False))
+            end = time.time()
+        self.synchronize_between_processes()
+        total_time = time.time() - start_time
+        total_time_str = tqdm_class.format_interval(total_time)
+
+        _dict = {k: v for k, v in self.meters.items()}
+        if verbose > 2 and torch.cuda.is_available():
+            _dict.update(memory=f"{str(self.memory)} MB")
+        if verbose > 1:
+            _dict.update(
+                iter=f"{str(self.iter_time)} s", data=f"{str(self.data_time)} s"
+            )
+        _dict.update(time=total_time_str)
+        prints(self.delimiter.join([header, self.get_str(**_dict)]), indent=indent)

modules/finetune/utils/model.py

@@ -0,0 +1,19 @@
+import torch
+from einops import rearrange
+from vector_quantize_pytorch.residual_fsq import GroupedResidualFSQ
+
+
+def quantize(
+    quantizer: GroupedResidualFSQ,
+    audio_latents: torch.Tensor,  # (batch_size, audio_len, audio_dim=1024)
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # feat shape (batch_size, audio_len, audio_dim)
+    # ind shape (GFSQ.G, batch_size, audio_len, GFSQ.R)
+    # num_vq=GFSQ.G*GFSQ.R
+    feat, ind = quantizer(audio_latents)
+    audio_quantized_latents = feat  # (batch_size, audio_len, audio_dim)
+    audio_input_ids = rearrange(  # (batch_size, audio_len, num_vq)
+        ind,
+        "g b t r ->b t (g r)",
+    )
+    return audio_quantized_latents, audio_input_ids

modules/finetune/utils/output.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+
+import re
+import sys
+from contextlib import contextmanager
+
+
+class ANSI:
+    ansi_color = {
+        "black": "\033[30m",
+        "red": "\033[31m",
+        "green": "\033[32m",
+        "yellow": "\033[33m",
+        "blue": "\033[34m",
+        "purple": "\033[35m",
+        "blue_light": "\033[36m",
+        "white": "\033[37m",
+        "reset": "\033[0m",
+        "upline": "\033[1A",
+        "clear_line": "\033[2K",
+        "clear": "\033[2J",
+    }
+    ansi_nocolor = {
+        "black": "",
+        "red": "",
+        "green": "",
+        "yellow": "",
+        "blue": "",
+        "purple": "",
+        "blue_light": "",
+        "white": "",
+        "reset": "",
+        "upline": "\033[1A\033[",
+        "clear_line": "\033[K",
+        "clear": "\033[2J",
+    }
+
+    def __init__(self):
+        self._dict = ANSI.ansi_color if ("--color" in sys.argv) else ANSI.ansi_nocolor
+
+    def switch(self, color: bool):
+        self._dict = ANSI.ansi_color if color else ANSI.ansi_nocolor
+
+    def keys(self):
+        return self._dict.keys()
+
+    def items(self):
+        return self._dict.items()
+
+    def __getitem__(self, key):
+        return self._dict[key]
+
+    def __str__(self):
+        return str(self._dict)
+
+    def __repr__(self):
+        return repr(self._dict)
+
+
+ansi = ANSI()
+
+
+def remove_ansi(s: str) -> str:
+    ansi_escape = re.compile(r"(\x9B|\x1B\[)[0-?]*[ -\/]*[@-~]")
+    return ansi_escape.sub("", s)
+
+
+def get_ansi_len(s: str) -> int:
+    return len(s) - len(remove_ansi(s))
+
+
+def prints(*args: str, indent: int = 0, prefix: str = "", **kwargs):
+    assert indent >= 0
+    new_args = []
+    for arg in args:
+        new_args.append(indent_str(str(arg), indent=indent))
+    if len(new_args):
+        new_args[0] = prefix + str(new_args[0])
+    print(*new_args, **kwargs)
+
+
+def output_iter(_iter: int, iteration: int = None, iter_len: int = 4) -> str:
+    if iteration is None:
+        pattern = "{blue_light}[ {red}{0}{blue_light} ]{reset}"
+        return pattern.format(str(_iter).rjust(iter_len), **ansi)
+    else:
+        iter_str = str(iteration)
+        length = len(iter_str)
+        pattern = (
+            "{blue_light}[ {red}{0}{blue_light} " "/ {red}{1}{blue_light} ]{reset}"
+        )
+        return pattern.format(str(_iter).rjust(length), iter_str, **ansi)
+
+
+def indent_str(s_: str, indent: int = 0) -> str:
+    # modified from torch.nn.modules._addindent
+    if indent > 0 and s_:
+        s_ = indent * " " + str(s_[:-1]).replace("\n", "\n" + indent * " ") + s_[-1]
+    return s_
+
+
+class IndentRedirect:  # TODO: inherit TextIOWrapper?
+    def __init__(self, buffer: bool = True, indent: int = 0):
+        self.__console__ = sys.stdout
+        self.indent = indent
+        self.__buffer: str = None
+        if buffer:
+            self.__buffer = ""
+
+    def write(self, text: str, indent: int = None):
+        indent = indent if indent is not None else self.indent
+        text = indent_str(text, indent=indent)
+        if self.__buffer is None:
+            self.__console__.write(text)
+        else:
+            self.__buffer += text
+
+    def flush(self):
+        if self.__buffer is not None:
+            self.__console__.write(self.__buffer)
+            self.__buffer = ""
+        self.__console__.flush()
+
+    @contextmanager
+    def __call__(self) -> None:
+        try:
+            sys.stdout = self
+            yield
+        finally:
+            sys.stdout = self.__console__
+            self.__buffer = ""
+
+    def enable(self):
+        sys.stdout = self
+
+    def disable(self):
+        if self.__buffer is not None:
+            self.__buffer = ""
+        sys.stdout = self.__console__
+
+    @property
+    def buffer(self) -> str:
+        return self.__buffer
+
+
+redirect = IndentRedirect()

modules/generate_audio.py

@@ -76,6 +76,8 @@ def generate_audio_batch(
             params_infer_code["spk_emb"] = chat_tts.sample_random_speaker()
         logger.debug(("spk", spk))
     elif isinstance(spk, Speaker):
+        if not isinstance(spk.emb, torch.Tensor):
+            raise ValueError("spk.pt is broken, please retrain the model.")
         params_infer_code["spk_emb"] = spk.emb
         logger.debug(("spk", spk.name))
     else:

modules/normalization.py

@@ -120,6 +120,7 @@ character_map = {
     "~": " ",
     "～": " ",
     "/": " ",
+    "·": " ",
 }
 
 character_to_word = {
@@ -282,6 +283,9 @@ def text_normalize(text, is_end=False):
 
 
 if __name__ == "__main__":
+    from modules.devices import devices
+
+    devices.reset_device()
     test_cases = [
         "ChatTTS是专门为对话场景设计的文本转语音模型，例如LLM助手对话任务。它支持英文和中文两种语言。最大的模型使用了10万小时以上的中英文数据进行训练。在HuggingFace中开源的版本为4万小时训练且未SFT的版本.",
         " [oral_9] [laugh_0] [break_0] 电 [speed_0] 影 [speed_0] 中 梁朝伟 [speed_9] 扮演的陈永仁的编号27149",
@@ -319,6 +323,7 @@ State-of-the-art Machine Learning for PyTorch, TensorFlow, and JAX.
         """
 120米
 有12%的概率会下雨
+埃隆·马斯克
 """,
     ]
 

modules/repos_static/resemble_enhance/data/distorter/base.py

@@ -2,6 +2,7 @@ import itertools
 import os
 import random
 import time
+from typing import Union
 import warnings
 
 import numpy as np
@@ -87,7 +88,7 @@ class Choice(Effect):
 
 
 class Permutation(Effect):
-    def __init__(self, *effects, n: int | None = None):
+    def __init__(self, *effects, n: Union[int, None] = None):
         super().__init__()
         self.effects = effects
         self.n = n

modules/repos_static/resemble_enhance/data/distorter/custom.py

@@ -3,6 +3,7 @@ import random
 from dataclasses import dataclass
 from functools import cached_property
 from pathlib import Path
+from typing import Union
 
 import librosa
 import numpy as np
@@ -16,7 +17,7 @@ _logger = logging.getLogger(__name__)
 
 @dataclass
 class RandomRIR(Effect):
-    rir_dir: Path | None
+    rir_dir: Union[Path, None]
     rir_rate: int = 44_000
     rir_suffix: str = ".npy"
     deterministic: bool = False
@@ -49,7 +50,9 @@ class RandomRIR(Effect):
 
         length = len(wav)
 
-        wav = librosa.resample(wav, orig_sr=sr, target_sr=self.rir_rate, res_type="kaiser_fast")
+        wav = librosa.resample(
+            wav, orig_sr=sr, target_sr=self.rir_rate, res_type="kaiser_fast"
+        )
         rir = self._sample_rir()
 
         wav = signal.convolve(wav, rir, mode="same")
@@ -58,7 +61,9 @@ class RandomRIR(Effect):
         if actlev > 0.99:
             wav = (wav / actlev) * 0.98
 
-        wav = librosa.resample(wav, orig_sr=self.rir_rate, target_sr=sr, res_type="kaiser_fast")
+        wav = librosa.resample(
+            wav, orig_sr=self.rir_rate, target_sr=sr, res_type="kaiser_fast"
+        )
 
         if abs(length - len(wav)) > 10:
             _logger.warning(f"length mismatch: {length} vs {len(wav)}")

modules/repos_static/resemble_enhance/data/distorter/sox.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import random
+from typing import Union
 import warnings
 from functools import partial
 
@@ -29,7 +30,9 @@ class AttachableEffect(Effect):
         chain = augment.EffectChain()
         chain = self.attach(chain)
         tensor = torch.from_numpy(wav)[None].float()  # (1, T)
-        tensor = chain.apply(tensor, src_info={"rate": sr}, target_info={"channels": 1, "rate": sr})
+        tensor = chain.apply(
+            tensor, src_info={"rate": sr}, target_info={"channels": 1, "rate": sr}
+        )
         wav = tensor.numpy()[0]  # (T,)
         return wav
 
@@ -41,7 +44,9 @@ class SoxEffect(AttachableEffect):
         self.kwargs = kwargs
 
     def attach(self, chain: augment.EffectChain) -> augment.EffectChain:
-        _logger.debug(f"Attaching {self.effect_name} with {self.args} and {self.kwargs}")
+        _logger.debug(
+            f"Attaching {self.effect_name} with {self.args} and {self.kwargs}"
+        )
         if not hasattr(chain, self.effect_name):
             raise ValueError(f"EffectChain has no attribute {self.effect_name}")
         return getattr(chain, self.effect_name)(*self.args, **self.kwargs)
@@ -115,21 +120,30 @@ class Randint(Generator):
 
 
 class Concat(Generator):
-    def __init__(self, *parts: Generator | str):
+    def __init__(self, *parts: Union[Generator, str]):
         self.parts = parts
 
     def __call__(self):
-        return "".join([part if isinstance(part, str) else part() for part in self.parts])
+        return "".join(
+            [part if isinstance(part, str) else part() for part in self.parts]
+        )
 
 
 class RandomLowpassDistorter(SoxEffect):
     def __init__(self, low=2000, high=16000):
-        super().__init__("sinc", "-n", Randint(50, 200), Concat("-", Uniform(low, high)))
+        super().__init__(
+            "sinc", "-n", Randint(50, 200), Concat("-", Uniform(low, high))
+        )
 
 
 class RandomBandpassDistorter(SoxEffect):
     def __init__(self, low=100, high=1000, min_width=2000, max_width=4000):
-        super().__init__("sinc", "-n", Randint(50, 200), partial(self._fn, low, high, min_width, max_width))
+        super().__init__(
+            "sinc",
+            "-n",
+            Randint(50, 200),
+            partial(self._fn, low, high, min_width, max_width),
+        )
 
     @staticmethod
     def _fn(low, high, min_width, max_width):
@@ -139,7 +153,15 @@ class RandomBandpassDistorter(SoxEffect):
 
 
 class RandomEqualizer(SoxEffect):
-    def __init__(self, low=100, high=4000, q_low=1, q_high=5, db_low: int = -30, db_high: int = 30):
+    def __init__(
+        self,
+        low=100,
+        high=4000,
+        q_low=1,
+        q_high=5,
+        db_low: int = -30,
+        db_high: int = 30,
+    ):
         super().__init__(
             "equalizer",
             Uniform(low, high),
@@ -150,7 +172,9 @@ class RandomEqualizer(SoxEffect):
 
 class RandomOverdrive(SoxEffect):
     def __init__(self, gain_low=5, gain_high=40, colour_low=20, colour_high=80):
-        super().__init__("overdrive", Uniform(gain_low, gain_high), Uniform(colour_low, colour_high))
+        super().__init__(
+            "overdrive", Uniform(gain_low, gain_high), Uniform(colour_low, colour_high)
+        )
 
 
 class RandomReverb(Chain):

modules/repos_static/resemble_enhance/data/utils.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Callable
+from typing import Callable, Union
 
 from torch import Tensor
 
@@ -16,7 +16,9 @@ def rglob_audio_files(path: Path):
     return list(walk_paths(path, ".wav")) + list(walk_paths(path, ".flac"))
 
 
-def mix_fg_bg(fg: Tensor, bg: Tensor, alpha: float | Callable[..., float] = 0.5, eps=1e-7):
+def mix_fg_bg(
+    fg: Tensor, bg: Tensor, alpha: Union[float, Callable[..., float]] = 0.5, eps=1e-7
+):
     """
     Args:
         fg: (b, t)

modules/repos_static/resemble_enhance/denoiser/denoiser.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Union
 
 import torch
 import torch.nn.functional as F
@@ -154,7 +155,7 @@ class Denoiser(nn.Module):
         sep_sin = sin * cos_res + cos * sin_res
         return sep_mag, sep_cos, sep_sin
 
-    def forward(self, x: Tensor, y: Tensor | None = None):
+    def forward(self, x: Tensor, y: Union[Tensor, None] = None):
         """
         Args:
             x: (b t), a mixed audio

modules/repos_static/resemble_enhance/enhancer/download.py

@@ -1,5 +1,6 @@
 import logging
 from pathlib import Path
+from typing import Union
 
 import torch
 
@@ -12,14 +13,18 @@ def get_source_url(relpath):
     return f"https://huggingface.co/ResembleAI/resemble-enhance/resolve/main/{RUN_NAME}/{relpath}?download=true"
 
 
-def get_target_path(relpath: str | Path, run_dir: str | Path | None = None):
+def get_target_path(relpath: Union[str, Path], run_dir: Union[str, Path, None] = None):
     if run_dir is None:
         run_dir = Path(__file__).parent.parent / "model_repo" / RUN_NAME
     return Path(run_dir) / relpath
 
 
-def download(run_dir: str | Path | None = None):
-    relpaths = ["hparams.yaml", "ds/G/latest", "ds/G/default/mp_rank_00_model_states.pt"]
+def download(run_dir: Union[str, Path, None] = None):
+    relpaths = [
+        "hparams.yaml",
+        "ds/G/latest",
+        "ds/G/default/mp_rank_00_model_states.pt",
+    ]
     for relpath in relpaths:
         path = get_target_path(relpath, run_dir=run_dir)
         if path.exists():

modules/repos_static/resemble_enhance/enhancer/enhancer.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Union
 
 import matplotlib.pyplot as plt
 import pandas as pd
@@ -109,7 +110,7 @@ class Enhancer(nn.Module):
             return self.mel_fn(x)[..., :-1]  # (b d t)
         return self.mel_fn(x)
 
-    def _may_denoise(self, x: Tensor, y: Tensor | None = None):
+    def _may_denoise(self, x: Tensor, y: Union[Tensor, None] = None):
         if self.hp.lcfm_training_mode == "cfm":
             return self.denoiser(x, y)
         return x
@@ -126,7 +127,9 @@ class Enhancer(nn.Module):
         self.lcfm.eval_tau_(tau)
         self._eval_lambd = lambd
 
-    def forward(self, x: Tensor, y: Tensor | None = None, z: Tensor | None = None):
+    def forward(
+        self, x: Tensor, y: Union[Tensor, None] = None, z: Union[Tensor, None] = None
+    ):
         """
         Args:
             x: (b t), mix wavs (fg + bg)

modules/repos_static/resemble_enhance/enhancer/hparams.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Union
 
 from ..hparams import HParams as HParamsBase
 
@@ -17,7 +18,7 @@ class HParams(HParamsBase):
 
     vocoder_extra_dim: int = 32
 
-    gan_training_start_step: int | None = 5_000
-    enhancer_stage1_run_dir: Path | None = None
+    gan_training_start_step: Union[int, None] = 5_000
+    enhancer_stage1_run_dir: Union[Path, None] = None
 
-    denoiser_run_dir: Path | None = None
+    denoiser_run_dir: Union[Path, None] = None

modules/repos_static/resemble_enhance/enhancer/inference.py

@@ -1,6 +1,7 @@
 import logging
 from functools import cache
 from pathlib import Path
+from typing import Union
 
 import torch
 
@@ -13,7 +14,7 @@ logger = logging.getLogger(__name__)
 
 
 @cache
-def load_enhancer(run_dir: str | Path | None, device):
+def load_enhancer(run_dir: Union[str, Path, None], device):
     run_dir = download(run_dir)
     hp = HParams.load(run_dir)
     enhancer = Enhancer(hp)

modules/repos_static/resemble_enhance/enhancer/lcfm/cfm.py

@@ -1,7 +1,7 @@
 import logging
 from dataclasses import dataclass
 from functools import partial
-from typing import Protocol
+from typing import Protocol, Union
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -17,8 +17,7 @@ logger = logging.getLogger(__name__)
 
 
 class VelocityField(Protocol):
-    def __call__(self, *, t: Tensor, ψt: Tensor, dt: Tensor) -> Tensor:
-        ...
+    def __call__(self, *, t: Tensor, ψt: Tensor, dt: Tensor) -> Tensor: ...
 
 
 class Solver:
@@ -40,7 +39,9 @@ class Solver:
 
         self._camera = None
         self._mel_fn = mel_fn
-        self._time_mapping = partial(self.exponential_decay_mapping, n=time_mapping_divisor)
+        self._time_mapping = partial(
+            self.exponential_decay_mapping, n=time_mapping_divisor
+        )
 
     def configurate_(self, nfe=None, method=None):
         if nfe is None:
@@ -50,7 +51,9 @@ class Solver:
             method = self.method
 
         if nfe == 1 and method in ("midpoint", "rk4"):
-            logger.warning(f"1 NFE is not supported for {method}, using euler method instead.")
+            logger.warning(
+                f"1 NFE is not supported for {method}, using euler method instead."
+            )
             method = "euler"
 
         self.nfe = nfe
@@ -105,7 +108,9 @@ class Solver:
                 )
             else:
                 # Spectrogram, b c t
-                plt.imshow(ψt.detach().cpu().numpy()[0], origin="lower", interpolation="none")
+                plt.imshow(
+                    ψt.detach().cpu().numpy()[0], origin="lower", interpolation="none"
+                )
             ax = plt.gca()
             ax.text(0.5, 1.01, f"t={t:.2f}", transform=ax.transAxes, ha="center")
             camera.snap()
@@ -271,7 +276,7 @@ class CFM(nn.Module):
             global_dim=self.time_emb_dim,
         )
 
-    def _perturb(self, ψ1: Tensor, t: Tensor | None = None):
+    def _perturb(self, ψ1: Tensor, t: Union[Tensor, None] = None):
         """
         Perturb ψ1 to ψt.
         """
@@ -311,7 +316,7 @@ class CFM(nn.Module):
         """
         return ψ1 - ψ0
 
-    def _to_v(self, *, ψt, x, t: float | Tensor):
+    def _to_v(self, *, ψt, x, t: Union[float, Tensor]):
         """
         Args:
             ψt: (b c t)
@@ -364,7 +369,13 @@ class CFM(nn.Module):
         ψ1 = self.solver(f=f, ψ0=ψ0, t0=t0)
         return ψ1
 
-    def forward(self, x: Tensor, y: Tensor | None = None, ψ0: Tensor | None = None, t0=0.0):
+    def forward(
+        self,
+        x: Tensor,
+        y: Union[Tensor, None] = None,
+        ψ0: Union[Tensor, None] = None,
+        t0=0.0,
+    ):
         if y is None:
             y = self.sample(x, ψ0=ψ0, t0=t0)
         else:

modules/repos_static/resemble_enhance/enhancer/lcfm/irmae.py

@@ -1,5 +1,6 @@
 import logging
 from dataclasses import dataclass
+from typing import Union
 
 import torch.nn as nn
 import torch.nn.functional as F
@@ -14,7 +15,7 @@ logger = logging.getLogger(__name__)
 @dataclass
 class IRMAEOutput:
     latent: Tensor  # latent vector
-    decoded: Tensor | None  # decoder output, include extra dim
+    decoded: Union[Tensor, None]  # decoder output, include extra dim
 
 
 class ResBlock(nn.Sequential):

modules/repos_static/resemble_enhance/enhancer/lcfm/lcfm.py

@@ -1,5 +1,6 @@
 import logging
 from enum import Enum
+from typing import Union
 
 import matplotlib.pyplot as plt
 import torch
@@ -70,19 +71,34 @@ class LCFM(nn.Module):
             return
 
         plt.subplot(221)
-        plt.imshow(y[0].detach().cpu().numpy(), aspect="auto", origin="lower", interpolation="none")
+        plt.imshow(
+            y[0].detach().cpu().numpy(),
+            aspect="auto",
+            origin="lower",
+            interpolation="none",
+        )
         plt.title("GT")
 
         plt.subplot(222)
         y_ = y_[:, : y.shape[1]]
-        plt.imshow(y_[0].detach().cpu().numpy(), aspect="auto", origin="lower", interpolation="none")
+        plt.imshow(
+            y_[0].detach().cpu().numpy(),
+            aspect="auto",
+            origin="lower",
+            interpolation="none",
+        )
         plt.title("Posterior")
 
         plt.subplot(223)
         z_ = self.cfm(x)
         y__ = self.ae.decode(z_)
         y__ = y__[:, : y.shape[1]]
-        plt.imshow(y__[0].detach().cpu().numpy(), aspect="auto", origin="lower", interpolation="none")
+        plt.imshow(
+            y__[0].detach().cpu().numpy(),
+            aspect="auto",
+            origin="lower",
+            interpolation="none",
+        )
         plt.title("C-Prior")
         del y__
 
@@ -90,7 +106,12 @@ class LCFM(nn.Module):
         z_ = torch.randn_like(z_)
         y__ = self.ae.decode(z_)
         y__ = y__[:, : y.shape[1]]
-        plt.imshow(y__[0].detach().cpu().numpy(), aspect="auto", origin="lower", interpolation="none")
+        plt.imshow(
+            y__[0].detach().cpu().numpy(),
+            aspect="auto",
+            origin="lower",
+            interpolation="none",
+        )
         plt.title("Prior")
         del z_, y__
 
@@ -109,7 +130,7 @@ class LCFM(nn.Module):
     def eval_tau_(self, tau):
         self._eval_tau = tau
 
-    def forward(self, x, y: Tensor | None = None, ψ0: Tensor | None = None):
+    def forward(self, x, y: Union[Tensor, None] = None, ψ0: Union[Tensor, None] = None):
         """
         Args:
             x: (b d t), condition mel
@@ -139,14 +160,20 @@ class LCFM(nn.Module):
 
             h = self.ae.decode(z)
         else:
-            ae_output: IRMAEOutput = self.ae(y, skip_decoding=self.mode == self.Mode.CFM)
+            ae_output: IRMAEOutput = self.ae(
+                y, skip_decoding=self.mode == self.Mode.CFM
+            )
 
             if self.mode == self.Mode.CFM:
                 _ = self.cfm(x, self._scale(ae_output.latent.detach()), ψ0=ψ0)
 
             h = ae_output.decoded
 
-            if h is not None and self.global_step is not None and self.global_step % 100 == 0:
+            if (
+                h is not None
+                and self.global_step is not None
+                and self.global_step % 100 == 0
+            ):
                 self._visualize(x[:1], y[:1], h[:1])
 
         return h

modules/repos_static/resemble_enhance/enhancer/univnet/univnet.py

@@ -1,3 +1,4 @@
+from typing import Union
 import numpy as np
 import torch
 import torch.nn.functional as F
@@ -50,7 +51,9 @@ class UnivNet(nn.Module):
             ]
         )
 
-        self.conv_pre = weight_norm(nn.Conv1d(self.d_noise, self.nc, 7, padding=3, padding_mode="reflect"))
+        self.conv_pre = weight_norm(
+            nn.Conv1d(self.d_noise, self.nc, 7, padding=3, padding_mode="reflect")
+        )
 
         self.conv_post = nn.Sequential(
             nn.LeakyReLU(0.2),
@@ -64,7 +67,7 @@ class UnivNet(nn.Module):
     def eps(self):
         return 1e-5
 
-    def forward(self, x: Tensor, y: Tensor | None = None, npad=10):
+    def forward(self, x: Tensor, y: Union[Tensor, None] = None, npad=10):
         """
         Args:
             x: (b c t), acoustic features
@@ -74,7 +77,9 @@ class UnivNet(nn.Module):
         """
         assert x.ndim == 3, "x must be 3D tensor"
         assert y is None or y.ndim == 2, "y must be 2D tensor"
-        assert x.shape[1] == self.d_input, f"x.shape[1] must be {self.d_input}, but got {x.shape}"
+        assert (
+            x.shape[1] == self.d_input
+        ), f"x.shape[1] must be {self.d_input}, but got {x.shape}"
         assert npad >= 0, "npad must be positive or zero"
 
         x = F.pad(x, (0, npad), "constant", 0)

modules/repos_static/resemble_enhance/hparams.py

@@ -1,6 +1,7 @@
 import logging
 from dataclasses import asdict, dataclass
 from pathlib import Path
+from typing import Union
 
 from omegaconf import OmegaConf
 from rich.console import Console
@@ -102,7 +103,7 @@ class HParams:
         OmegaConf.save(asdict(self), str(path))
 
     @classmethod
-    def load(cls, run_dir, yaml: Path | None = None):
+    def load(cls, run_dir, yaml: Union[Path, None] = None):
         hps = []
 
         if (run_dir / "hparams.yaml").exists():
@@ -120,7 +121,9 @@ class HParams:
                 for k, v in asdict(hp).items():
                     if getattr(hps[0], k) != v:
                         errors[k] = f"{getattr(hps[0], k)} != {v}"
-                raise ValueError(f"Found inconsistent hparams: {errors}, consider deleting {run_dir}")
+                raise ValueError(
+                    f"Found inconsistent hparams: {errors}, consider deleting {run_dir}"
+                )
 
         return hps[0]
 

modules/speaker.py

@@ -29,13 +29,15 @@ class Speaker:
         speaker.emb = tensor
         return speaker
 
-    def __init__(self, seed, name="", gender="", describe=""):
+    def __init__(
+        self, seed_or_tensor: Union[int, torch.Tensor], name="", gender="", describe=""
+    ):
         self.id = uuid.uuid4()
-        self.seed = seed
+        self.seed = -2 if isinstance(seed_or_tensor, torch.Tensor) else seed_or_tensor
         self.name = name
         self.gender = gender
         self.describe = describe
-        self.emb = None
+        self.emb = None if isinstance(seed_or_tensor, int) else seed_or_tensor
 
         # TODO replace emb => tokens
         self.tokens = []

modules/ssml_parser/SSMLParser.py

@@ -11,8 +11,8 @@ import copy
 
 
 class SSMLContext(Box):
-    def __init__(self, parent=None):
-        self.parent: Union[SSMLContext, None] = parent
+    def __init__(self, *args, **kwargs):
+        self.parent: Union[SSMLContext, None] = None
 
         self.style = None
         self.spk = None
@@ -29,18 +29,14 @@ class SSMLContext(Box):
         self.prompt2 = None
         self.prefix = None
 
-    def clone(self):
-        ctx = SSMLContext()
-        for k, v in self.items():
-            ctx[k] = v
-        return ctx
+        super().__init__(*args, **kwargs)
 
 
 class SSMLSegment(Box):
-    def __init__(self, text: str, attrs=SSMLContext()):
-        self.attrs = attrs
+    def __init__(self, text: str, attrs=SSMLContext(), params=None):
+        self.attrs = SSMLContext(**attrs)
         self.text = text
-        self.params = None
+        self.params = params
 
 
 class SSMLBreak:
@@ -68,7 +64,7 @@ class SSMLParser:
         root = etree.fromstring(ssml)
 
         root_ctx = SSMLContext()
-        segments = []
+        segments: List[Union[SSMLSegment, SSMLBreak]] = []
         self.resolve(root, root_ctx, segments)
 
         return segments
@@ -89,8 +85,13 @@ def create_ssml_parser():
     parser = SSMLParser()
 
     @parser.resolver("speak")
-    def tag_speak(element, context, segments, parser):
-        ctx = context.clone() if context is not None else SSMLContext()
+    def tag_speak(
+        element: etree.Element,
+        context: Box,
+        segments: List[Union[SSMLSegment, SSMLBreak]],
+        parser: SSMLParser,
+    ):
+        ctx = context.copy() if context is not None else SSMLContext()
 
         version = element.get("version")
         if version != "0.1":
@@ -100,8 +101,13 @@ def create_ssml_parser():
             parser.resolve(child, ctx, segments)
 
     @parser.resolver("voice")
-    def tag_voice(element, context, segments, parser):
-        ctx = context.clone() if context is not None else SSMLContext()
+    def tag_voice(
+        element: etree.Element,
+        context: Box,
+        segments: List[Union[SSMLSegment, SSMLBreak]],
+        parser: SSMLParser,
+    ):
+        ctx = context.copy() if context is not None else SSMLContext()
 
         ctx.spk = element.get("spk", ctx.spk)
         ctx.style = element.get("style", ctx.style)
@@ -131,13 +137,23 @@ def create_ssml_parser():
                 segments.append(SSMLSegment(child.tail.strip(), ctx))
 
     @parser.resolver("break")
-    def tag_break(element, context, segments, parser):
+    def tag_break(
+        element: etree.Element,
+        context: Box,
+        segments: List[Union[SSMLSegment, SSMLBreak]],
+        parser: SSMLParser,
+    ):
         time_ms = int(element.get("time", "0").replace("ms", ""))
         segments.append(SSMLBreak(time_ms))
 
     @parser.resolver("prosody")
-    def tag_prosody(element, context, segments, parser):
-        ctx = context.clone() if context is not None else SSMLContext()
+    def tag_prosody(
+        element: etree.Element,
+        context: Box,
+        segments: List[Union[SSMLSegment, SSMLBreak]],
+        parser: SSMLParser,
+    ):
+        ctx = context.copy() if context is not None else SSMLContext()
 
         ctx.spk = element.get("spk", ctx.spk)
         ctx.style = element.get("style", ctx.style)

modules/synthesize_audio.py

@@ -7,6 +7,7 @@ from modules import generate_audio as generate
 
 
 from modules.speaker import Speaker
+from modules.ssml_parser.SSMLParser import SSMLSegment
 from modules.utils import audio
 
 
@@ -23,45 +24,33 @@ def synthesize_audio(
     prefix: str = "",
     batch_size: int = 1,
     spliter_threshold: int = 100,
+    end_of_sentence="",
 ):
-    if batch_size == 1:
-        return generate.generate_audio(
-            text,
-            temperature=temperature,
-            top_P=top_P,
-            top_K=top_K,
-            spk=spk,
-            infer_seed=infer_seed,
-            use_decoder=use_decoder,
-            prompt1=prompt1,
-            prompt2=prompt2,
-            prefix=prefix,
+    spliter = SentenceSplitter(spliter_threshold)
+    sentences = spliter.parse(text)
+
+    text_segments = [
+        SSMLSegment(
+            text=s,
+            params={
+                "temperature": temperature,
+                "top_P": top_P,
+                "top_K": top_K,
+                "spk": spk,
+                "infer_seed": infer_seed,
+                "use_decoder": use_decoder,
+                "prompt1": prompt1,
+                "prompt2": prompt2,
+                "prefix": prefix,
+            },
         )
-    else:
-        spliter = SentenceSplitter(spliter_threshold)
-        sentences = spliter.parse(text)
+        for s in sentences
+    ]
+    synthesizer = SynthesizeSegments(
+        batch_size=batch_size, eos=end_of_sentence, spliter_thr=spliter_threshold
+    )
+    audio_segments = synthesizer.synthesize_segments(text_segments)
 
-        text_segments = [
-            {
-                "text": s,
-                "params": {
-                    "text": s,
-                    "temperature": temperature,
-                    "top_P": top_P,
-                    "top_K": top_K,
-                    "spk": spk,
-                    "infer_seed": infer_seed,
-                    "use_decoder": use_decoder,
-                    "prompt1": prompt1,
-                    "prompt2": prompt2,
-                    "prefix": prefix,
-                },
-            }
-            for s in sentences
-        ]
-        synthesizer = SynthesizeSegments(batch_size)
-        audio_segments = synthesizer.synthesize_segments(text_segments)
+    combined_audio = combine_audio_segments(audio_segments)
 
-        combined_audio = combine_audio_segments(audio_segments)
-
-        return audio.pydub_to_np(combined_audio)
+    return audio.pydub_to_np(combined_audio)

modules/utils/audio.py

@@ -95,7 +95,11 @@ def pitch_shift(
 
 
 def apply_prosody_to_audio_data(
-    audio_data: np.ndarray, rate: float, volume: float, pitch: float, sr: int
+    audio_data: np.ndarray,
+    rate: float = 1,
+    volume: float = 0,
+    pitch: float = 0,
+    sr: int = 24000,
 ) -> np.ndarray:
     if rate != 1:
         audio_data = pyrb.time_stretch(audio_data, sr=sr, rate=rate)

modules/webui/app.py

@@ -7,6 +7,7 @@ from modules import config
 from modules.webui import gradio_extensions, webui_config
 
 from modules.webui.changelog_tab import create_changelog_tab
+from modules.webui.finetune.ft_tab import create_ft_tabs
 from modules.webui.localization_runtime import ENLocalizationVars, ZHLocalizationVars
 from modules.webui.ssml.podcast_tab import create_ssml_podcast_tab
 from modules.webui.system_tab import create_system_tab
@@ -118,6 +119,8 @@ def create_interface():
                 gr.Markdown("🚧 Under construction")
             with gr.TabItem("ASR", visible=webui_config.experimental):
                 gr.Markdown("🚧 Under construction")
+            with gr.TabItem("Finetune", visible=webui_config.experimental):
+                create_ft_tabs(demo)
 
             with gr.TabItem("System"):
                 create_system_tab()

modules/webui/finetune/ProcessMonitor.py

@@ -0,0 +1,92 @@
+import os
+import sys
+import subprocess
+import threading
+
+
+class ProcessMonitor:
+    def __init__(self):
+        self.process = None
+        self.stdout = ""
+        self.stderr = ""
+        self.lock = threading.Lock()
+
+    def start_process(self, command):
+        self.process = subprocess.Popen(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            bufsize=1,
+            universal_newlines=True,
+        )
+
+        # Set pipes to non-blocking mode
+        fd_out = self.process.stdout.fileno()
+        fd_err = self.process.stderr.fileno()
+
+        if sys.platform != "win32":
+            import fcntl
+
+            fl_out = fcntl.fcntl(fd_out, fcntl.F_GETFL)
+            fl_err = fcntl.fcntl(fd_err, fcntl.F_GETFL)
+            fcntl.fcntl(fd_out, fcntl.F_SETFL, fl_out | os.O_NONBLOCK)
+            fcntl.fcntl(fd_err, fcntl.F_SETFL, fl_err | os.O_NONBLOCK)
+
+        # Start threads to read stdout and stderr
+        threading.Thread(target=self._read_stdout).start()
+        threading.Thread(target=self._read_stderr).start()
+
+    def _read_stdout(self):
+        while self.process is not None and self.process.poll() is None:
+            try:
+                output = self.process.stdout.read()
+                if output:
+                    with self.lock:
+                        self.stdout += output
+            except:
+                pass
+
+    def _read_stderr(self):
+        while self.process is not None and self.process.poll() is None:
+            try:
+                error = self.process.stderr.read()
+                if error:
+                    with self.lock:
+                        self.stderr += error
+            except:
+                pass
+
+    def get_output(self):
+        with self.lock:
+            return self.stdout, self.stderr
+
+    def stop_process(self):
+        if self.process:
+            self.process.terminate()
+            self.process = None
+
+
+if __name__ == "__main__":
+    import time
+
+    pm = ProcessMonitor()
+    pm.start_process(
+        [
+            "python",
+            "-u",
+            "-c",
+            "import time; [print(i) or time.sleep(1) for i in range(5)]",
+        ]
+    )
+
+    while pm.process and pm.process.poll() is None:
+        stdout, stderr = pm.get_output()
+        if stdout:
+            print("STDOUT:", stdout)
+        if stderr:
+            print("STDERR:", stderr)
+        time.sleep(1)
+
+    stdout, stderr = pm.get_output()
+    print("Final STDOUT:", stdout)
+    print("Final STDERR:", stderr)

modules/webui/finetune/ft_tab.py

@@ -0,0 +1,13 @@
+import gradio as gr
+
+from modules.webui.finetune.speaker_ft_tab import create_speaker_ft_tab
+
+
+def create_ft_tabs(demo):
+    with gr.Tabs():
+        with gr.TabItem("Speaker"):
+            create_speaker_ft_tab(demo)
+        with gr.TabItem("GPT"):
+            gr.Markdown("🚧 Under construction")
+        with gr.TabItem("AE"):
+            gr.Markdown("🚧 Under construction")

modules/webui/finetune/ft_ui_utils.py

@@ -0,0 +1,49 @@
+import os
+from typing import IO, Union
+from modules.speaker import Speaker, speaker_mgr
+import subprocess
+
+
+def get_datasets_dir():
+    """
+    列出 ./datasets/data_* 文件夹
+    """
+    dataset_path = "./datasets"
+    dataset_list = os.listdir(dataset_path)
+    dataset_list = [
+        d for d in dataset_list if os.path.isdir(os.path.join(dataset_path, d))
+    ]
+    dataset_list = [d for d in dataset_list if d.startswith("data_")]
+    return dataset_list
+
+
+def get_datasets_listfile():
+    datasets = get_datasets_dir()
+    listfiles = []
+    for d in datasets:
+        dir_path = os.path.join("./datasets", d)
+        files = os.listdir(dir_path)
+        for f in files:
+            if f.endswith(".list"):
+                listfiles.append(os.path.join(dir_path, f))
+    return listfiles
+
+
+def run_speaker_ft(
+    batch_size: int, epochs: int, train_text: bool, data_path: str, init_speaker: str
+):
+    command = ["python3", "-m", "modules.finetune.train_speaker"]
+    command += [
+        f"--batch_size={batch_size}",
+        f"--epochs={epochs}",
+        f"--data_path={data_path}",
+    ]
+    if train_text:
+        command.append("--train_text")
+    if init_speaker:
+        command.append(f"--init_speaker={init_speaker}")
+    process = subprocess.Popen(
+        command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1
+    )
+
+    return process

modules/webui/finetune/speaker_ft_tab.py

@@ -0,0 +1,130 @@
+import gradio as gr
+
+from modules.Enhancer.ResembleEnhance import unload_enhancer
+from modules.webui import webui_config
+from modules.webui.webui_utils import get_speaker_names
+from .ft_ui_utils import get_datasets_listfile, run_speaker_ft
+from .ProcessMonitor import ProcessMonitor
+from modules.speaker import speaker_mgr
+from modules.models import unload_chat_tts
+
+
+class SpeakerFt:
+    def __init__(self):
+        self.process_monitor = ProcessMonitor()
+        self.status_str = "idle"
+
+    def unload_main_thread_models(self):
+        unload_chat_tts()
+        unload_enhancer()
+
+    def run(
+        self,
+        batch_size: int,
+        epochs: int,
+        lr: str,
+        train_text: bool,
+        data_path: str,
+        select_speaker: str = "",
+    ):
+        if self.process_monitor.process:
+            return
+        self.unload_main_thread_models()
+        spk_path = None
+        if select_speaker != "" and select_speaker != "none":
+            select_speaker = select_speaker.split(" : ")[1].strip()
+            spk = speaker_mgr.get_speaker(select_speaker)
+            if spk is None:
+                return ["Speaker not found"]
+            spk_filename = speaker_mgr.get_speaker_filename(spk.id)
+            spk_path = f"./data/speakers/{spk_filename}"
+
+        command = ["python3", "-m", "modules.finetune.train_speaker"]
+        command += [
+            f"--batch_size={batch_size}",
+            f"--epochs={epochs}",
+            f"--data_path={data_path}",
+        ]
+        if train_text:
+            command.append("--train_text")
+        if spk_path:
+            command.append(f"--init_speaker={spk_path}")
+
+        self.status("Training process starting")
+
+        self.process_monitor.start_process(command)
+
+        self.status("Training started")
+
+    def status(self, text: str):
+        self.status_str = text
+
+    def flush(self):
+        stdout, stderr = self.process_monitor.get_output()
+        return f"{self.status_str}\n{stdout}\n{stderr}"
+
+    def clear(self):
+        self.process_monitor.stdout = ""
+        self.process_monitor.stderr = ""
+        self.status("Logs cleared")
+
+    def stop(self):
+        self.process_monitor.stop_process()
+        self.status("Training stopped")
+
+
+def create_speaker_ft_tab(demo: gr.Blocks):
+    spk_ft = SpeakerFt()
+    speakers, speaker_names = get_speaker_names()
+    speaker_names = ["none"] + speaker_names
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            with gr.Group():
+                gr.Markdown("🎛️hparams")
+                dataset_input = gr.Dropdown(
+                    label="Dataset", choices=get_datasets_listfile()
+                )
+                lr_input = gr.Textbox(label="Learning Rate", value="1e-2")
+                epochs_input = gr.Slider(
+                    label="Epochs", value=10, minimum=1, maximum=100, step=1
+                )
+                batch_size_input = gr.Slider(
+                    label="Batch Size", value=4, minimum=1, maximum=64, step=1
+                )
+                train_text_checkbox = gr.Checkbox(label="Train text_loss", value=True)
+                init_spk_dropdown = gr.Dropdown(
+                    label="Initial Speaker",
+                    choices=speaker_names,
+                    value="none",
+                )
+
+            with gr.Group():
+                start_train_btn = gr.Button("Start Training")
+                stop_train_btn = gr.Button("Stop Training")
+                clear_train_btn = gr.Button("Clear logs")
+        with gr.Column(scale=5):
+            with gr.Group():
+                # log
+                gr.Markdown("📜logs")
+                log_output = gr.Textbox(
+                    show_label=False, label="Log", value="", lines=20, interactive=True
+                )
+
+    start_train_btn.click(
+        spk_ft.run,
+        inputs=[
+            batch_size_input,
+            epochs_input,
+            lr_input,
+            train_text_checkbox,
+            dataset_input,
+            init_spk_dropdown,
+        ],
+        outputs=[],
+    )
+    stop_train_btn.click(spk_ft.stop)
+    clear_train_btn.click(spk_ft.clear)
+
+    if webui_config.experimental:
+        demo.load(spk_ft.flush, every=1, outputs=[log_output])

modules/webui/localization_runtime.py

@@ -7,6 +7,7 @@ class LocalizationVars:
 
         self.ssml_examples = []
         self.tts_examples = []
+        self.podcast_default = []
 
 
 class ZHLocalizationVars(LocalizationVars):
@@ -167,6 +168,69 @@ class ZHLocalizationVars(LocalizationVars):
             },
         ]
 
+        self.podcast_default = [
+            [
+                1,
+                "female2",
+                "你好，欢迎收听今天的播客内容。今天我们要聊的是中华料理。",
+                "podcast",
+            ],
+            [
+                2,
+                "Alice",
+                "嗨，我特别期待这个话题！中华料理真的是博大精深。",
+                "podcast",
+            ],
+            [
+                3,
+                "Bob",
+                "没错，中华料理有着几千年的历史，而且每个地区都有自己的特色菜。",
+                "podcast",
+            ],
+            [
+                4,
+                "female2",
+                "那我们先从最有名的川菜开始吧。川菜以其麻辣著称，是很多人的最爱。",
+                "podcast",
+            ],
+            [
+                5,
+                "Alice",
+                "对，我特别喜欢吃麻婆豆腐和辣子鸡。那种麻辣的感觉真是让人难以忘怀。",
+                "podcast",
+            ],
+            [
+                6,
+                "Bob",
+                "除了川菜，粤菜也是很受欢迎的。粤菜讲究鲜美，像是白切鸡和蒸鱼都是经典。",
+                "podcast",
+            ],
+            [
+                7,
+                "female2",
+                "对啊，粤菜的烹饪方式比较清淡，更注重食材本身的味道。",
+                "podcast",
+            ],
+            [
+                8,
+                "Alice",
+                "还有北京的京菜，像北京烤鸭，那可是来北京必吃的美食。",
+                "podcast",
+            ],
+            [
+                9,
+                "Bob",
+                "不仅如此，还有淮扬菜、湘菜、鲁菜等等，每个菜系都有其独特的风味。",
+                "podcast",
+            ],
+            [
+                10,
+                "female2",
+                "对对对，像淮扬菜的狮子头，湘菜的剁椒鱼头，都是让人垂涎三尺的美味。",
+                "podcast",
+            ],
+        ]
+
 
 class ENLocalizationVars(LocalizationVars):
     def __init__(self):
@@ -224,3 +288,65 @@ class ENLocalizationVars(LocalizationVars):
                 "text": "Don't ever let somebody tell you you can't do something. Not even me. Alright? You got a dream, you gotta protect it. When people can't do something themselves, they're gonna tell you that you can't do it. You want something, go get it. Period.",
             },
         ]
+        self.podcast_default = [
+            [
+                1,
+                "female2",
+                "Hello, welcome to today's podcast. Today, we're going to talk about global cuisine.",
+                "podcast",
+            ],
+            [
+                2,
+                "Alice",
+                "Hi, I'm really excited about this topic! Global cuisine is incredibly diverse and fascinating.",
+                "podcast",
+            ],
+            [
+                3,
+                "Bob",
+                "Absolutely, every country has its own unique culinary traditions and specialties.",
+                "podcast",
+            ],
+            [
+                4,
+                "female2",
+                "Let's start with Italian cuisine. Italian food is loved worldwide, especially for its pasta and pizza.",
+                "podcast",
+            ],
+            [
+                5,
+                "Alice",
+                "Yes, I especially love a good Margherita pizza and a hearty plate of spaghetti carbonara. The flavors are simply amazing.",
+                "podcast",
+            ],
+            [
+                6,
+                "Bob",
+                "Besides Italian cuisine, Japanese cuisine is also very popular. Dishes like sushi and ramen have become global favorites.",
+                "podcast",
+            ],
+            [
+                7,
+                "female2",
+                "Exactly, Japanese cuisine is known for its emphasis on fresh ingredients and delicate presentation.",
+                "podcast",
+            ],
+            [
+                8,
+                "Alice",
+                "And then there's Mexican cuisine, with its bold flavors and colorful dishes like tacos and guacamole.",
+                "podcast",
+            ],
+            [
+                9,
+                "Bob",
+                "Not to mention, there's also Indian cuisine, Thai cuisine, French cuisine, and so many more, each with its own distinctive flavors and techniques.",
+                "podcast",
+            ],
+            [
+                10,
+                "female2",
+                "Yes, like Indian curry, Thai tom yum soup, and French croissants, these are all mouth-watering dishes that are loved by people all over the world.",
+                "podcast",
+            ],
+        ]

modules/webui/ssml/podcast_tab.py

@@ -3,72 +3,9 @@ import pandas as pd
 import torch
 
 from modules.normalization import text_normalize
-from modules.webui import webui_utils
+from modules.webui import webui_config, webui_utils
 from modules.utils.hf import spaces
 
-podcast_default_case = [
-    [
-        1,
-        "female2",
-        "你好，欢迎收听今天的播客内容。今天我们要聊的是中华料理。 [lbreak]",
-        "podcast",
-    ],
-    [
-        2,
-        "Alice",
-        "嗨，我特别期待这个话题！中华料理真的是博大精深。 [lbreak]",
-        "podcast",
-    ],
-    [
-        3,
-        "Bob",
-        "没错，中华料理有着几千年的历史，而且每个地区都有自己的特色菜。 [lbreak]",
-        "podcast",
-    ],
-    [
-        4,
-        "female2",
-        "那我们先从最有名的川菜开始吧。川菜以其麻辣著称，是很多人的最爱。 [lbreak]",
-        "podcast",
-    ],
-    [
-        5,
-        "Alice",
-        "对，我特别喜欢吃麻婆豆腐和辣子鸡。那种麻辣的感觉真是让人难以忘怀。 [lbreak]",
-        "podcast",
-    ],
-    [
-        6,
-        "Bob",
-        "除了川菜，粤菜也是很受欢迎的。粤菜讲究鲜美，像是白切鸡和蒸鱼都是经典。 [lbreak]",
-        "podcast",
-    ],
-    [
-        7,
-        "female2",
-        "对啊，粤菜的烹饪方式比较清淡，更注重食材本身的味道。 [lbreak]",
-        "podcast",
-    ],
-    [
-        8,
-        "Alice",
-        "还有北京的京菜，像北京烤鸭，那可是来北京必吃的美食。 [lbreak]",
-        "podcast",
-    ],
-    [
-        9,
-        "Bob",
-        "不仅如此，还有淮扬菜、湘菜、鲁菜等等，每个菜系都有其独特的风味。 [lbreak]",
-        "podcast",
-    ],
-    [
-        10,
-        "female2",
-        "对对对，像淮扬菜的狮子头，湘菜的剁椒鱼头，都是让人垂涎三尺的美味。 [lbreak]",
-        "podcast",
-    ],
-]
-
 
 # NOTE: 因为 text_normalize 需要使用 tokenizer
 @torch.inference_mode()
@@ -133,7 +70,7 @@ def create_ssml_podcast_tab(ssml_input: gr.Textbox, tabs1: gr.Tabs, tabs2: gr.Ta
                     datatype=["number", "str", "str", "str"],
                     interactive=True,
                     wrap=True,
-                    value=podcast_default_case,
+                    value=webui_config.localization.podcast_default,
                     row_count=(0, "dynamic"),
                     col_count=(4, "fixed"),
                 )

modules/webui/ssml/ssml_tab.py

@@ -22,7 +22,6 @@ def create_ssml_interface():
                 ssml_button = gr.Button("🔊Synthesize SSML", variant="primary")
         with gr.Column(scale=1):
             with gr.Group():
-                # 参数
                 gr.Markdown("🎛️Parameters")
                 # batch size
                 batch_size_input = gr.Slider(
@@ -32,6 +31,19 @@ def create_ssml_interface():
                     maximum=webui_config.max_batch_size,
                     step=1,
                 )
+            with gr.Group():
+                gr.Markdown("🎛️Spliter")
+                eos_input = gr.Textbox(
+                    label="eos",
+                    value="[uv_break]",
+                )
+                spliter_thr_input = gr.Slider(
+                    label="Spliter Threshold",
+                    value=100,
+                    minimum=50,
+                    maximum=1000,
+                    step=1,
+                )
 
             with gr.Group():
                 gr.Markdown("💪🏼Enhance")
@@ -49,7 +61,14 @@ def create_ssml_interface():
 
     ssml_button.click(
         synthesize_ssml,
-        inputs=[ssml_input, batch_size_input, enable_enhance, enable_de_noise],
+        inputs=[
+            ssml_input,
+            batch_size_input,
+            enable_enhance,
+            enable_de_noise,
+            eos_input,
+            spliter_thr_input,
+        ],
         outputs=ssml_output,
     )