ChatTTS-Forge / modules /generate_audio.py
zhzluke96
update
bed01bd
import gc
import logging
from typing import Generator, Union
import numpy as np
import torch
from modules import config, models
from modules.ChatTTS import ChatTTS
from modules.devices import devices
from modules.speaker import Speaker
from modules.utils.cache import conditional_cache
from modules.utils.SeedContext import SeedContext
logger = logging.getLogger(__name__)
SAMPLE_RATE = 24000
def generate_audio(
text: str,
temperature: float = 0.3,
top_P: float = 0.7,
top_K: float = 20,
spk: Union[int, Speaker] = -1,
infer_seed: int = -1,
use_decoder: bool = True,
prompt1: str = "",
prompt2: str = "",
prefix: str = "",
):
(sample_rate, wav) = generate_audio_batch(
[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,
)[0]
return (sample_rate, wav)
def parse_infer_params(
texts: list[str],
chat_tts: ChatTTS.Chat,
temperature: float = 0.3,
top_P: float = 0.7,
top_K: float = 20,
spk: Union[int, Speaker] = -1,
infer_seed: int = -1,
prompt1: str = "",
prompt2: str = "",
prefix: str = "",
):
params_infer_code = {
"spk_emb": None,
"temperature": temperature,
"top_P": top_P,
"top_K": top_K,
"prompt1": prompt1 or "",
"prompt2": prompt2 or "",
"prefix": prefix or "",
"repetition_penalty": 1.0,
"disable_tqdm": config.runtime_env_vars.off_tqdm,
}
if isinstance(spk, int):
with SeedContext(spk, True):
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:
logger.warn(
f"spk must be int or Speaker, but: <{type(spk)}> {spk}, wiil set to default voice"
)
with SeedContext(2, True):
params_infer_code["spk_emb"] = chat_tts.sample_random_speaker()
logger.debug(
{
"text": texts,
"infer_seed": infer_seed,
"temperature": temperature,
"top_P": top_P,
"top_K": top_K,
"prompt1": prompt1 or "",
"prompt2": prompt2 or "",
"prefix": prefix or "",
}
)
return params_infer_code
@torch.inference_mode()
def generate_audio_batch(
texts: list[str],
temperature: float = 0.3,
top_P: float = 0.7,
top_K: float = 20,
spk: Union[int, Speaker] = -1,
infer_seed: int = -1,
use_decoder: bool = True,
prompt1: str = "",
prompt2: str = "",
prefix: str = "",
):
chat_tts = models.load_chat_tts()
params_infer_code = parse_infer_params(
texts=texts,
chat_tts=chat_tts,
temperature=temperature,
top_P=top_P,
top_K=top_K,
spk=spk,
infer_seed=infer_seed,
prompt1=prompt1,
prompt2=prompt2,
prefix=prefix,
)
with SeedContext(infer_seed, True):
wavs = chat_tts.generate_audio(
prompt=texts, params_infer_code=params_infer_code, use_decoder=use_decoder
)
if config.auto_gc:
devices.torch_gc()
gc.collect()
return [(SAMPLE_RATE, np.array(wav).flatten().astype(np.float32)) for wav in wavs]
# TODO: generate_audio_stream 也应该支持 lru cache
@torch.inference_mode()
def generate_audio_stream(
text: str,
temperature: float = 0.3,
top_P: float = 0.7,
top_K: float = 20,
spk: Union[int, Speaker] = -1,
infer_seed: int = -1,
use_decoder: bool = True,
prompt1: str = "",
prompt2: str = "",
prefix: str = "",
) -> Generator[tuple[int, np.ndarray], None, None]:
chat_tts = models.load_chat_tts()
texts = [text]
params_infer_code = parse_infer_params(
texts=texts,
chat_tts=chat_tts,
temperature=temperature,
top_P=top_P,
top_K=top_K,
spk=spk,
infer_seed=infer_seed,
prompt1=prompt1,
prompt2=prompt2,
prefix=prefix,
)
with SeedContext(infer_seed, True):
wavs_gen = chat_tts.generate_audio(
prompt=texts,
params_infer_code=params_infer_code,
use_decoder=use_decoder,
stream=True,
)
for wav in wavs_gen:
yield [SAMPLE_RATE, np.array(wav).flatten().astype(np.float32)]
if config.auto_gc:
devices.torch_gc()
gc.collect()
return
lru_cache_enabled = False
def setup_lru_cache():
global generate_audio_batch
global lru_cache_enabled
if lru_cache_enabled:
return
lru_cache_enabled = True
def should_cache(*args, **kwargs):
spk_seed = kwargs.get("spk", -1)
infer_seed = kwargs.get("infer_seed", -1)
return spk_seed != -1 and infer_seed != -1
lru_size = config.runtime_env_vars.lru_size
if isinstance(lru_size, int):
generate_audio_batch = conditional_cache(lru_size, should_cache)(
generate_audio_batch
)
logger.info(f"LRU cache enabled with size {lru_size}")
else:
logger.debug(f"LRU cache failed to enable, invalid size {lru_size}")
if __name__ == "__main__":
import soundfile as sf
# 测试batch生成
inputs = ["你好[lbreak]", "再见[lbreak]", "长度不同的文本片段[lbreak]"]
outputs = generate_audio_batch(inputs, spk=5, infer_seed=42)
for i, (sample_rate, wav) in enumerate(outputs):
print(i, sample_rate, wav.shape)
sf.write(f"batch_{i}.wav", wav, sample_rate, format="wav")
# 单独生成
for i, text in enumerate(inputs):
sample_rate, wav = generate_audio(text, spk=5, infer_seed=42)
print(i, sample_rate, wav.shape)
sf.write(f"one_{i}.wav", wav, sample_rate, format="wav")