diff --git a/.env.webui b/.env.webui index a807cd0298c85378bbefe6e65990411d0b9ee51b..22d73786f170f6fa0741125fe3841a6eed8f0128 100644 --- a/.env.webui +++ b/.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 diff --git a/language/zh-CN.json b/language/zh-CN.json index 31e5890a575dcac42c44b75b7675650a38f22553..f4f41cf038fe73c2194b311b9357b8f5a3b77d6d 100644 --- a/language/zh-CN.json +++ b/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]`", diff --git a/modules/ChatTTS/ChatTTS/core.py b/modules/ChatTTS/ChatTTS/core.py index 225de72f07cb237cfa7210872d316936afe808f8..549973e0c5dcdf9869ae1237a65fc7762ceae244 100644 --- a/modules/ChatTTS/ChatTTS/core.py +++ b/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: diff --git a/modules/SynthesizeSegments.py b/modules/SynthesizeSegments.py index e7589ab85d6ebdbd3618d6a668a0bf60c6a51b56..de8a7778a27b2d89e6058a15716a0c53538c9a71 100644 --- a/modules/SynthesizeSegments.py +++ b/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) diff --git a/modules/api/api_setup.py b/modules/api/api_setup.py index e7de2a62e4131afe6fb5db0280feca8288f4d79d..bfe07f4de7b7f9ddc4a2c625579f1ecd07aa1e2a 100644 --- a/modules/api/api_setup.py +++ b/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) diff --git a/modules/api/impl/google_api.py b/modules/api/impl/google_api.py index cd66e036578866f35b2cf8a9f1f559e1e2ca7e90..6244bc9fa8ab4c625b5ec8a04982be20798f8ed5 100644 --- a/modules/api/impl/google_api.py +++ b/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." diff --git a/modules/api/impl/openai_api.py b/modules/api/impl/openai_api.py index f1e21e1241c958047d3b8488105981528fde82eb..7c0e012c093c13bb6ce19415a3e504cc3fcafe46 100644 --- a/modules/api/impl/openai_api.py +++ b/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: diff --git a/modules/api/impl/ssml_api.py b/modules/api/impl/ssml_api.py index 2696470d6afaa6f5ba6bac9a75b36e2bd6164ce8..c6277b6214fe18a5f9e271c766f1f16d9d3f981f 100644 --- a/modules/api/impl/ssml_api.py +++ b/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) diff --git a/modules/api/impl/tts_api.py b/modules/api/impl/tts_api.py index 7330b4820612341e8ca57fea0a04b3dbe3eadfa5..b91f5493ca898c5d7522110004a65443849536b8 100644 --- a/modules/api/impl/tts_api.py +++ b/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() diff --git a/modules/api/impl/xtts_v2_api.py b/modules/api/impl/xtts_v2_api.py new file mode 100644 index 0000000000000000000000000000000000000000..0b660562b2ae1041f46c62fe17355040723a6590 --- /dev/null +++ b/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)) diff --git a/modules/devices/devices.py b/modules/devices/devices.py index c9e1862cdd427f19edf669b2b6c4daee9d6f3340..e11f6eba5da29af778ee1d254778acaa7492b6bb 100644 --- a/modules/devices/devices.py +++ b/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 diff --git a/modules/finetune/__init__.py b/modules/finetune/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/modules/finetune/model/__init__.py b/modules/finetune/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/modules/finetune/model/encoder.py b/modules/finetune/model/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..d445ad83d07ebdd9bfaa25e9e6d9c64001471d08 --- /dev/null +++ b/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) diff --git a/modules/finetune/model/wavenet.py b/modules/finetune/model/wavenet.py new file mode 100644 index 0000000000000000000000000000000000000000..828aa41c8afdbbb59f038d5053357a317721b5c0 --- /dev/null +++ b/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 diff --git a/modules/finetune/train_gpt.py b/modules/finetune/train_gpt.py new file mode 100644 index 0000000000000000000000000000000000000000..9642d37016a0b99348d11a79dcdf7b2bdd3c0aef --- /dev/null +++ b/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() diff --git a/modules/finetune/train_speaker.py b/modules/finetune/train_speaker.py new file mode 100644 index 0000000000000000000000000000000000000000..343d743d6c2acff51c170fde6fc3ae32bb47c482 --- /dev/null +++ b/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 +""" diff --git a/modules/finetune/utils/__init__.py b/modules/finetune/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/modules/finetune/utils/dataset.py b/modules/finetune/utils/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..e5d86dae96de753e9c648b62329844958c3a33eb --- /dev/null +++ b/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}") diff --git a/modules/finetune/utils/logger.py b/modules/finetune/utils/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..7e59b5cfd272cfb4294822924aba108e143d7310 --- /dev/null +++ b/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) diff --git a/modules/finetune/utils/model.py b/modules/finetune/utils/model.py new file mode 100644 index 0000000000000000000000000000000000000000..416cb1bcc7084c8d0e065e4de36a75e43ab47fa6 --- /dev/null +++ b/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 diff --git a/modules/finetune/utils/output.py b/modules/finetune/utils/output.py new file mode 100644 index 0000000000000000000000000000000000000000..541092ddfa1f33848e1d8ff914ffbeab312db44f --- /dev/null +++ b/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() diff --git a/modules/generate_audio.py b/modules/generate_audio.py index 9fcabe3954c3912b1f008ed8865ca04c12e18a14..a2e4552b9103d2bb13dd030724f93f740ab7f1b8 100644 --- a/modules/generate_audio.py +++ b/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: diff --git a/modules/normalization.py b/modules/normalization.py index 1d740e1ca6b914a37deceb515409e088cb5c29d2..cc6e941f78143b2b46b5eb8f886f55f68417c77f 100644 --- a/modules/normalization.py +++ b/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%的概率会下雨 +埃隆·马斯克 """, ] diff --git a/modules/repos_static/resemble_enhance/data/distorter/base.py b/modules/repos_static/resemble_enhance/data/distorter/base.py index d43d84fa840dd25804d9c5e5dc9673f65fdc5b94..f07ef407fa92190234ead9f7de43d7c5ea3c6b4d 100644 --- a/modules/repos_static/resemble_enhance/data/distorter/base.py +++ b/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 diff --git a/modules/repos_static/resemble_enhance/data/distorter/custom.py b/modules/repos_static/resemble_enhance/data/distorter/custom.py index 28428f7789cebb2d174c581111711f4d73f6565b..fdabed6aac1647de9a7ee887f84308effa71c8da 100644 --- a/modules/repos_static/resemble_enhance/data/distorter/custom.py +++ b/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)}") diff --git a/modules/repos_static/resemble_enhance/data/distorter/sox.py b/modules/repos_static/resemble_enhance/data/distorter/sox.py index 92a2d74033d33b975141c1ece7ac5619d1dfcc39..3e08376087683222dd5db98f4c4b25ad0e38b847 100644 --- a/modules/repos_static/resemble_enhance/data/distorter/sox.py +++ b/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): diff --git a/modules/repos_static/resemble_enhance/data/utils.py b/modules/repos_static/resemble_enhance/data/utils.py index 77f59d345b75cac76c6c423c734ae9c70a626590..38ca25fe36074d615962dc229599cf1b3a548aaa 100644 --- a/modules/repos_static/resemble_enhance/data/utils.py +++ b/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) diff --git a/modules/repos_static/resemble_enhance/denoiser/denoiser.py b/modules/repos_static/resemble_enhance/denoiser/denoiser.py index c0d9c2b6ffbc471029cee620216a2d080b9dd100..4d672df3431d1877d3c8cb882aa8606d6d8b5d1f 100644 --- a/modules/repos_static/resemble_enhance/denoiser/denoiser.py +++ b/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 diff --git a/modules/repos_static/resemble_enhance/enhancer/download.py b/modules/repos_static/resemble_enhance/enhancer/download.py index 614b9a4b4f9a1a10b79f12ca1a25821247ea2a16..089181893229ba67c9202e204f994d512975f9fc 100644 --- a/modules/repos_static/resemble_enhance/enhancer/download.py +++ b/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(): diff --git a/modules/repos_static/resemble_enhance/enhancer/enhancer.py b/modules/repos_static/resemble_enhance/enhancer/enhancer.py index c7ab9417deb429855b7fce43962426f6b6c4a9c0..1ea3f351752d8e8e13040fef842372367926c3e4 100644 --- a/modules/repos_static/resemble_enhance/enhancer/enhancer.py +++ b/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) diff --git a/modules/repos_static/resemble_enhance/enhancer/hparams.py b/modules/repos_static/resemble_enhance/enhancer/hparams.py index ca89bea6f5d7d4ec4f543f8bde88b29dcae69f6a..7878e4172b5772c59aea0de54d2537f0523d9437 100644 --- a/modules/repos_static/resemble_enhance/enhancer/hparams.py +++ b/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 diff --git a/modules/repos_static/resemble_enhance/enhancer/inference.py b/modules/repos_static/resemble_enhance/enhancer/inference.py index af57a2c7d3e5cc7b08b00f85f0135e881e50fcbe..dc7712cb6a4d2126bb4d740d24ed9355312741ef 100644 --- a/modules/repos_static/resemble_enhance/enhancer/inference.py +++ b/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) diff --git a/modules/repos_static/resemble_enhance/enhancer/lcfm/cfm.py b/modules/repos_static/resemble_enhance/enhancer/lcfm/cfm.py index a5125267b7f32e11c58e4b96bffa3ba1e96fdc4f..09b4a3e45ce3b50cca7ce7debe77ddb230ee9783 100644 --- a/modules/repos_static/resemble_enhance/enhancer/lcfm/cfm.py +++ b/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: diff --git a/modules/repos_static/resemble_enhance/enhancer/lcfm/irmae.py b/modules/repos_static/resemble_enhance/enhancer/lcfm/irmae.py index 91f5dbb506187271c67c7bbbf55475021854ab27..aa82827c8809b001d31827d76bbee731e11ae2e2 100644 --- a/modules/repos_static/resemble_enhance/enhancer/lcfm/irmae.py +++ b/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): diff --git a/modules/repos_static/resemble_enhance/enhancer/lcfm/lcfm.py b/modules/repos_static/resemble_enhance/enhancer/lcfm/lcfm.py index 4c2f5f88718e2f42f82e2f4714ea510b4677b450..8d1c241312f96525fcf7630e805560cbe9b84406 100644 --- a/modules/repos_static/resemble_enhance/enhancer/lcfm/lcfm.py +++ b/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 diff --git a/modules/repos_static/resemble_enhance/enhancer/univnet/univnet.py b/modules/repos_static/resemble_enhance/enhancer/univnet/univnet.py index bb20217f048f398236698f6a38927310d0c1ba9b..602f08851095b7a25a1bddc8a2daa7e48fc10cb1 100644 --- a/modules/repos_static/resemble_enhance/enhancer/univnet/univnet.py +++ b/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) diff --git a/modules/repos_static/resemble_enhance/hparams.py b/modules/repos_static/resemble_enhance/hparams.py index a8e716175fa962ada1d98cd755430e2ea770278c..9f796e97c3ab1c3d540d9aed14c8bf0796a7d39b 100644 --- a/modules/repos_static/resemble_enhance/hparams.py +++ b/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] diff --git a/modules/speaker.py b/modules/speaker.py index 46ceaf947a5be1c2915c51d29c5a48707388af82..764cf48839fe8127af63704f9ea947710ea1b3a4 100644 --- a/modules/speaker.py +++ b/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 = [] diff --git a/modules/ssml_parser/SSMLParser.py b/modules/ssml_parser/SSMLParser.py index 5db290006a51a809178c7c0198a51a9d6324a888..bda224cd2e9846cff34eb21e7430f1a5b7f9e42a 100644 --- a/modules/ssml_parser/SSMLParser.py +++ b/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) diff --git a/modules/synthesize_audio.py b/modules/synthesize_audio.py index a07c7bc69de6352cfec6a9c01d6c0203ac4f8d94..c032abc2ee99237b729c23b9e4e2cd1b0cf683d9 100644 --- a/modules/synthesize_audio.py +++ b/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) diff --git a/modules/utils/audio.py b/modules/utils/audio.py index 48f38c598db590bad30687e519db78f1b0b491af..b1a97eeee49ea8aa9d877fc3e9cdeb6e8f1ea1cf 100644 --- a/modules/utils/audio.py +++ b/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) diff --git a/modules/webui/app.py b/modules/webui/app.py index 1f64fa8d33ac67801307735ac3cfe26c63e384fa..b1acbf9e5f1fa6be03466d739b8f8445bbd853f7 100644 --- a/modules/webui/app.py +++ b/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() diff --git a/modules/webui/finetune/ProcessMonitor.py b/modules/webui/finetune/ProcessMonitor.py new file mode 100644 index 0000000000000000000000000000000000000000..a92c187ae0de80ec0ad93f56d65e623d4a916c55 --- /dev/null +++ b/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) diff --git a/modules/webui/finetune/ft_tab.py b/modules/webui/finetune/ft_tab.py new file mode 100644 index 0000000000000000000000000000000000000000..ac7147f40e598abf846ef5ee2d6e3e8a6bf005b9 --- /dev/null +++ b/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") diff --git a/modules/webui/finetune/ft_ui_utils.py b/modules/webui/finetune/ft_ui_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c2a8e8ca09ebe5a8c9c6df6f6e336a77de10fb17 --- /dev/null +++ b/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 diff --git a/modules/webui/finetune/speaker_ft_tab.py b/modules/webui/finetune/speaker_ft_tab.py new file mode 100644 index 0000000000000000000000000000000000000000..f652ec422a5b71acfaec98aa74f93c72b61ae32b --- /dev/null +++ b/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]) diff --git a/modules/webui/localization_runtime.py b/modules/webui/localization_runtime.py index 9689c960e900ad1ec93c3e85e5c09d8bb5a54626..273eb05c66676f525fb484b6a6a64a1129091462 100644 --- a/modules/webui/localization_runtime.py +++ b/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", + ], + ] diff --git a/modules/webui/ssml/podcast_tab.py b/modules/webui/ssml/podcast_tab.py index 440e0b60eb5f8b394a86cc6dac7266ec018fb6ec..32e732d34fb9c4b4e86cf88b5e549312c88174b4 100644 --- a/modules/webui/ssml/podcast_tab.py +++ b/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"), ) diff --git a/modules/webui/ssml/ssml_tab.py b/modules/webui/ssml/ssml_tab.py index f2de84c2e28bcd35e7c14611cc9ff3fdb57bdfc7..6fa6dd861daa3fd246d56aec4ded84ef068537d3 100644 --- a/modules/webui/ssml/ssml_tab.py +++ b/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, ) diff --git a/modules/webui/tts_tab.py b/modules/webui/tts_tab.py index c39e7ed284cca32eae897cfcfaf80559a1c3d49b..ab81e12243bbd83eccf0a0f026cb9650f86df690 100644 --- a/modules/webui/tts_tab.py +++ b/modules/webui/tts_tab.py @@ -29,32 +29,6 @@ def create_tts_interface(): with gr.Row(): with gr.Column(scale=1): - with gr.Group(): - gr.Markdown("🎛️Sampling") - temperature_input = gr.Slider( - 0.01, 2.0, value=0.3, step=0.01, label="Temperature" - ) - top_p_input = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Top P") - top_k_input = gr.Slider(1, 50, value=20, step=1, label="Top K") - batch_size_input = gr.Slider( - 1, - webui_config.max_batch_size, - value=4, - step=1, - label="Batch Size", - ) - - with gr.Row(): - with gr.Group(): - gr.Markdown("🎭Style") - gr.Markdown("TTS_STYLE_GUIDE") - style_input_dropdown = gr.Dropdown( - choices=styles, - # label="Choose Style", - interactive=True, - show_label=False, - value="*auto", - ) with gr.Row(): with gr.Group(): gr.Markdown("🗣️Speaker") @@ -102,7 +76,47 @@ def create_tts_interface(): fn=load_spk_info, inputs=[spk_file_upload], outputs=[infos], - ), + ) + + with gr.Row(): + with gr.Group(): + gr.Markdown("🎭Style") + gr.Markdown("TTS_STYLE_GUIDE") + style_input_dropdown = gr.Dropdown( + choices=styles, + # label="Choose Style", + interactive=True, + show_label=False, + value="*auto", + ) + + with gr.Group(): + gr.Markdown("🎛️Sampling") + temperature_input = gr.Slider( + 0.01, 2.0, value=0.3, step=0.01, label="Temperature" + ) + top_p_input = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Top P") + top_k_input = gr.Slider(1, 50, value=20, step=1, label="Top K") + batch_size_input = gr.Slider( + 1, + webui_config.max_batch_size, + value=4, + step=1, + label="Batch Size", + ) + 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("💃Inference Seed") @@ -202,7 +216,8 @@ def create_tts_interface(): ) refine_button = gr.Button("✍️Refine Text") - with gr.Group(): + # 由于使用不是很方便，所以列为实验性功能 + with gr.Group(visible=webui_config.experimental): gr.Markdown("🔧Prompt engineering") prompt1_input = gr.Textbox(label="Prompt 1") prompt2_input = gr.Textbox(label="Prompt 2") @@ -253,6 +268,8 @@ def create_tts_interface(): enable_enhance, enable_de_noise, spk_file_upload, + spliter_thr_input, + eos_input, ], outputs=tts_output, ) diff --git a/modules/webui/webui_utils.py b/modules/webui/webui_utils.py index 4a6d6dedf17c6a37c16a3546afab089a69599d38..cf57c80c597d20508a7e70a3d544e90f9bafe92a 100644 --- a/modules/webui/webui_utils.py +++ b/modules/webui/webui_utils.py @@ -95,6 +95,8 @@ def synthesize_ssml( batch_size=4, enable_enhance=False, enable_denoise=False, + eos: str = "[uv_break]", + spliter_thr: int = 100, ): try: batch_size = int(batch_size) @@ -114,7 +116,9 @@ def synthesize_ssml( if len(segments) == 0: return None - synthesize = SynthesizeSegments(batch_size=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) @@ -151,6 +155,8 @@ def tts_generate( enable_enhance=False, enable_denoise=False, spk_file=None, + spliter_thr: int = 100, + eos: str = "[uv_break]", ): try: batch_size = int(batch_size) @@ -199,6 +205,8 @@ def tts_generate( prompt2=prompt2, prefix=prefix, batch_size=batch_size, + end_of_sentence=eos, + spliter_threshold=spliter_thr, ) audio_data, sample_rate = apply_audio_enhance( diff --git a/webui.py b/webui.py index 055874043e85b90485fee64bc4982cd4d3373486..d2b7fdb075b9eb61ca9f0e37c63cb9e4c7469ea2 100644 --- a/webui.py +++ b/webui.py @@ -1,4 +1,5 @@ import os +import sys import logging from modules.api.api_setup import ( @@ -106,6 +107,7 @@ def process_webui_args(args): auth=auth, show_api=False, prevent_thread_lock=True, + inbrowser=sys.platform == "win32", app_kwargs={ "title": app_title, "description": app_description,