update

modules/Denoiser/AudioDenoiser.py

@@ -0,0 +1,140 @@
+import logging
+import math
+from typing import Union
+import torch
+import torchaudio
+from torch import nn
+from audio_denoiser.helpers.torch_helper import batched_apply
+from modules.Denoiser.AudioNosiseModel import load_audio_denosier_model
+from audio_denoiser.helpers.audio_helper import (
+    create_spectrogram,
+    reconstruct_from_spectrogram,
+)
+
+_expected_t_std = 0.23
+_recommended_backend = "soundfile"
+
+
+# ref: https://github.com/jose-solorzano/audio-denoiser
+class AudioDenoiser:
+    def __init__(
+        self,
+        local_dir: str,
+        device: Union[str, torch.device] = None,
+        num_iterations: int = 100,
+    ):
+        super().__init__()
+        if device is None:
+            is_cuda = torch.cuda.is_available()
+            if not is_cuda:
+                logging.warning("CUDA not available. Will use CPU.")
+            device = torch.device("cuda:0") if is_cuda else torch.device("cpu")
+        self.device = device
+        self.model = load_audio_denosier_model(dir_path=local_dir, device=device)
+        self.model.eval()
+        self.model_sample_rate = self.model.sample_rate
+        self.scaler = self.model.scaler
+        self.n_fft = self.model.n_fft
+        self.segment_num_frames = self.model.num_frames
+        self.num_iterations = num_iterations
+
+    @staticmethod
+    def _sp_log(spectrogram: torch.Tensor, eps=0.01):
+        return torch.log(spectrogram + eps)
+
+    @staticmethod
+    def _sp_exp(log_spectrogram: torch.Tensor, eps=0.01):
+        return torch.clamp(torch.exp(log_spectrogram) - eps, min=0)
+
+    @staticmethod
+    def _trimmed_dev(waveform: torch.Tensor, q: float = 0.90) -> float:
+        # Expected for training data is ~0.23
+        abs_waveform = torch.abs(waveform)
+        quantile_value = torch.quantile(abs_waveform, q).item()
+        trimmed_values = waveform[abs_waveform >= quantile_value]
+        return torch.std(trimmed_values).item()
+
+    def process_waveform(
+        self,
+        waveform: torch.Tensor,
+        sample_rate: int,
+        return_cpu_tensor: bool = False,
+        auto_scale: bool = False,
+    ) -> torch.Tensor:
+        """
+        Denoises a waveform.
+        @param waveform: A waveform tensor. Use torchaudio structure.
+        @param sample_rate: The sample rate of the waveform in Hz.
+        @param return_cpu_tensor: Whether the returned tensor must be a CPU tensor.
+        @param auto_scale: Normalize the scale of the waveform before processing. Recommended for low-volume audio.
+        @return: A denoised waveform.
+        """
+        waveform = waveform.cpu()
+        if auto_scale:
+            w_t_std = self._trimmed_dev(waveform)
+            waveform = waveform * _expected_t_std / w_t_std
+        if sample_rate != self.model_sample_rate:
+            transform = torchaudio.transforms.Resample(
+                orig_freq=sample_rate, new_freq=self.model_sample_rate
+            )
+            waveform = transform(waveform)
+        hop_len = self.n_fft // 2
+        spectrogram = create_spectrogram(waveform, n_fft=self.n_fft, hop_length=hop_len)
+        spectrogram = spectrogram.to(self.device)
+        num_a_channels = spectrogram.size(0)
+        with torch.no_grad():
+            results = []
+            for c in range(num_a_channels):
+                c_spectrogram = spectrogram[c]
+                # c_spectrogram: (257, num_frames)
+                fft_size, num_frames = c_spectrogram.shape
+                num_segments = math.ceil(num_frames / self.segment_num_frames)
+                adj_num_frames = num_segments * self.segment_num_frames
+                if adj_num_frames > num_frames:
+                    c_spectrogram = nn.functional.pad(
+                        c_spectrogram, (0, adj_num_frames - num_frames)
+                    )
+                c_spectrogram = c_spectrogram.view(
+                    fft_size, num_segments, self.segment_num_frames
+                )
+                # c_spectrogram: (257, num_segments, 32)
+                c_spectrogram = torch.permute(c_spectrogram, (1, 0, 2))
+                # c_spectrogram: (num_segments, 257, 32)
+                log_c_spectrogram = self._sp_log(c_spectrogram)
+                scaled_log_c_sp = self.scaler(log_c_spectrogram)
+                pred_noise_log_sp = batched_apply(
+                    self.model, scaled_log_c_sp, detached=True
+                )
+                log_denoised_sp = log_c_spectrogram - pred_noise_log_sp
+                denoised_sp = self._sp_exp(log_denoised_sp)
+                # denoised_sp: (num_segments, 257, 32)
+                denoised_sp = torch.permute(denoised_sp, (1, 0, 2))
+                # denoised_sp: (257, num_segments, 32)
+                denoised_sp = denoised_sp.contiguous().view(1, fft_size, adj_num_frames)
+                # denoised_sp: (1, 257, adj_num_frames)
+                denoised_sp = denoised_sp[:, :, :num_frames]
+                denoised_sp = denoised_sp.cpu()
+                denoised_waveform = reconstruct_from_spectrogram(
+                    denoised_sp, num_iterations=self.num_iterations
+                )
+                # denoised_waveform: (1, num_samples)
+                results.append(denoised_waveform)
+            cpu_results = torch.cat(results)
+            return cpu_results if return_cpu_tensor else cpu_results.to(self.device)
+
+    def process_audio_file(
+        self, in_audio_file: str, out_audio_file: str, auto_scale: bool = False
+    ):
+        """
+        Denoises an audio file.
+        @param in_audio_file: An input audio file with a format supported by torchaudio.
+        @param out_audio_file: Am output audio file with a format supported by torchaudio.
+        @param auto_scale: Whether the input waveform scale should be normalized before processing. Recommended for low-volume audio.
+        """
+        waveform, sample_rate = torchaudio.load(in_audio_file)
+        denoised_waveform = self.process_waveform(
+            waveform, sample_rate, return_cpu_tensor=True, auto_scale=auto_scale
+        )
+        torchaudio.save(
+            out_audio_file, denoised_waveform, sample_rate=self.model_sample_rate
+        )

modules/Denoiser/AudioNosiseModel.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn as nn
+
+from audio_denoiser.modules.Permute import Permute
+from audio_denoiser.modules.SimpleRoberta import SimpleRoberta
+from audio_denoiser.modules.SpectrogramScaler import SpectrogramScaler
+
+import json
+
+
+class AudioNoiseModel(nn.Module):
+    def __init__(self, config: dict):
+        super(AudioNoiseModel, self).__init__()
+
+        # Encoder layers
+        self.config = config
+        scaler_dict = config["scaler"]
+        self.scaler = SpectrogramScaler.from_dict(scaler_dict)
+        self.in_channels = config.get("in_channels", 257)
+        self.roberta_hidden_size = config.get("roberta_hidden_size", 768)
+        self.model1 = nn.Sequential(
+            nn.Conv1d(self.in_channels, 1024, kernel_size=1),
+            nn.ELU(),
+            nn.Conv1d(1024, 1024, kernel_size=1),
+            nn.ELU(),
+            nn.Conv1d(1024, self.in_channels, kernel_size=1),
+        )
+        self.model2 = nn.Sequential(
+            Permute(0, 2, 1),
+            nn.Linear(self.in_channels, self.roberta_hidden_size),
+            SimpleRoberta(num_hidden_layers=5, hidden_size=self.roberta_hidden_size),
+            nn.Linear(self.roberta_hidden_size, self.in_channels),
+            Permute(0, 2, 1),
+        )
+
+    @property
+    def sample_rate(self) -> int:
+        return self.config.get("sample_rate", 16000)
+
+    @property
+    def n_fft(self) -> int:
+        return self.config.get("n_fft", 512)
+
+    @property
+    def num_frames(self) -> int:
+        return self.config.get("num_frames", 32)
+
+    def forward(self, x, use_scaler: bool = False, out_scale: float = 1.0):
+        if use_scaler:
+            x = self.scaler(x)
+        x1 = self.model1(x)
+        x2 = self.model2(x)
+        x = x1 + x2
+        return x * out_scale
+
+
+def load_audio_denosier_model(dir_path: str, device) -> AudioNoiseModel:
+    config = json.load(open(f"{dir_path}/config.json", "r"))
+    model = AudioNoiseModel(config)
+    model.load_state_dict(torch.load(f"{dir_path}/pytorch_model.bin"))
+
+    model.to(device)
+    model.model1.to(device)
+    model.model2.to(device)
+
+    return model

modules/Enhancer/ResembleEnhance.py

