Spaces:
Sleeping
Sleeping
# A unified script for inference process | |
# Make adjustments inside functions, and consider both gradio and cli scripts if need to change func output format | |
import re | |
import tempfile | |
import numpy as np | |
import torch | |
import torchaudio | |
import tqdm | |
from pydub import AudioSegment, silence | |
from transformers import pipeline | |
from vocos import Vocos | |
from model import CFM | |
from model.utils import ( | |
load_checkpoint, | |
get_tokenizer, | |
convert_char_to_pinyin, | |
) | |
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" | |
vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz") | |
# ----------------------------------------- | |
target_sample_rate = 24000 | |
n_mel_channels = 100 | |
hop_length = 256 | |
target_rms = 0.1 | |
cross_fade_duration = 0.15 | |
ode_method = "euler" | |
nfe_step = 32 # 16, 32 | |
cfg_strength = 2.0 | |
sway_sampling_coef = -1.0 | |
speed = 1.0 | |
fix_duration = None | |
# ----------------------------------------- | |
# chunk text into smaller pieces | |
def chunk_text(text, max_chars=135): | |
""" | |
Splits the input text into chunks, each with a maximum number of characters. | |
Args: | |
text (str): The text to be split. | |
max_chars (int): The maximum number of characters per chunk. | |
Returns: | |
List[str]: A list of text chunks. | |
""" | |
chunks = [] | |
current_chunk = "" | |
# Split the text into sentences based on punctuation followed by whitespace | |
sentences = re.split(r"(?<=[;:,.!?])\s+|(?<=[;:,。!?])", text) | |
for sentence in sentences: | |
if len(current_chunk.encode("utf-8")) + len(sentence.encode("utf-8")) <= max_chars: | |
current_chunk += sentence + " " if sentence and len(sentence[-1].encode("utf-8")) == 1 else sentence | |
else: | |
if current_chunk: | |
chunks.append(current_chunk.strip()) | |
current_chunk = sentence + " " if sentence and len(sentence[-1].encode("utf-8")) == 1 else sentence | |
if current_chunk: | |
chunks.append(current_chunk.strip()) | |
return chunks | |
# load vocoder | |
def load_vocoder(is_local=False, local_path="", device=device): | |
if is_local: | |
print(f"Load vocos from local path {local_path}") | |
vocos = Vocos.from_hparams(f"{local_path}/config.yaml") | |
state_dict = torch.load(f"{local_path}/pytorch_model.bin", map_location=device) | |
vocos.load_state_dict(state_dict) | |
vocos.eval() | |
else: | |
print("Download Vocos from huggingface charactr/vocos-mel-24khz") | |
vocos = Vocos.from_pretrained("charactr/vocos-mel-24khz") | |
return vocos | |
# load asr pipeline | |
asr_pipe = None | |
def initialize_asr_pipeline(device=device): | |
global asr_pipe | |
asr_pipe = pipeline( | |
"automatic-speech-recognition", | |
model="openai/whisper-large-v3-turbo", | |
torch_dtype=torch.float16, | |
device=device, | |
) | |
# load model for inference | |
def load_model(model_cls, model_cfg, ckpt_path, vocab_file="", ode_method=ode_method, use_ema=True, device=device): | |
if vocab_file == "": | |
vocab_file = "Emilia_ZH_EN" | |
tokenizer = "pinyin" | |
else: | |
tokenizer = "custom" | |
print("\nvocab : ", vocab_file) | |
print("tokenizer : ", tokenizer) | |
print("model : ", ckpt_path, "\n") | |
vocab_char_map, vocab_size = get_tokenizer(vocab_file, tokenizer) | |
model = CFM( | |
transformer=model_cls(**model_cfg, text_num_embeds=vocab_size, mel_dim=n_mel_channels), | |
mel_spec_kwargs=dict( | |
target_sample_rate=target_sample_rate, | |
n_mel_channels=n_mel_channels, | |
hop_length=hop_length, | |
), | |
odeint_kwargs=dict( | |
method=ode_method, | |
), | |
vocab_char_map=vocab_char_map, | |
).to(device) | |
model = load_checkpoint(model, ckpt_path, device, use_ema=use_ema) | |
return model | |
# preprocess reference audio and text | |
def preprocess_ref_audio_text(ref_audio_orig, ref_text, show_info=print, device=device): | |
show_info("Converting audio...") | |
with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f: | |
aseg = AudioSegment.from_file(ref_audio_orig) | |
non_silent_segs = silence.split_on_silence(aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=1000) | |
non_silent_wave = AudioSegment.silent(duration=0) | |
for non_silent_seg in non_silent_segs: | |
non_silent_wave += non_silent_seg | |
aseg = non_silent_wave | |
audio_duration = len(aseg) | |
if audio_duration > 15000: | |
show_info("Audio is over 15s, clipping to only first 15s.") | |
aseg = aseg[:15000] | |
aseg.export(f.name, format="wav") | |
ref_audio = f.name | |
if not ref_text.strip(): | |
global asr_pipe | |
if asr_pipe is None: | |
initialize_asr_pipeline(device=device) | |
show_info("No reference text provided, transcribing reference audio...") | |
ref_text = asr_pipe( | |
ref_audio, | |
chunk_length_s=30, | |
batch_size=128, | |
generate_kwargs={"task": "transcribe"}, | |
return_timestamps=False, | |
)["text"].strip() | |
show_info("Finished transcription") | |
else: | |
show_info("Using custom reference text...") | |
# Add the functionality to ensure it ends with ". " | |
if not ref_text.endswith(". ") and not ref_text.endswith("。"): | |
if ref_text.endswith("."): | |
ref_text += " " | |
else: | |
ref_text += ". " | |
return ref_audio, ref_text | |
# infer process: chunk text -> infer batches [i.e. infer_batch_process()] | |
def infer_process( | |
ref_audio, | |
ref_text, | |
gen_text, | |
model_obj, | |
show_info=print, | |
progress=tqdm, | |
target_rms=target_rms, | |
cross_fade_duration=cross_fade_duration, | |
nfe_step=nfe_step, | |
cfg_strength=cfg_strength, | |
sway_sampling_coef=sway_sampling_coef, | |
speed=speed, | |
fix_duration=fix_duration, | |
device=device, | |
): | |
# Split the input text into batches | |
audio, sr = torchaudio.load(ref_audio) | |
max_chars = int(len(ref_text.encode("utf-8")) / (audio.shape[-1] / sr) * (25 - audio.shape[-1] / sr)) | |
gen_text_batches = chunk_text(gen_text, max_chars=max_chars) | |
for i, gen_text in enumerate(gen_text_batches): | |
print(f"gen_text {i}", gen_text) | |
show_info(f"Generating audio in {len(gen_text_batches)} batches...") | |
return infer_batch_process( | |
(audio, sr), | |
ref_text, | |
gen_text_batches, | |
model_obj, | |
progress=progress, | |
target_rms=target_rms, | |
cross_fade_duration=cross_fade_duration, | |
nfe_step=nfe_step, | |
cfg_strength=cfg_strength, | |
sway_sampling_coef=sway_sampling_coef, | |
speed=speed, | |
fix_duration=fix_duration, | |
device=device, | |
) | |
# infer batches | |
def infer_batch_process( | |
ref_audio, | |
ref_text, | |
gen_text_batches, | |
model_obj, | |
progress=tqdm, | |
target_rms=0.1, | |
cross_fade_duration=0.15, | |
nfe_step=32, | |
cfg_strength=2.0, | |
sway_sampling_coef=-1, | |
speed=1, | |
fix_duration=None, | |
device=None, | |
): | |
audio, sr = ref_audio | |
if audio.shape[0] > 1: | |
audio = torch.mean(audio, dim=0, keepdim=True) | |
rms = torch.sqrt(torch.mean(torch.square(audio))) | |
if rms < target_rms: | |
audio = audio * target_rms / rms | |
if sr != target_sample_rate: | |
resampler = torchaudio.transforms.Resample(sr, target_sample_rate) | |
audio = resampler(audio) | |
audio = audio.to(device) | |
generated_waves = [] | |
spectrograms = [] | |
if len(ref_text[-1].encode("utf-8")) == 1: | |
ref_text = ref_text + " " | |
for i, gen_text in enumerate(progress.tqdm(gen_text_batches)): | |
# Prepare the text | |
text_list = [ref_text + gen_text] | |
final_text_list = convert_char_to_pinyin(text_list) | |
ref_audio_len = audio.shape[-1] // hop_length | |
if fix_duration is not None: | |
duration = int(fix_duration * target_sample_rate / hop_length) | |
else: | |
# Calculate duration | |
ref_text_len = len(ref_text.encode("utf-8")) | |
gen_text_len = len(gen_text.encode("utf-8")) | |
duration = ref_audio_len + int(ref_audio_len / ref_text_len * gen_text_len / speed) | |
# inference | |
with torch.inference_mode(): | |
generated, _ = model_obj.sample( | |
cond=audio, | |
text=final_text_list, | |
duration=duration, | |
steps=nfe_step, | |
cfg_strength=cfg_strength, | |
sway_sampling_coef=sway_sampling_coef, | |
) | |
generated = generated.to(torch.float32) | |
generated = generated[:, ref_audio_len:, :] | |
generated_mel_spec = generated.permute(0, 2, 1) | |
generated_wave = vocos.decode(generated_mel_spec.cpu()) | |
if rms < target_rms: | |
generated_wave = generated_wave * rms / target_rms | |
# wav -> numpy | |
generated_wave = generated_wave.squeeze().cpu().numpy() | |
generated_waves.append(generated_wave) | |
spectrograms.append(generated_mel_spec[0].cpu().numpy()) | |
# Combine all generated waves with cross-fading | |
if cross_fade_duration <= 0: | |
# Simply concatenate | |
final_wave = np.concatenate(generated_waves) | |
else: | |
final_wave = generated_waves[0] | |
for i in range(1, len(generated_waves)): | |
prev_wave = final_wave | |
next_wave = generated_waves[i] | |
# Calculate cross-fade samples, ensuring it does not exceed wave lengths | |
cross_fade_samples = int(cross_fade_duration * target_sample_rate) | |
cross_fade_samples = min(cross_fade_samples, len(prev_wave), len(next_wave)) | |
if cross_fade_samples <= 0: | |
# No overlap possible, concatenate | |
final_wave = np.concatenate([prev_wave, next_wave]) | |
continue | |
# Overlapping parts | |
prev_overlap = prev_wave[-cross_fade_samples:] | |
next_overlap = next_wave[:cross_fade_samples] | |
# Fade out and fade in | |
fade_out = np.linspace(1, 0, cross_fade_samples) | |
fade_in = np.linspace(0, 1, cross_fade_samples) | |
# Cross-faded overlap | |
cross_faded_overlap = prev_overlap * fade_out + next_overlap * fade_in | |
# Combine | |
new_wave = np.concatenate( | |
[prev_wave[:-cross_fade_samples], cross_faded_overlap, next_wave[cross_fade_samples:]] | |
) | |
final_wave = new_wave | |
# Create a combined spectrogram | |
combined_spectrogram = np.concatenate(spectrograms, axis=1) | |
return final_wave, target_sample_rate, combined_spectrogram | |
# remove silence from generated wav | |
def remove_silence_for_generated_wav(filename): | |
aseg = AudioSegment.from_file(filename) | |
non_silent_segs = silence.split_on_silence(aseg, min_silence_len=1000, silence_thresh=-50, keep_silence=500) | |
non_silent_wave = AudioSegment.silent(duration=0) | |
for non_silent_seg in non_silent_segs: | |
non_silent_wave += non_silent_seg | |
aseg = non_silent_wave | |
aseg.export(filename, format="wav") | |