import copy
import io
import json
import logging
import re
from typing import List, Union
import numpy as np
from box import Box
from pydub import AudioSegment
from scipy.io import wavfile
from modules import generate_audio
from modules.api.utils import calc_spk_style
from modules.normalization import text_normalize
from modules.SentenceSplitter import SentenceSplitter
from modules.speaker import Speaker
from modules.ssml_parser.SSMLParser import SSMLBreak, SSMLContext, SSMLSegment
from modules.utils import rng
from modules.utils.audio import apply_prosody_to_audio_segment
logger = logging.getLogger(__name__)
def audio_data_to_segment_slow(audio_data, sr):
byte_io = io.BytesIO()
wavfile.write(byte_io, rate=sr, data=audio_data)
byte_io.seek(0)
return AudioSegment.from_file(byte_io, format="wav")
def clip_audio(audio_data: np.ndarray, threshold: float = 0.99):
audio_data = np.clip(audio_data, -threshold, threshold)
return audio_data
def normalize_audio(audio_data: np.ndarray, norm_factor: float = 0.8):
max_amplitude = np.max(np.abs(audio_data))
if max_amplitude > 0:
audio_data = audio_data / max_amplitude * norm_factor
return audio_data
def audio_data_to_segment(audio_data: np.ndarray, sr: int):
"""
optimize: https://github.com/lenML/ChatTTS-Forge/issues/57
"""
audio_data = normalize_audio(audio_data)
audio_data = clip_audio(audio_data)
audio_data = (audio_data * 32767).astype(np.int16)
audio_segment = AudioSegment(
audio_data.tobytes(),
frame_rate=sr,
sample_width=audio_data.dtype.itemsize,
channels=1,
)
return audio_segment
def combine_audio_segments(audio_segments: list[AudioSegment]) -> AudioSegment:
combined_audio = AudioSegment.empty()
for segment in audio_segments:
combined_audio += segment
return combined_audio
def to_number(value, t, default=0):
try:
number = t(value)
return number
except (ValueError, TypeError) as e:
return default
class TTSAudioSegment(Box):
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, 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]
) -> TTSAudioSegment:
if isinstance(segment, SSMLBreak):
return TTSAudioSegment(_type="break")
if segment.get("params", None) is not None:
params = segment.get("params")
text = segment.get("text", None) or segment.text or ""
return TTSAudioSegment(**params, text=text)
text = segment.get("text", None) or segment.text or ""
is_end = segment.get("is_end", False)
text = str(text).strip()
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"
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:
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)
)
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]
def append_eos(text: str):
text = text.strip()
eos_arr = ["[uv_break]", "[v_break]", "[lbreak]", "[llbreak]"]
has_eos = False
for eos in eos_arr:
if eos in text:
has_eos = True
break
if not has_eos:
text += self.eos
return text
# 这里会添加 end_of_text 到 text 之后
texts = [append_eos(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"))
audio_segment = audio_data_to_segment(audio_data, sr)
audio_segment = apply_prosody_to_audio_segment(
audio_segment, rate=rate, volume=volume, pitch=pitch
)
# compare by Box object
original_index = src_segments.index(segment)
audio_segments[original_index] = audio_segment
def bucket_segments(
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)
if isinstance(params.spk, Speaker):
params.spk = str(params.spk.id)
key = json.dumps(
{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)
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:
seg = SSMLSegment(
text=sentence,
attrs=segment.attrs.copy(),
params=copy.copy(segment.params),
)
ret_segments.append(seg)
setattr(seg, "_idx", len(ret_segments) - 1)
def is_none_speak_segment(segment: SSMLSegment):
text = segment.text.strip()
regexp = r"\[[^\]]+?\]"
text = re.sub(regexp, "", text)
text = text.strip()
if not text:
return True
return False
# 将 none_speak 合并到前一个 speak segment
for i in range(1, len(ret_segments)):
if is_none_speak_segment(ret_segments[i]):
ret_segments[i - 1].text += ret_segments[i].text
ret_segments[i].text = ""
# 移除空的 segment
ret_segments = [seg for seg in ret_segments if seg.text.strip()]
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)
break_segments = buckets.pop("<break>")
self.process_break_segments(segments, break_segments, audio_segments)
buckets = list(buckets.values())
for bucket in buckets:
self.process_voice_segments(segments, bucket, audio_segments)
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 = [
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")