@@ -0,0 +1,116 @@
+import os
+from typing import List
+from resemble_enhance.enhancer.enhancer import Enhancer
+from resemble_enhance.enhancer.hparams import HParams
+from resemble_enhance.inference import inference
+
+import torch
+
+from modules.utils.constants import MODELS_DIR
+from pathlib import Path
+
+from threading import Lock
+
+resemble_enhance = None
+lock = Lock()
+
+
+def load_enhancer(device: torch.device):
+    global resemble_enhance
+    with lock:
+        if resemble_enhance is None:
+            resemble_enhance = ResembleEnhance(device)
+            resemble_enhance.load_model()
+    return resemble_enhance
+
+
+class ResembleEnhance:
+    hparams: HParams
+    enhancer: Enhancer
+
+    def __init__(self, device: torch.device):
+        self.device = device
+
+        self.enhancer = None
+        self.hparams = None
+
+    def load_model(self):
+        hparams = HParams.load(Path(MODELS_DIR) / "resemble-enhance")
+        enhancer = Enhancer(hparams)
+        state_dict = torch.load(
+            Path(MODELS_DIR) / "resemble-enhance" / "mp_rank_00_model_states.pt",
+            map_location="cpu",
+        )["module"]
+        enhancer.load_state_dict(state_dict)
+        enhancer.eval()
+        enhancer.to(self.device)
+        enhancer.denoiser.to(self.device)
+
+        self.hparams = hparams
+        self.enhancer = enhancer
+
+    @torch.inference_mode()
+    def denoise(self, dwav, sr, device) -> tuple[torch.Tensor, int]:
+        assert self.enhancer is not None, "Model not loaded"
+        assert self.enhancer.denoiser is not None, "Denoiser not loaded"
+        enhancer = self.enhancer
+        return inference(model=enhancer.denoiser, dwav=dwav, sr=sr, device=device)
+
+    @torch.inference_mode()
+    def enhance(
+        self,
+        dwav,
+        sr,
+        device,
+        nfe=32,
+        solver="midpoint",
+        lambd=0.5,
+        tau=0.5,
+    ) -> tuple[torch.Tensor, int]:
+        assert 0 < nfe <= 128, f"nfe must be in (0, 128], got {nfe}"
+        assert solver in (
+            "midpoint",
+            "rk4",
+            "euler",
+        ), f"solver must be in ('midpoint', 'rk4', 'euler'), got {solver}"
+        assert 0 <= lambd <= 1, f"lambd must be in [0, 1], got {lambd}"
+        assert 0 <= tau <= 1, f"tau must be in [0, 1], got {tau}"
+        assert self.enhancer is not None, "Model not loaded"
+        enhancer = self.enhancer
+        enhancer.configurate_(nfe=nfe, solver=solver, lambd=lambd, tau=tau)
+        return inference(model=enhancer, dwav=dwav, sr=sr, device=device)
+
+
+if __name__ == "__main__":
+    import torchaudio
+    from modules.models import load_chat_tts
+
+    load_chat_tts()
+
+    device = torch.device("cuda")
+    ench = ResembleEnhance(device)
+    ench.load_model()
+
+    wav, sr = torchaudio.load("test.wav")
+
+    print(wav.shape, type(wav), sr, type(sr))
+    exit()
+
+    wav = wav.squeeze(0).cuda()
+
+    print(wav.device)
+
+    denoised, d_sr = ench.denoise(wav.cpu(), sr, device)
+    denoised = denoised.unsqueeze(0)
+    print(denoised.shape)
+    torchaudio.save("denoised.wav", denoised, d_sr)
+
+    for solver in ("midpoint", "rk4", "euler"):
+        for lambd in (0.1, 0.5, 0.9):
+            for tau in (0.1, 0.5, 0.9):
+                enhanced, e_sr = ench.enhance(
+                    wav.cpu(), sr, device, solver=solver, lambd=lambd, tau=tau, nfe=128
+                )
+                enhanced = enhanced.unsqueeze(0)
+                print(enhanced.shape)
+                torchaudio.save(f"enhanced_{solver}_{lambd}_{tau}.wav", enhanced, e_sr)

modules/SynthesizeSegments.py

@@ -1,17 +1,18 @@
+from box import Box
 from pydub import AudioSegment
-from typing import Any, List, Dict, Union
+from typing import List, Union
 from scipy.io.wavfile import write
 import io
+from modules.api.utils import calc_spk_style
+from modules.ssml_parser.SSMLParser import SSMLSegment, SSMLBreak, SSMLContext
 from modules.utils import rng
 from modules.utils.audio import time_stretch, pitch_shift
 from modules import generate_audio
 from modules.normalization import text_normalize
 import logging
 import json
-import copy
-import numpy as np
 
-from modules.speaker import Speaker
+from modules.speaker import Speaker, speaker_mgr
 
 logger = logging.getLogger(__name__)
 
@@ -24,7 +25,7 @@ def audio_data_to_segment(audio_data, sr):
     return AudioSegment.from_file(byte_io, format="wav")
 
 
-def combine_audio_segments(audio_segments: list) -> AudioSegment:
+def combine_audio_segments(audio_segments: list[AudioSegment]) -> AudioSegment:
     combined_audio = AudioSegment.empty()
     for segment in audio_segments:
         combined_audio += segment
@@ -54,230 +55,191 @@ def to_number(value, t, default=0):
         return default
 
 
+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)
+
+
 class SynthesizeSegments:
     def __init__(self, batch_size: int = 8):
         self.batch_size = batch_size
         self.batch_default_spk_seed = rng.np_rng()
         self.batch_default_infer_seed = rng.np_rng()
 
-    def segment_to_generate_params(self, segment: Dict[str, Any]) -> Dict[str, Any]:
+    def segment_to_generate_params(
+        self, segment: Union[SSMLSegment, SSMLBreak]
+    ) -> TTSAudioSegment:
+        if isinstance(segment, SSMLBreak):
+            return TTSAudioSegment(_type="break")
+
         if segment.get("params", None) is not None:
-            return segment["params"]
+            return TTSAudioSegment(**segment.get("params"))
 
         text = segment.get("text", "")
         is_end = segment.get("is_end", False)
 
         text = str(text).strip()
 
-        attrs = segment.get("attrs", {})
-        spk = attrs.get("spk", "")
-        if isinstance(spk, str):
-            spk = int(spk)
-        seed = to_number(attrs.get("seed", ""), int, -1)
-        top_k = to_number(attrs.get("top_k", ""), int, None)
-        top_p = to_number(attrs.get("top_p", ""), float, None)
-        temp = to_number(attrs.get("temp", ""), float, None)
-
-        prompt1 = attrs.get("prompt1", "")
-        prompt2 = attrs.get("prompt2", "")
-        prefix = attrs.get("prefix", "")
+        attrs = segment.attrs
+        spk = attrs.spk
+        style = attrs.style
+
+        ss_params = calc_spk_style(spk, style)
+
+        if "spk" in ss_params:
+            spk = ss_params["spk"]
+
+        seed = to_number(attrs.seed, int, ss_params.get("seed") or -1)
+        top_k = to_number(attrs.top_k, int, None)
+        top_p = to_number(attrs.top_p, float, None)
+        temp = to_number(attrs.temp, float, None)
+
+        prompt1 = attrs.prompt1 or ss_params.get("prompt1")
+        prompt2 = attrs.prompt2 or ss_params.get("prompt2")
+        prefix = attrs.prefix or ss_params.get("prefix")
         disable_normalize = attrs.get("normalize", "") == "False"
 
-        params = {
-            "text": text,
-            "temperature": temp if temp is not None else 0.3,
-            "top_P": top_p if top_p is not None else 0.5,
-            "top_K": top_k if top_k is not None else 20,
-            "spk": spk if spk else -1,
-            "infer_seed": seed if seed else -1,
-            "prompt1": prompt1 if prompt1 else "",
-            "prompt2": prompt2 if prompt2 else "",
-            "prefix": prefix if prefix else "",
-        }
+        seg = TTSAudioSegment(
+            _type="voice",
+            text=text,
+            temperature=temp if temp is not None else 0.3,
+            top_P=top_p if top_p is not None else 0.5,
+            top_K=top_k if top_k is not None else 20,
+            spk=spk if spk else -1,
+            infer_seed=seed if seed else -1,
+            prompt1=prompt1 if prompt1 else "",
+            prompt2=prompt2 if prompt2 else "",
+            prefix=prefix if prefix else "",
+        )
 
         if not disable_normalize:
-            params["text"] = text_normalize(text, is_end=is_end)
+            seg.text = text_normalize(text, is_end=is_end)
+
+        # NOTE 每个batch的默认seed保证前后一致即使是没设置spk的情况
+        if seg.spk == -1:
+            seg.spk = self.batch_default_spk_seed
+        if seg.infer_seed == -1:
+            seg.infer_seed = self.batch_default_infer_seed
+
+        return seg
+
+    def process_break_segments(
+        self,
+        src_segments: List[SSMLBreak],
+        bucket_segments: List[SSMLBreak],
+        audio_segments: List[AudioSegment],
+    ):
+        for segment in bucket_segments:
+            index = src_segments.index(segment)
+            audio_segments[index] = AudioSegment.silent(
+                duration=int(segment.attrs.duration)
+            )
 
-        # Set default values for spk and infer_seed
-        if params["spk"] == -1:
-            params["spk"] = self.batch_default_spk_seed
-        if params["infer_seed"] == -1:
-            params["infer_seed"] = self.batch_default_infer_seed
+    def process_voice_segments(
+        self,
+        src_segments: List[SSMLSegment],
+        bucket: List[SSMLSegment],
+        audio_segments: List[AudioSegment],
+    ):
+        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]
+
+            params = param_arr[0]
+            audio_datas = generate_audio.generate_audio_batch(
+                texts=texts,
+                temperature=params.temperature,
+                top_P=params.top_P,
+                top_K=params.top_K,
+                spk=params.spk,
+                infer_seed=params.infer_seed,
+                prompt1=params.prompt1,
+                prompt2=params.prompt2,
+                prefix=params.prefix,
+            )
+            for idx, segment in enumerate(batch):
+                sr, audio_data = audio_datas[idx]
+                rate = float(segment.get("rate", "1.0"))
+                volume = float(segment.get("volume", "0"))
+                pitch = float(segment.get("pitch", "0"))
 
-        return params
+                audio_segment = audio_data_to_segment(audio_data, sr)
+                audio_segment = apply_prosody(audio_segment, rate, volume, pitch)
+                original_index = src_segments.index(segment)
+                audio_segments[original_index] = audio_segment
 
     def bucket_segments(
-        self, segments: List[Dict[str, Any]]
-    ) -> List[List[Dict[str, Any]]]:
-        # Create a dictionary to hold buckets
-        buckets = {}
+        self, segments: List[Union[SSMLSegment, SSMLBreak]]
+    ) -> List[List[Union[SSMLSegment, SSMLBreak]]]:
+        buckets = {"<break>": []}
         for segment in segments:
+            if isinstance(segment, SSMLBreak):
+                buckets["<break>"].append(segment)
+                continue
+
             params = self.segment_to_generate_params(segment)
 
-            key_params = copy.copy(params)
-            if isinstance(key_params.get("spk"), Speaker):
-                key_params["spk"] = str(key_params["spk"].id)
+            if isinstance(params.spk, Speaker):
+                params.spk = str(params.spk.id)
+
             key = json.dumps(
-                {k: v for k, v in key_params.items() if k != "text"}, sort_keys=True
+                {k: v for k, v in params.items() if k != "text"}, sort_keys=True
             )
             if key not in buckets:
                 buckets[key] = []
             buckets[key].append(segment)
 
-        # Convert dictionary to list of buckets
-        bucket_list = list(buckets.values())
-        return bucket_list
+        return buckets
 
-    def synthesize_segments(self, segments: List[Dict[str, Any]]) -> List[AudioSegment]:
-        audio_segments = [None] * len(
-            segments
-        )  # Create a list with the same length as segments
+    def synthesize_segments(
+        self, segments: List[Union[SSMLSegment, SSMLBreak]]
+    ) -> List[AudioSegment]:
+        audio_segments = [None] * len(segments)
         buckets = self.bucket_segments(segments)
-        logger.debug(f"segments len: {len(segments)}")
-        logger.debug(f"bucket pool size: {len(buckets)}")
-        for bucket in buckets:
-            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]
-
-                params = param_arr[0]  # Use the first segment to get the parameters
-                audio_datas = generate_audio.generate_audio_batch(
-                    texts=texts,
-                    temperature=params["temperature"],
-                    top_P=params["top_P"],
-                    top_K=params["top_K"],
-                    spk=params["spk"],
-                    infer_seed=params["infer_seed"],
-                    prompt1=params["prompt1"],
-                    prompt2=params["prompt2"],
-                    prefix=params["prefix"],
-                )
-                for idx, segment in enumerate(batch):
-                    (sr, audio_data) = audio_datas[idx]
-                    rate = float(segment.get("rate", "1.0"))
-                    volume = float(segment.get("volume", "0"))
-                    pitch = float(segment.get("pitch", "0"))
-
-                    audio_segment = audio_data_to_segment(audio_data, sr)
-                    audio_segment = apply_prosody(audio_segment, rate, volume, pitch)
-                    original_index = segments.index(
-                        segment
-                    )  # Get the original index of the segment
-                    audio_segments[original_index] = (
-                        audio_segment  # Place the audio_segment in the correct position
-                    )
-
-        return audio_segments
 
+        break_segments = buckets.pop("<break>")
+        self.process_break_segments(segments, break_segments, audio_segments)
 
-def generate_audio_segment(
-    text: str,
-    spk: int = -1,
-    seed: int = -1,
-    top_p: float = 0.5,
-    top_k: int = 20,
-    temp: float = 0.3,
-    prompt1: str = "",
-    prompt2: str = "",
-    prefix: str = "",
-    enable_normalize=True,
-    is_end: bool = False,
-) -> AudioSegment:
-    if enable_normalize:
-        text = text_normalize(text, is_end=is_end)
-
-    logger.debug(f"generate segment: {text}")
-
-    sample_rate, audio_data = generate_audio.generate_audio(
-        text=text,
-        temperature=temp if temp is not None else 0.3,
-        top_P=top_p if top_p is not None else 0.5,
-        top_K=top_k if top_k is not None else 20,
-        spk=spk if spk else -1,
-        infer_seed=seed if seed else -1,
-        prompt1=prompt1 if prompt1 else "",
-        prompt2=prompt2 if prompt2 else "",
-        prefix=prefix if prefix else "",
-    )
-
-    byte_io = io.BytesIO()
-    write(byte_io, sample_rate, audio_data)
-    byte_io.seek(0)
+        buckets = list(buckets.values())
 
-    return AudioSegment.from_file(byte_io, format="wav")
-
-
-def synthesize_segment(segment: Dict[str, Any]) -> Union[AudioSegment, None]:
-    if "break" in segment:
-        pause_segment = AudioSegment.silent(duration=segment["break"])
-        return pause_segment
-
-    attrs = segment.get("attrs", {})
-    text = segment.get("text", "")
-    is_end = segment.get("is_end", False)
-
-    text = str(text).strip()
-
-    if text == "":
-        return None
-
-    spk = attrs.get("spk", "")
-    if isinstance(spk, str):
-        spk = int(spk)
-    seed = to_number(attrs.get("seed", ""), int, -1)
-    top_k = to_number(attrs.get("top_k", ""), int, None)
-    top_p = to_number(attrs.get("top_p", ""), float, None)
-    temp = to_number(attrs.get("temp", ""), float, None)
-
-    prompt1 = attrs.get("prompt1", "")
-    prompt2 = attrs.get("prompt2", "")
-    prefix = attrs.get("prefix", "")
-    disable_normalize = attrs.get("normalize", "") == "False"
-
-    audio_segment = generate_audio_segment(
-        text,
-        enable_normalize=not disable_normalize,
-        spk=spk,
-        seed=seed,
-        top_k=top_k,
-        top_p=top_p,
-        temp=temp,
-        prompt1=prompt1,
-        prompt2=prompt2,
-        prefix=prefix,
-        is_end=is_end,
-    )
-
-    rate = float(attrs.get("rate", "1.0"))
-    volume = float(attrs.get("volume", "0"))
-    pitch = float(attrs.get("pitch", "0"))
-
-    audio_segment = apply_prosody(audio_segment, rate, volume, pitch)
+        for bucket in buckets:
+            self.process_voice_segments(segments, bucket, audio_segments)
 
-    return audio_segment
+        return audio_segments
 
 
 # 示例使用
 if __name__ == "__main__":
+    ctx1 = SSMLContext()
+    ctx1.spk = 1
+    ctx1.seed = 42
+    ctx1.temp = 0.1
+    ctx2 = SSMLContext()
+    ctx2.spk = 2
+    ctx2.seed = 42
+    ctx2.temp = 0.1
     ssml_segments = [
-        {
-            "text": "大🍌，一条大🍌，嘿，你的感觉真的很奇妙  [lbreak]",
-            "attrs": {"spk": 2, "temp": 0.1, "seed": 42},
-        },
-        {
-            "text": "大🍉，一个大🍉，嘿，你的感觉真的很奇妙  [lbreak]",
-            "attrs": {"spk": 2, "temp": 0.1, "seed": 42},
-        },
-        {
-            "text": "大🍌，一条大🍌，嘿，你的感觉真的很奇妙  [lbreak]",
-            "attrs": {"spk": 2, "temp": 0.3, "seed": 42},
-        },
+        SSMLSegment(text="大🍌，一条大🍌，嘿，你的感觉真的很奇妙", attrs=ctx1.copy()),
+        SSMLBreak(duration_ms=1000),
+        SSMLSegment(text="大🍉，一个大🍉，嘿，你的感觉真的很奇妙", attrs=ctx1.copy()),
+        SSMLSegment(text="大🍊，一个大🍊，嘿，你的感觉真的很奇妙", attrs=ctx2.copy()),
     ]
 
     synthesizer = SynthesizeSegments(batch_size=2)
     audio_segments = synthesizer.synthesize_segments(ssml_segments)
+    print(audio_segments)
     combined_audio = combine_audio_segments(audio_segments)
     combined_audio.export("output.wav", format="wav")

modules/api/impl/google_api.py

@@ -18,7 +18,6 @@ from modules.ssml import parse_ssml
 from modules.SynthesizeSegments import (
     SynthesizeSegments,
     combine_audio_segments,
-    synthesize_segment,
 )
 
 from modules.api import utils as api_utils

modules/api/impl/speaker_api.py

@@ -7,11 +7,11 @@ from modules.api.Api import APIManager
 
 
 class CreateSpeaker(BaseModel):
-    seed: int
     name: str
     gender: str
     describe: str
-    tensor: list
+    tensor: list = None
+    seed: int = None
 
 
 class UpdateSpeaker(BaseModel):
@@ -76,7 +76,7 @@ def setup(app: APIManager):
                 gender=request.gender,
                 describe=request.describe,
             )
-        else:
+        elif request.seed:
             # from seed
             speaker = speaker_mgr.create_speaker_from_seed(
                 seed=request.seed,
@@ -84,6 +84,10 @@ def setup(app: APIManager):
                 gender=request.gender,
                 describe=request.describe,
             )
+        else:
+            raise HTTPException(
+                status_code=400, detail="Missing tensor or seed in request"
+            )
         return {"message": "ok", "data": speaker.to_json()}
 
     @app.post("/v1/speaker/refresh", response_model=api_utils.BaseResponse)

modules/api/impl/ssml_api.py

@@ -10,7 +10,6 @@ from modules.normalization import text_normalize
 from modules.ssml import parse_ssml
 from modules.SynthesizeSegments import (
     SynthesizeSegments,
-    synthesize_segment,
     combine_audio_segments,
 )
 
@@ -23,6 +22,8 @@ from modules.api.Api import APIManager
 class SSMLRequest(BaseModel):
     ssml: str
     format: str = "mp3"
+
+    # NOTE: 🤔 也许这个值应该配置成系统变量？ 传进来有点奇怪
     batch_size: int = 4
 
 
@@ -48,29 +49,15 @@ async def synthesize_ssml(
         for seg in segments:
             seg["text"] = text_normalize(seg["text"], is_end=True)
 
-        if batch_size != 1:
-            synthesize = SynthesizeSegments(batch_size)
-            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)
-            if format == "mp3":
-                buffer = api_utils.wav_to_mp3(buffer)
-            return StreamingResponse(buffer, media_type=f"audio/{format}")
-        else:
-
-            def audio_streamer():
-                for segment in segments:
-                    audio_segment = synthesize_segment(segment=segment)
-                    buffer = io.BytesIO()
-                    audio_segment.export(buffer, format="wav")
-                    buffer.seek(0)
-                    if format == "mp3":
-                        buffer = api_utils.wav_to_mp3(buffer)
-                    yield buffer.read()
-
-            return StreamingResponse(audio_streamer(), media_type=f"audio/{format}")
+        synthesize = SynthesizeSegments(batch_size)
+        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)
+        if format == "mp3":
+            buffer = api_utils.wav_to_mp3(buffer)
+        return StreamingResponse(buffer, media_type=f"audio/{format}")
 
     except Exception as e:
         import logging

modules/api/utils.py

@@ -52,7 +52,6 @@ def to_number(value, t, default=0):
 def calc_spk_style(spk: Union[str, int], style: Union[str, int]):
     voice_attrs = {
         "spk": None,
-        "seed": None,
         "prompt1": None,
         "prompt2": None,
         "prefix": None,
@@ -85,7 +84,6 @@ def calc_spk_style(spk: Union[str, int], style: Union[str, int]):
     merge_prompt(voice_attrs, params)
 
     voice_attrs["spk"] = params.get("spk", voice_attrs.get("spk", None))
-    voice_attrs["seed"] = params.get("seed", voice_attrs.get("seed", None))
     voice_attrs["temperature"] = params.get(
         "temp", voice_attrs.get("temperature", None)
     )

modules/denoise.py

@@ -1,7 +1,51 @@
-from audio_denoiser.AudioDenoiser import AudioDenoiser
+import os
+from typing import Union
+
 import torch
 import torchaudio
+from modules.Denoiser.AudioDenoiser import AudioDenoiser
+
+from modules.utils.constants import MODELS_DIR
+
+from modules.devices import devices
+
+import soundfile as sf
+
+ad: Union[AudioDenoiser, None] = None
 
 
 class TTSAudioDenoiser:
-    pass
+
+    def load_ad(self):
+        global ad
+        if ad is None:
+            ad = AudioDenoiser(
+                os.path.join(
+                    MODELS_DIR,
+                    "Denoise",
+                    "audio-denoiser-512-32-v1",
+                ),
+                device=devices.device,
+            )
+            ad.model.to(devices.device)
+        return ad
+
+    def denoise(self, audio_data, sample_rate, auto_scale=False):
+        ad = self.load_ad()
+        sr = ad.model_sample_rate
+        return sr, ad.process_waveform(audio_data, sample_rate, auto_scale)
+
+
+if __name__ == "__main__":
+    tts_deno = TTSAudioDenoiser()
+    data, sr = sf.read("test.wav")
+    audio_tensor = torch.from_numpy(data).unsqueeze(0).float()
+    print(audio_tensor)
+
+    # data, sr = torchaudio.load("test.wav")
+    # print(data)
+    # data = data.to(devices.device)
+
+    sr, denoised = tts_deno.denoise(audio_data=audio_tensor, sample_rate=sr)
+    denoised = denoised.cpu()
+    torchaudio.save("denoised.wav", denoised, sample_rate=sr)

modules/generate_audio.py

@@ -79,7 +79,7 @@ def generate_audio_batch(
         params_infer_code["spk_emb"] = spk.emb
         logger.info(("spk", spk.name))
     else:
-        raise ValueError("spk must be int or Speaker")
+        raise ValueError(f"spk must be int or Speaker, but: <{type(spk)}> {spk}")
 
     logger.info(
         {

modules/models.py

@@ -37,17 +37,9 @@ def load_chat_tts_in_thread():
     logger.info("ChatTTS models loaded")
 
 
-def initialize_chat_tts():
+def load_chat_tts():
     with lock:
         if chat_tts is None:
-            model_thread = threading.Thread(target=load_chat_tts_in_thread)
-            model_thread.start()
-            model_thread.join()
-
-
-def load_chat_tts():
-    if chat_tts is None:
-        with lock:
             load_chat_tts_in_thread()
     if chat_tts is None:
         raise Exception("Failed to load ChatTTS models")

modules/speaker.py

@@ -1,5 +1,6 @@
 import os
 from typing import Union
+from box import Box
 import torch
 
 from modules import models
@@ -16,6 +17,18 @@ def create_speaker_from_seed(seed):
 
 
 class Speaker:
+    @staticmethod
+    def from_file(file_like):
+        speaker = torch.load(file_like, map_location=torch.device("cpu"))
+        speaker.fix()
+        return speaker
+
+    @staticmethod
+    def from_tensor(tensor):
+        speaker = Speaker(seed=-2)
+        speaker.emb = tensor
+        return speaker
+
     def __init__(self, seed, name="", gender="", describe=""):
         self.id = uuid.uuid4()
         self.seed = seed
@@ -24,15 +37,20 @@ class Speaker:
         self.describe = describe
         self.emb = None
 
+        # TODO replace emb => tokens
+        self.tokens = []
+
     def to_json(self, with_emb=False):
-        return {
-            "id": str(self.id),
-            "seed": self.seed,
-            "name": self.name,
-            "gender": self.gender,
-            "describe": self.describe,
-            "emb": self.emb.tolist() if with_emb else None,
-        }
+        return Box(
+            **{
+                "id": str(self.id),
+                "seed": self.seed,
+                "name": self.name,
+                "gender": self.gender,
+                "describe": self.describe,
+                "emb": self.emb.tolist() if with_emb else None,
+            }
+        )
 
     def fix(self):
         is_update = False
@@ -78,14 +96,9 @@ class SpeakerManager:
         self.speakers = {}
         for speaker_file in os.listdir(self.speaker_dir):
             if speaker_file.endswith(".pt"):
-                speaker = torch.load(
-                    self.speaker_dir + speaker_file, map_location=torch.device("cpu")
+                self.speakers[speaker_file] = Speaker.from_file(
+                    self.speaker_dir + speaker_file
                 )
-                self.speakers[speaker_file] = speaker
-
-                is_update = speaker.fix()
-                if is_update:
-                    torch.save(speaker, self.speaker_dir + speaker_file)
 
     def list_speakers(self):
         return list(self.speakers.values())
@@ -103,8 +116,8 @@ class SpeakerManager:
     def create_speaker_from_tensor(
         self, tensor, filename="", name="", gender="", describe=""
     ):
-        if name == "":
-            name = filename
+        if filename == "":
+            filename = name
         speaker = Speaker(seed=-2, name=name, gender=gender, describe=describe)
         if isinstance(tensor, torch.Tensor):
             speaker.emb = tensor

modules/ssml_parser/SSMLParser.py

@@ -0,0 +1,178 @@
+from lxml import etree
+
+
+from typing import Any, List, Dict, Union
+import logging
+
+from modules.data import styles_mgr
+from modules.speaker import speaker_mgr
+from box import Box
+import copy
+
+
+class SSMLContext(Box):
+    def __init__(self, parent=None):
+        self.parent: Union[SSMLContext, None] = parent
+
+        self.style = None
+        self.spk = None
+        self.volume = None
+        self.rate = None
+        self.pitch = None
+        # tempurature
+        self.temp = None
+        self.top_p = None
+        self.top_k = None
+        self.seed = None
+        self.noramalize = None
+        self.prompt1 = None
+        self.prompt2 = None
+        self.prefix = None
+
+
+class SSMLSegment(Box):
+    def __init__(self, text: str, attrs=SSMLContext()):
+        self.attrs = attrs
+        self.text = text
+        self.params = None
+
+
+class SSMLBreak:
+    def __init__(self, duration_ms: Union[str, int, float]):
+        # TODO 支持其他单位
+        duration_ms = int(str(duration_ms).replace("ms", ""))
+        self.attrs = Box(**{"duration": duration_ms})
+
+
+class SSMLParser:
+
+    def __init__(self):
+        self.logger = logging.getLogger(__name__)
+        self.logger.debug("SSMLParser.__init__()")
+        self.resolvers = []
+
+    def resolver(self, tag: str):
+        def decorator(func):
+            self.resolvers.append((tag, func))
+            return func
+
+        return decorator
+
+    def parse(self, ssml: str) -> List[Union[SSMLSegment, SSMLBreak]]:
+        root = etree.fromstring(ssml)
+
+        root_ctx = SSMLContext()
+        segments = []
+        self.resolve(root, root_ctx, segments)
+
+        return segments
+
+    def resolve(
+        self, element: etree.Element, context: SSMLContext, segments: List[SSMLSegment]
+    ):
+        resolver = [resolver for tag, resolver in self.resolvers if tag == element.tag]
+        if len(resolver) == 0:
+            raise NotImplementedError(f"Tag {element.tag} not supported.")
+        else:
+            resolver = resolver[0]
+
+        resolver(element, context, segments, self)
+
+
+def create_ssml_parser():
+    parser = SSMLParser()
+
+    @parser.resolver("speak")
+    def tag_speak(element, context, segments, parser):
+        ctx = copy.deepcopy(context)
+
+        version = element.get("version")
+        if version != "0.1":
+            raise ValueError(f"Unsupported SSML version {version}")
+
+        for child in element:
+            parser.resolve(child, ctx, segments)
+
+    @parser.resolver("voice")
+    def tag_voice(element, context, segments, parser):
+        ctx = copy.deepcopy(context)
+
+        ctx.spk = element.get("spk", ctx.spk)
+        ctx.style = element.get("style", ctx.style)
+        ctx.spk = element.get("spk", ctx.spk)
+        ctx.volume = element.get("volume", ctx.volume)
+        ctx.rate = element.get("rate", ctx.rate)
+        ctx.pitch = element.get("pitch", ctx.pitch)
+        # tempurature
+        ctx.temp = element.get("temp", ctx.temp)
+        ctx.top_p = element.get("top_p", ctx.top_p)
+        ctx.top_k = element.get("top_k", ctx.top_k)
+        ctx.seed = element.get("seed", ctx.seed)
+        ctx.noramalize = element.get("noramalize", ctx.noramalize)
+        ctx.prompt1 = element.get("prompt1", ctx.prompt1)
+        ctx.prompt2 = element.get("prompt2", ctx.prompt2)
+        ctx.prefix = element.get("prefix", ctx.prefix)
+
+        # 处理 voice 开头的文本
+        if element.text and element.text.strip():
+            segments.append(SSMLSegment(element.text.strip(), ctx))
+
+        for child in element:
+            parser.resolve(child, ctx, segments)
+
+            # 处理 voice 结尾的文本
+            if child.tail and child.tail.strip():
+                segments.append(SSMLSegment(child.tail.strip(), ctx))
+
+    @parser.resolver("break")
+    def tag_break(element, context, segments, parser):
+        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 = copy.deepcopy(context)
+
+        ctx.spk = element.get("spk", ctx.spk)
+        ctx.style = element.get("style", ctx.style)
+        ctx.spk = element.get("spk", ctx.spk)
+        ctx.volume = element.get("volume", ctx.volume)
+        ctx.rate = element.get("rate", ctx.rate)
+        ctx.pitch = element.get("pitch", ctx.pitch)
+        # tempurature
+        ctx.temp = element.get("temp", ctx.temp)
+        ctx.top_p = element.get("top_p", ctx.top_p)
+        ctx.top_k = element.get("top_k", ctx.top_k)
+        ctx.seed = element.get("seed", ctx.seed)
+        ctx.noramalize = element.get("noramalize", ctx.noramalize)
+        ctx.prompt1 = element.get("prompt1", ctx.prompt1)
+        ctx.prompt2 = element.get("prompt2", ctx.prompt2)
+        ctx.prefix = element.get("prefix", ctx.prefix)
+
+        if element.text and element.text.strip():
+            segments.append(SSMLSegment(element.text.strip(), ctx))
+
+    return parser
+
+
+if __name__ == "__main__":
+    parser = create_ssml_parser()
+
+    ssml = """
+    <speak version="0.1">
+        <voice spk="xiaoyan" style="news">
+            <prosody rate="fast">你好</prosody>
+            <break time="500ms"/>
+            <prosody rate="slow">你好</prosody>
+        </voice>
+    </speak>
+    """
+
+    segments = parser.parse(ssml)
+    for segment in segments:
+        if isinstance(segment, SSMLBreak):
+            print("<break>", segment.attrs)
+        elif isinstance(segment, SSMLSegment):
+            print(segment.text, segment.attrs)
+        else:
+            raise ValueError("Unknown segment type")

modules/ssml_parser/test_ssml_parser.py

@@ -0,0 +1,104 @@
+import pytest
+from lxml import etree
+from modules.ssml_parser.SSMLParser import (
+    create_ssml_parser,
+    SSMLSegment,
+    SSMLBreak,
+    SSMLContext,
+)
+
+
+@pytest.fixture
+def parser():
+    return create_ssml_parser()
+
+
+@pytest.mark.ssml_parser
+def test_speak_tag(parser):
+    ssml = """
+    <speak version="0.1">
+        <voice spk="xiaoyan" style="news">
+            <prosody rate="fast">你好</prosody>
+            <break time="500ms"/>
+            <prosody rate="slow">你好</prosody>
+        </voice>
+    </speak>
+    """
+    segments = parser.parse(ssml)
+    assert len(segments) == 3
+    assert isinstance(segments[0], SSMLSegment)
+    assert segments[0].text == "你好"
+    assert segments[0].params.rate == "fast"
+    assert isinstance(segments[1], SSMLBreak)
+    assert segments[1].duration == 500
+    assert isinstance(segments[2], SSMLSegment)
+    assert segments[2].text == "你好"
+    assert segments[2].params.rate == "slow"
+
+
+@pytest.mark.ssml_parser
+def test_voice_tag(parser):
+    ssml = """
+    <speak version="0.1">
+        <voice spk="xiaoyan" style="news">你好</voice>
+    </speak>
+    """
+    segments = parser.parse(ssml)
+    assert len(segments) == 1
+    assert isinstance(segments[0], SSMLSegment)
+    assert segments[0].text == "你好"
+    assert segments[0].params.spk == "xiaoyan"
+    assert segments[0].params.style == "news"
+
+
+@pytest.mark.ssml_parser
+def test_break_tag(parser):
+    ssml = """
+    <speak version="0.1">
+        <break time="500ms"/>
+    </speak>
+    """
+    segments = parser.parse(ssml)
+    assert len(segments) == 1
+    assert isinstance(segments[0], SSMLBreak)
+    assert segments[0].duration == 500
+
+
+@pytest.mark.ssml_parser
+def test_prosody_tag(parser):
+    ssml = """
+    <speak version="0.1">
+        <prosody rate="fast">你好</prosody>
+    </speak>
+    """
+    segments = parser.parse(ssml)
+    assert len(segments) == 1
+    assert isinstance(segments[0], SSMLSegment)
+    assert segments[0].text == "你好"
+    assert segments[0].params.rate == "fast"
+
+
+@pytest.mark.ssml_parser
+def test_unsupported_version(parser):
+    ssml = """
+    <speak version="0.2">
+        <voice spk="xiaoyan" style="news">你好</voice>
+    </speak>
+    """
+    with pytest.raises(ValueError, match=r"Unsupported SSML version 0.2"):
+        parser.parse(ssml)
+
+
+@pytest.mark.ssml_parser
+def test_unsupported_tag(parser):
+    ssml = """
+    <speak version="0.1">
+        <unsupported>你好</unsupported>
+    </speak>
+    """
+    with pytest.raises(NotImplementedError, match=r"Tag unsupported not supported."):
+        parser.parse(ssml)
+
+
+if __name__ == "__main__":
+    pytest.main()

modules/utils/JsonObject.py

@@ -8,6 +8,9 @@ class JsonObject:
         # If no initial dictionary is provided, use an empty dictionary
         self._dict_obj = initial_dict if initial_dict is not None else {}
 
+        if self._dict_obj is self:
+            raise ValueError("JsonObject cannot be initialized with itself")
+
     def __getattr__(self, name):
         """
         Get an attribute value. If the attribute does not exist,
@@ -111,3 +114,19 @@ class JsonObject:
         :return: A list of values.
         """
         return self._dict_obj.values()
+
+    def clone(self):
+        """
+        Clone the JsonObject.
+
+        :return: A new JsonObject with the same internal dictionary.
+        """
+        return JsonObject(self._dict_obj.copy())
+
+    def merge(self, other):
+        """
+        Merge the internal dictionary with another dictionary.
+
+        :param other: The other dictionary to merge.
+        """
+        self._dict_obj.update(other)

modules/utils/constants.py

@@ -10,4 +10,4 @@ DATA_DIR = os.path.join(ROOT_DIR, "data")
 
 MODELS_DIR = os.path.join(ROOT_DIR, "models")
 
-speakers_dir = os.path.join(DATA_DIR, "speakers")
+SPEAKERS_DIR = os.path.join(DATA_DIR, "speakers")

modules/webui/app.py

@@ -5,7 +5,9 @@ import torch
 import gradio as gr
 
 from modules import config
+from modules.webui import webui_config
 
+from modules.webui.system_tab import create_system_tab
 from modules.webui.tts_tab import create_tts_interface
 from modules.webui.ssml_tab import create_ssml_interface
 from modules.webui.spliter_tab import create_spliter_tab
@@ -93,15 +95,15 @@ def create_interface():
             with gr.TabItem("Spilter"):
                 create_spliter_tab(ssml_input, tabs=tabs)
 
-            if config.runtime_env_vars.webui_experimental:
-                with gr.TabItem("Speaker"):
-                    create_speaker_panel()
-                with gr.TabItem("Denoise"):
-                    gr.Markdown("🚧 Under construction")
-                with gr.TabItem("Inpainting"):
-                    gr.Markdown("🚧 Under construction")
-                with gr.TabItem("ASR"):
-                    gr.Markdown("🚧 Under construction")
+            with gr.TabItem("Speaker"):
+                create_speaker_panel()
+            with gr.TabItem("Inpainting", visible=webui_config.experimental):
+                gr.Markdown("🚧 Under construction")
+            with gr.TabItem("ASR", visible=webui_config.experimental):
+                gr.Markdown("🚧 Under construction")
+
+            with gr.TabItem("System"):
+                create_system_tab()
 
             with gr.TabItem("README"):
                 create_readme_tab()

modules/webui/speaker_tab.py

@@ -1,13 +1,259 @@
+import io
 import gradio as gr
+import torch
 
-from modules.webui.webui_utils import get_speakers
+from modules.hf import spaces
+from modules.webui.webui_utils import get_speakers, tts_generate
+from modules.speaker import speaker_mgr, Speaker
+
+import tempfile
+
+
+def spk_to_tensor(spk):
+    spk = spk.split(" : ")[1].strip() if " : " in spk else spk
+    if spk == "None" or spk == "":
+        return None
+    return speaker_mgr.get_speaker(spk).emb
+
+
+def get_speaker_show_name(spk):
+    if spk.gender == "*" or spk.gender == "":
+        return spk.name
+    return f"{spk.gender} : {spk.name}"
+
+
+def merge_spk(
+    spk_a,
+    spk_a_w,
+    spk_b,
+    spk_b_w,
+    spk_c,
+    spk_c_w,
+    spk_d,
+    spk_d_w,
+):
+    tensor_a = spk_to_tensor(spk_a)
+    tensor_b = spk_to_tensor(spk_b)
+    tensor_c = spk_to_tensor(spk_c)
+    tensor_d = spk_to_tensor(spk_d)
+
+    assert (
+        tensor_a is not None
+        or tensor_b is not None
+        or tensor_c is not None
+        or tensor_d is not None
+    ), "At least one speaker should be selected"
+
+    merge_tensor = torch.zeros_like(
+        tensor_a
+        if tensor_a is not None
+        else (
+            tensor_b
+            if tensor_b is not None
+            else tensor_c if tensor_c is not None else tensor_d
+        )
+    )
+
+    total_weight = 0
+    if tensor_a is not None:
+        merge_tensor += spk_a_w * tensor_a
+        total_weight += spk_a_w
+    if tensor_b is not None:
+        merge_tensor += spk_b_w * tensor_b
+        total_weight += spk_b_w
+    if tensor_c is not None:
+        merge_tensor += spk_c_w * tensor_c
+        total_weight += spk_c_w
+    if tensor_d is not None:
+        merge_tensor += spk_d_w * tensor_d
+        total_weight += spk_d_w
+
+    if total_weight > 0:
+        merge_tensor /= total_weight
+
+    merged_spk = Speaker.from_tensor(merge_tensor)
+    merged_spk.name = "<MIX>"
+
+    return merged_spk
+
+
+@torch.inference_mode()
+@spaces.GPU
+def merge_and_test_spk_voice(
+    spk_a, spk_a_w, spk_b, spk_b_w, spk_c, spk_c_w, spk_d, spk_d_w, test_text
+):
+    merged_spk = merge_spk(
+        spk_a,
+        spk_a_w,
+        spk_b,
+        spk_b_w,
+        spk_c,
+        spk_c_w,
+        spk_d,
+        spk_d_w,
+    )
+    return tts_generate(
+        spk=merged_spk,
+        text=test_text,
+    )
+
+
+@torch.inference_mode()
+@spaces.GPU
+def merge_spk_to_file(
+    spk_a,
+    spk_a_w,
+    spk_b,
+    spk_b_w,
+    spk_c,
+    spk_c_w,
+    spk_d,
+    spk_d_w,
+    speaker_name,
+    speaker_gender,
+    speaker_desc,
+):
+    merged_spk = merge_spk(
+        spk_a, spk_a_w, spk_b, spk_b_w, spk_c, spk_c_w, spk_d, spk_d_w
+    )
+    merged_spk.name = speaker_name
+    merged_spk.gender = speaker_gender
+    merged_spk.desc = speaker_desc
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".pt") as tmp_file:
+        torch.save(merged_spk, tmp_file)
+        tmp_file_path = tmp_file.name
+
+    return tmp_file_path
+
+
+merge_desc = """
+## Speaker Merger
+
+在本面板中，您可以选择多个说话人并指定他们的权重，合成新的语音并进行测试。以下是各个功能的详细说明：
+
+### 1. 选择说话人
+您可以从下拉菜单中选择最多四个说话人（A、B、C、D），每个说话人都有一个对应的权重滑块，范围从0到10。权重决定了每个说话人在合成语音中的影响程度。
+
+### 2. 合成语音
+在选择好说话人和设置好权重后，您可以在“测试文本”框中输入要测试的文本，然后点击“测试语音”按钮来生成并播放合成的语音。
+
+### 3. 保存说话人
+您还可以在右侧的“说话人信息”部分填写新的说话人的名称、性别和描述，并点击“保存说话人”按钮来保存合成的说话人。保存后的说话人文件将显示在“合成说话人”栏中，供下载使用。
+"""
 
 
 # 显示 a b c d 四个选择框，选择一个或多个，然后可以试音，并导出
 def create_speaker_panel():
     speakers = get_speakers()
 
-    def get_speaker_show_name(spk):
-        pass
+    speaker_names = ["None"] + [get_speaker_show_name(speaker) for speaker in speakers]
+
+    with gr.Tabs():
+        with gr.TabItem("Merger"):
+            gr.Markdown(merge_desc)
+
+            with gr.Row():
+                with gr.Column(scale=5):
+                    with gr.Row():
+                        with gr.Group():
+                            spk_a = gr.Dropdown(
+                                choices=speaker_names, value="None", label="Speaker A"
+                            )
+                            spk_a_w = gr.Slider(
+                                value=1, minimum=0, maximum=10, step=1, label="Weight A"
+                            )
+
+                        with gr.Group():
+                            spk_b = gr.Dropdown(
+                                choices=speaker_names, value="None", label="Speaker B"
+                            )
+                            spk_b_w = gr.Slider(
+                                value=1, minimum=0, maximum=10, step=1, label="Weight B"
+                            )
+
+                        with gr.Group():
+                            spk_c = gr.Dropdown(
+                                choices=speaker_names, value="None", label="Speaker C"
+                            )
+                            spk_c_w = gr.Slider(
+                                value=1, minimum=0, maximum=10, step=1, label="Weight C"
+                            )
+
+                        with gr.Group():
+                            spk_d = gr.Dropdown(
+                                choices=speaker_names, value="None", label="Speaker D"
+                            )
+                            spk_d_w = gr.Slider(
+                                value=1, minimum=0, maximum=10, step=1, label="Weight D"
+                            )
+
+                    with gr.Row():
+                        with gr.Column(scale=3):
+                            with gr.Group():
+                                gr.Markdown("🎤Test voice")
+                                with gr.Row():
+                                    test_voice_btn = gr.Button(
+                                        "Test Voice", variant="secondary"
+                                    )
+
+                                    with gr.Column(scale=4):
+                                        test_text = gr.Textbox(
+                                            label="Test Text",
+                                            placeholder="Please input test text",
+                                            value="说话人合并测试 123456789 [uv_break] ok, test done [lbreak]",
+                                        )
+
+                                        output_audio = gr.Audio(label="Output Audio")
+
+                with gr.Column(scale=1):
+                    with gr.Group():
+                        gr.Markdown("🗃️Save to file")
+
+                        speaker_name = gr.Textbox(
+                            label="Name", value="forge_speaker_merged"
+                        )
+                        speaker_gender = gr.Textbox(label="Gender", value="*")
+                        speaker_desc = gr.Textbox(
+                            label="Description", value="merged speaker"
+                        )
+
+                        save_btn = gr.Button("Save Speaker", variant="primary")
+
+                        merged_spker = gr.File(
+                            label="Merged Speaker", interactive=False, type="binary"
+                        )
+
+            test_voice_btn.click(
+                merge_and_test_spk_voice,
+                inputs=[
+                    spk_a,
+                    spk_a_w,
+                    spk_b,
+                    spk_b_w,
+                    spk_c,
+                    spk_c_w,
+                    spk_d,
+                    spk_d_w,
+                    test_text,
+                ],
+                outputs=[output_audio],
+            )
 
-    gr.Markdown("🚧 Under construction")
+            save_btn.click(
+                merge_spk_to_file,
+                inputs=[
+                    spk_a,
+                    spk_a_w,
+                    spk_b,
+                    spk_b_w,
+                    spk_c,
+                    spk_c_w,
+                    spk_d,
+                    spk_d_w,
+                    speaker_name,
+                    speaker_gender,
+                    speaker_desc,
+                ],
+                outputs=[merged_spker],
+            )

modules/webui/spliter_tab.py

@@ -9,6 +9,7 @@ from modules.webui.webui_utils import (
 from modules.hf import spaces
 
 
+# NOTE: 因为 text_normalize 需要使用 tokenizer
 @torch.inference_mode()
 @spaces.GPU
 def merge_dataframe_to_ssml(dataframe, spk, style, seed):
@@ -31,7 +32,7 @@ def merge_dataframe_to_ssml(dataframe, spk, style, seed):
         if seed:
             ssml += f' seed="{seed}"'
         ssml += ">\n"
-        ssml += f"{indent}{indent}{text_normalize(row[1])}\n"
+        ssml += f"{indent}{indent}{text_normalize(row.iloc[1])}\n"
         ssml += f"{indent}</voice>\n"
     return f"<speak version='0.1'>\n{ssml}</speak>"
 

modules/webui/system_tab.py

@@ -0,0 +1,15 @@
+import gradio as gr
+from modules.webui import webui_config
+
+
+def create_system_tab():
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown(f"info")
+
+        with gr.Column(scale=5):
+            toggle_experimental = gr.Checkbox(
+                label="Enable Experimental Features",
+                value=webui_config.experimental,
+                interactive=False,
+            )

modules/webui/tts_tab.py

@@ -3,6 +3,7 @@ import torch
 from modules.webui.webui_utils import (
     get_speakers,
     get_styles,
+    load_spk_info,
     refine_text,
     tts_generate,
 )
@@ -10,6 +11,13 @@ from modules.webui import webui_config
 from modules.webui.examples import example_texts
 from modules import config
 
+default_text_content = """
+chat T T S 是一款强大的对话式文本转语音模型。它有中英混读和多说话人的能力。
+chat T T S 不仅能够生成自然流畅的语音，还能控制[laugh]笑声啊[laugh]，
+停顿啊[uv_break]语气词啊等副语言现象[uv_break]。这个韵律超越了许多开源模型[uv_break]。
+请注意，chat T T S 的使用应遵守法律和伦理准则，避免滥用的安全风险。[uv_break]
+"""
+
 
 def create_tts_interface():
     speakers = get_speakers()
@@ -90,15 +98,18 @@ def create_tts_interface():
                                 outputs=[spk_input_text],
                             )
 
-                        if config.runtime_env_vars.webui_experimental:
-                            with gr.Tab(label="Upload"):
-                                spk_input_upload = gr.File(label="Speaker (Upload)")
-                                # TODO 读取 speaker
-                                # spk_input_upload.change(
-                                #     fn=lambda x: x.read().decode("utf-8"),
-                                #     inputs=[spk_input_upload],
-                                #     outputs=[spk_input_text],
-                                # )
+                        with gr.Tab(label="Upload"):
+                            spk_file_upload = gr.File(label="Speaker (Upload)")
+
+                            gr.Markdown("📝Speaker info")
+                            infos = gr.Markdown("empty")
+
+                            spk_file_upload.change(
+                                fn=load_spk_info,
+                                inputs=[spk_file_upload],
+                                outputs=[infos],
+                            ),
+
             with gr.Group():
                 gr.Markdown("💃Inference Seed")
                 infer_seed_input = gr.Number(
@@ -122,85 +133,62 @@ def create_tts_interface():
                 prompt2_input = gr.Textbox(label="Prompt 2")
                 prefix_input = gr.Textbox(label="Prefix")
 
-                if config.runtime_env_vars.webui_experimental:
-                    prompt_audio = gr.File(label="prompt_audio")
+                prompt_audio = gr.File(
+                    label="prompt_audio", visible=webui_config.experimental
+                )
 
             infer_seed_rand_button.click(
                 lambda x: int(torch.randint(0, 2**32 - 1, (1,)).item()),
                 inputs=[infer_seed_input],
                 outputs=[infer_seed_input],
             )
-        with gr.Column(scale=3):
-            with gr.Row():
-                with gr.Column(scale=4):
-                    with gr.Group():
-                        input_title = gr.Markdown(
-                            "📝Text Input",
-                            elem_id="input-title",
-                        )
-                        gr.Markdown(
-                            f"- 字数限制{webui_config.tts_max:,}字，超过部分截断"
-                        )
-                        gr.Markdown("- 如果尾字吞字不读，可以试试结尾加上 `[lbreak]`")
-                        gr.Markdown(
-                            "- If the input text is all in English, it is recommended to check disable_normalize"
-                        )
-                        text_input = gr.Textbox(
-                            show_label=False,
-                            label="Text to Speech",
-                            lines=10,
-                            placeholder="输入文本或选择示例",
-                            elem_id="text-input",
-                        )
-                        # TODO 字数统计，其实实现很好写，但是就是会触发loading...并且还要和后端交互...
-                        # text_input.change(
-                        #     fn=lambda x: (
-                        #         f"📝Text Input ({len(x)} char)"
-                        #         if x
-                        #         else (
-                        #             "📝Text Input (0 char)"
-                        #             if not x
-                        #             else "📝Text Input (0 char)"
-                        #         )
-                        #     ),
-                        #     inputs=[text_input],
-                        #     outputs=[input_title],
-                        # )
-                        with gr.Row():
-                            contorl_tokens = [
-                                "[laugh]",
-                                "[uv_break]",
-                                "[v_break]",
-                                "[lbreak]",
-                            ]
-
-                            for tk in contorl_tokens:
-                                t_btn = gr.Button(tk)
-                                t_btn.click(
-                                    lambda text, tk=tk: text + " " + tk,
-                                    inputs=[text_input],
-                                    outputs=[text_input],
-                                )
-                with gr.Column(scale=1):
-                    with gr.Group():
-                        gr.Markdown("🎶Refiner")
-                        refine_prompt_input = gr.Textbox(
-                            label="Refine Prompt",
-                            value="[oral_2][laugh_0][break_6]",
-                        )
-                        refine_button = gr.Button("✍️Refine Text")
-                        # TODO 分割句子，使用当前配置拼接为SSML，然后发送到SSML tab
-                        # send_button = gr.Button("📩Split and send to SSML")
-
-                    with gr.Group():
-                        gr.Markdown("🔊Generate")
-                        disable_normalize_input = gr.Checkbox(
-                            value=False, label="Disable Normalize"
-                        )
-                        tts_button = gr.Button(
-                            "🔊Generate Audio",
-                            variant="primary",
-                            elem_classes="big-button",
+        with gr.Column(scale=4):
+            with gr.Group():
+                input_title = gr.Markdown(
+                    "📝Text Input",
+                    elem_id="input-title",
+                )
+                gr.Markdown(f"- 字数限制{webui_config.tts_max:,}字，超过部分截断")
+                gr.Markdown("- 如果尾字吞字不读，可以试试结尾加上 `[lbreak]`")
+                gr.Markdown(
+                    "- If the input text is all in English, it is recommended to check disable_normalize"
+                )
+                text_input = gr.Textbox(
+                    show_label=False,
+                    label="Text to Speech",
+                    lines=10,
+                    placeholder="输入文本或选择示例",
+                    elem_id="text-input",
+                    value=default_text_content,
+                )
+                # TODO 字数统计，其实实现很好写，但是就是会触发loading...并且还要和后端交互...
+                # text_input.change(
+                #     fn=lambda x: (
+                #         f"📝Text Input ({len(x)} char)"
+                #         if x
+                #         else (
+                #             "📝Text Input (0 char)"
+                #             if not x
+                #             else "📝Text Input (0 char)"
+                #         )
+                #     ),
+                #     inputs=[text_input],
+                #     outputs=[input_title],
+                # )
+                with gr.Row():
+                    contorl_tokens = [
+                        "[laugh]",
+                        "[uv_break]",
+                        "[v_break]",
+                        "[lbreak]",
+                    ]
+
+                    for tk in contorl_tokens:
+                        t_btn = gr.Button(tk)
+                        t_btn.click(
+                            lambda text, tk=tk: text + " " + tk,
+                            inputs=[text_input],
+                            outputs=[text_input],
                         )
 
             with gr.Group():
@@ -220,6 +208,31 @@ def create_tts_interface():
             with gr.Group():
                 gr.Markdown("🎨Output")
                 tts_output = gr.Audio(label="Generated Audio")
+        with gr.Column(scale=1):
+            with gr.Group():
+                gr.Markdown("🎶Refiner")
+                refine_prompt_input = gr.Textbox(
+                    label="Refine Prompt",
+                    value="[oral_2][laugh_0][break_6]",
+                )
+                refine_button = gr.Button("✍️Refine Text")
+
+            with gr.Group():
+                gr.Markdown("🔊Generate")
+                disable_normalize_input = gr.Checkbox(
+                    value=False, label="Disable Normalize"
+                )
+
+                # FIXME: 不知道为啥，就是非常慢，单独调脚本是很快的
+                with gr.Group(visible=webui_config.experimental):
+                    gr.Markdown("💪🏼Enhance")
+                    enable_enhance = gr.Checkbox(value=False, label="Enable Enhance")
+                    enable_de_noise = gr.Checkbox(value=False, label="Enable De-noise")
+                tts_button = gr.Button(
+                    "🔊Generate Audio",
+                    variant="primary",
+                    elem_classes="big-button",
+                )
 
     refine_button.click(
         refine_text,
@@ -243,6 +256,9 @@ def create_tts_interface():
             style_input_dropdown,
             disable_normalize_input,
             batch_size_input,
+            enable_enhance,
+            enable_de_noise,
+            spk_file_upload,
         ],
         outputs=tts_output,
     )

modules/webui/webui_config.py

@@ -1,4 +1,8 @@
+from typing import Literal
+
+
 tts_max = 1000
 ssml_max = 1000
 spliter_threshold = 100
 max_batch_size = 8
+experimental = False

modules/webui/webui_utils.py

@@ -1,37 +1,26 @@
-import os
-import logging
-import sys
-
+from typing import Union
 import numpy as np
 
+from modules.Enhancer.ResembleEnhance import load_enhancer
 from modules.devices import devices
 from modules.synthesize_audio import synthesize_audio
 from modules.hf import spaces
 from modules.webui import webui_config
 
-logging.basicConfig(
-    level=os.getenv("LOG_LEVEL", "INFO"),
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
-
-
-import gradio as gr
-
 import torch
 
-from modules.ssml import parse_ssml
+from modules.ssml_parser.SSMLParser import create_ssml_parser, SSMLBreak, SSMLSegment
 from modules.SynthesizeSegments import SynthesizeSegments, combine_audio_segments
 
-from modules.speaker import speaker_mgr
+from modules.speaker import speaker_mgr, Speaker
 from modules.data import styles_mgr
 
 from modules.api.utils import calc_spk_style
-import modules.generate_audio as generate
 
 from modules.normalization import text_normalize
-from modules import refiner, config
+from modules import refiner
 
-from modules.utils import env, audio
+from modules.utils import audio
 from modules.SentenceSplitter import SentenceSplitter
 
 
@@ -43,11 +32,30 @@ def get_styles():
     return styles_mgr.list_items()
 
 
-def segments_length_limit(segments, total_max: int):
+def load_spk_info(file):
+    if file is None:
+        return "empty"
+    try:
+
+        spk: Speaker = Speaker.from_file(file)
+        infos = spk.to_json()
+        return f"""
+- name: {infos.name}
+- gender: {infos.gender}
+- describe: {infos.describe}
+    """.strip()
+    except:
+        return "load failed"
+
+
+def segments_length_limit(
+    segments: list[Union[SSMLBreak, SSMLSegment]], total_max: int
+) -> list[Union[SSMLBreak, SSMLSegment]]:
     ret_segments = []
     total_len = 0
     for seg in segments:
-        if "text" not in seg:
+        if isinstance(seg, SSMLBreak):
+            ret_segments.append(seg)
             continue
         total_len += len(seg["text"])
         if total_len > total_max:
@@ -56,6 +64,28 @@ def segments_length_limit(segments, total_max: int):
     return ret_segments
 
 
+@torch.inference_mode()
+@spaces.GPU
+def apply_audio_enhance(audio_data, sr, enable_denoise, enable_enhance):
+    audio_data = torch.from_numpy(audio_data).float().squeeze().cpu()
+    if enable_denoise or enable_enhance:
+        enhancer = load_enhancer(devices.device)
+        if enable_denoise:
+            audio_data, sr = enhancer.denoise(audio_data, sr, devices.device)
+        if enable_enhance:
+            audio_data, sr = enhancer.enhance(
+                audio_data,
+                sr,
+                devices.device,
+                tau=0.9,
+                nfe=64,
+                solver="euler",
+                lambd=0.5,
+            )
+    audio_data = audio_data.cpu().numpy()
+    return audio_data, int(sr)
+
+
 @torch.inference_mode()
 @spaces.GPU
 def synthesize_ssml(ssml: str, batch_size=4):
@@ -69,7 +99,8 @@ def synthesize_ssml(ssml: str, batch_size=4):
     if ssml == "":
         return None
 
-    segments = parse_ssml(ssml)
+    parser = create_ssml_parser()
+    segments = parser.parse(ssml)
     max_len = webui_config.ssml_max
     segments = segments_length_limit(segments, max_len)
 
@@ -87,18 +118,21 @@ def synthesize_ssml(ssml: str, batch_size=4):
 @spaces.GPU
 def tts_generate(
     text,
-    temperature,
-    top_p,
-    top_k,
-    spk,
-    infer_seed,
-    use_decoder,
-    prompt1,
-    prompt2,
-    prefix,
-    style,
+    temperature=0.3,
+    top_p=0.7,
+    top_k=20,
+    spk=-1,
+    infer_seed=-1,
+    use_decoder=True,
+    prompt1="",
+    prompt2="",
+    prefix="",
+    style="",
     disable_normalize=False,
     batch_size=4,
+    enable_enhance=False,
+    enable_denoise=False,
+    spk_file=None,
 ):
     try:
         batch_size = int(batch_size)
@@ -126,12 +160,15 @@ def tts_generate(
     prompt1 = prompt1 or params.get("prompt1", "")
     prompt2 = prompt2 or params.get("prompt2", "")
 
-    infer_seed = np.clip(infer_seed, -1, 2**32 - 1, out=None, dtype=np.int64)
+    infer_seed = np.clip(infer_seed, -1, 2**32 - 1, out=None, dtype=np.float64)
     infer_seed = int(infer_seed)
 
     if not disable_normalize:
         text = text_normalize(text)
 
+    if spk_file:
+        spk = Speaker.from_file(spk_file)
+
     sample_rate, audio_data = synthesize_audio(
         text=text,
         temperature=temperature,
@@ -146,6 +183,10 @@ def tts_generate(
         batch_size=batch_size,
     )
 
+    audio_data, sample_rate = apply_audio_enhance(
+        audio_data, sample_rate, enable_denoise, enable_enhance
+    )
+
     audio_data = audio.audio_to_int16(audio_data)
     return sample_rate, audio_data
 

webui.py

@@ -93,8 +93,10 @@ if __name__ == "__main__":
     device_id = get_and_update_env(args, "device_id", None, str)
     use_cpu = get_and_update_env(args, "use_cpu", [], list)
     compile = get_and_update_env(args, "compile", False, bool)
-    webui_experimental = get_and_update_env(args, "webui_experimental", False, bool)
 
+    webui_config.experimental = get_and_update_env(
+        args, "webui_experimental", False, bool
+    )
     webui_config.tts_max = get_and_update_env(args, "tts_max_len", 1000, int)
     webui_config.ssml_max = get_and_update_env(args, "ssml_max_len", 5000, int)
     webui_config.max_batch_size = get_and_update_env(args, "max_batch_size", 8, int)