Spaces:
Runtime error
Runtime error
| import sys | |
| import traceback | |
| from pathlib import Path | |
| from time import perf_counter as timer | |
| import re | |
| import numpy as np | |
| import torch | |
| import soundfile as sf | |
| import librosa | |
| import spacy | |
| import encoder | |
| from encoder import inference as encoder_infer | |
| from synthesizer.inference import Synthesizer_infer | |
| from synthesizer.utils.cleaners import add_breaks, english_cleaners_predict | |
| from vocoder.display import save_attention_multiple, save_spectrogram, save_stop_tokens | |
| from synthesizer.hparams import syn_hparams | |
| from toolbox.ui import UI | |
| from toolbox.utterance import Utterance | |
| from vocoder import inference as vocoder | |
| from speed_changer.fixSpeed import * | |
| import time | |
| # Use this directory structure for your datasets, or modify it to fit your needs | |
| recognized_datasets = [ | |
| "LibriSpeech/dev-clean", | |
| "LibriSpeech/dev-other", | |
| "LibriSpeech/test-clean", | |
| "LibriSpeech/test-other", | |
| "LibriSpeech/train-clean-100", | |
| "LibriSpeech/train-clean-360", | |
| "LibriSpeech/train-other-500", | |
| "LibriTTS/dev-clean", | |
| "LibriTTS/dev-other", | |
| "LibriTTS/test-clean", | |
| "LibriTTS/test-other", | |
| "LibriTTS/train-clean-100", | |
| "LibriTTS/train-clean-360", | |
| "LibriTTS/train-other-500", | |
| "LJSpeech-1.1", | |
| "VoxCeleb1/wav", | |
| "VoxCeleb1/test_wav", | |
| "VoxCeleb2/dev/aac", | |
| "VoxCeleb2/test/aac", | |
| "VCTK-Corpus/wav48", | |
| ] | |
| # Maximum of generated wavs to keep on memory | |
| MAX_WAVS = 15 | |
| class Toolbox: | |
| def __init__(self, run_id: str, datasets_root: Path, models_dir: Path, seed: int=None): | |
| sys.excepthook = self.excepthook | |
| self.datasets_root = datasets_root | |
| self.utterances = set() | |
| self.current_generated = (None, None, None, None) # speaker_name, spec, breaks, wav | |
| self.synthesizer = None # type: Synthesizer_infer | |
| self.current_wav = None | |
| self.waves_list = [] | |
| self.waves_count = 0 | |
| self.waves_namelist = [] | |
| self.start_generate_time = None | |
| self.nlp = spacy.load('en_core_web_sm') | |
| if not os.path.exists("toolbox_results"): | |
| os.mkdir("toolbox_results") | |
| # Check for webrtcvad (enables removal of silences in vocoder output) | |
| try: | |
| import webrtcvad | |
| self.trim_silences = True | |
| except: | |
| self.trim_silences = False | |
| # Initialize the events and the interface | |
| self.ui = UI() | |
| self.reset_ui(run_id, models_dir, seed) | |
| self.setup_events() | |
| self.ui.start() | |
| def excepthook(self, exc_type, exc_value, exc_tb): | |
| traceback.print_exception(exc_type, exc_value, exc_tb) | |
| self.ui.log("Exception: %s" % exc_value) | |
| def setup_events(self): | |
| # Dataset, speaker and utterance selection | |
| self.ui.browser_load_button.clicked.connect(lambda: self.load_from_browser()) | |
| random_func = lambda level: lambda: self.ui.populate_browser(self.datasets_root, | |
| recognized_datasets, | |
| level) | |
| self.ui.random_dataset_button.clicked.connect(random_func(0)) | |
| self.ui.random_speaker_button.clicked.connect(random_func(1)) | |
| self.ui.random_utterance_button.clicked.connect(random_func(2)) | |
| self.ui.dataset_box.currentIndexChanged.connect(random_func(1)) | |
| self.ui.speaker_box.currentIndexChanged.connect(random_func(2)) | |
| # Model selection | |
| self.ui.encoder_box.currentIndexChanged.connect(self.init_encoder) | |
| def func(): | |
| self.synthesizer = None | |
| self.ui.synthesizer_box.currentIndexChanged.connect(func) | |
| self.ui.vocoder_box.currentIndexChanged.connect(self.init_vocoder) | |
| # Utterance selection | |
| func = lambda: self.load_from_browser(self.ui.browse_file()) | |
| self.ui.browser_browse_button.clicked.connect(func) | |
| func = lambda: self.ui.draw_utterance(self.ui.selected_utterance, "current") | |
| self.ui.utterance_history.currentIndexChanged.connect(func) | |
| func = lambda: self.ui.play(self.ui.selected_utterance.wav, Synthesizer_infer.sample_rate) | |
| self.ui.play_button.clicked.connect(func) | |
| self.ui.stop_button.clicked.connect(self.ui.stop) | |
| self.ui.record_button.clicked.connect(self.record) | |
| #Audio | |
| self.ui.setup_audio_devices(Synthesizer_infer.sample_rate) | |
| #Wav playback & save | |
| func = lambda: self.replay_last_wav() | |
| self.ui.replay_wav_button.clicked.connect(func) | |
| func = lambda: self.export_current_wave() | |
| self.ui.export_wav_button.clicked.connect(func) | |
| self.ui.waves_cb.currentIndexChanged.connect(self.set_current_wav) | |
| # Generation | |
| func = lambda: self.synthesize() or self.vocode() | |
| self.ui.generate_button.clicked.connect(func) | |
| self.ui.synthesize_button.clicked.connect(self.synthesize) | |
| self.ui.vocode_button.clicked.connect(self.vocode) | |
| self.ui.random_seed_checkbox.clicked.connect(self.update_seed_textbox) | |
| # UMAP legend | |
| self.ui.clear_button.clicked.connect(self.clear_utterances) | |
| def set_current_wav(self, index): | |
| self.current_wav = self.waves_list[index] | |
| def export_current_wave(self): | |
| self.ui.save_audio_file(self.current_wav, Synthesizer_infer.sample_rate) | |
| def replay_last_wav(self): | |
| self.ui.play(self.current_wav, Synthesizer_infer.sample_rate) | |
| def reset_ui(self, run_id: str, models_dir: Path, seed: int=None): | |
| self.ui.populate_browser(self.datasets_root, recognized_datasets, 0, True) | |
| self.ui.populate_models(run_id, models_dir) | |
| self.ui.populate_gen_options(seed, self.trim_silences) | |
| def load_from_browser(self, fpath=None): | |
| if fpath is None: | |
| fpath = Path(self.datasets_root, | |
| self.ui.current_dataset_name, | |
| self.ui.current_speaker_name, | |
| self.ui.current_utterance_name) | |
| name = str(fpath.relative_to(self.datasets_root)) | |
| speaker_name = self.ui.current_dataset_name + '_' + self.ui.current_speaker_name | |
| # Select the next utterance | |
| if self.ui.auto_next_checkbox.isChecked(): | |
| self.ui.browser_select_next() | |
| elif fpath == "": | |
| return | |
| else: | |
| name = fpath.name | |
| speaker_name = fpath.parent.name | |
| # Get the wav from the disk. We take the wav with the vocoder/synthesizer format for | |
| # playback, so as to have a fair comparison with the generated audio | |
| wav = Synthesizer_infer.load_preprocess_wav(fpath) | |
| self.ui.log("Loaded %s" % name) | |
| self.add_real_utterance(wav, name, speaker_name) | |
| def record(self): | |
| wav = self.ui.record_one(encoder_infer.sampling_rate, 5) | |
| if wav is None: | |
| return | |
| self.ui.play(wav, encoder_infer.sampling_rate) | |
| speaker_name = "user01" | |
| name = speaker_name + "_rec_%05d" % np.random.randint(100000) | |
| self.add_real_utterance(wav, name, speaker_name) | |
| def add_real_utterance(self, wav, name, speaker_name): | |
| # Compute the mel spectrogram | |
| spec = Synthesizer_infer.make_spectrogram(wav) | |
| self.ui.draw_spec(spec, "current") | |
| path_ori = os.getcwd() | |
| file_ori = 'temp.wav' | |
| fpath = os.path.join(path_ori, file_ori) | |
| sf.write(fpath, wav, samplerate=encoder.params_data.sampling_rate) | |
| # adjust the speed | |
| self.wav_ori_info = AudioAnalysis(path_ori, file_ori) | |
| DelFile(path_ori, '.TextGrid') | |
| os.remove(fpath) | |
| # Compute the embedding | |
| if not encoder_infer.is_loaded(): | |
| self.init_encoder() | |
| encoder_wav = encoder_infer.preprocess_wav(wav) | |
| embed, partial_embeds, _ = encoder_infer.embed_utterance(encoder_wav, return_partials=True) | |
| embed[embed < encoder.params_data.set_zero_thres]=0 # 噪声值置零 | |
| # Add the utterance | |
| utterance = Utterance(name, speaker_name, wav, spec, embed, partial_embeds, False) | |
| self.utterances.add(utterance) | |
| self.ui.register_utterance(utterance) | |
| # Plot it | |
| self.ui.draw_embed(embed, name, "current") | |
| self.ui.draw_umap_projections(self.utterances) | |
| self.ui.wav_ori_fig.savefig(f"toolbox_results/{name}_info.png", dpi=500) | |
| if len(self.utterances) >= self.ui.min_umap_points: | |
| self.ui.umap_fig.savefig(f"toolbox_results/umap_{len(self.utterances)}.png", dpi=500) | |
| def clear_utterances(self): | |
| self.utterances.clear() | |
| self.ui.draw_umap_projections(self.utterances) | |
| def synthesize(self): | |
| self.start_generate_time = time.time() | |
| self.ui.log("Generating the mel spectrogram...") | |
| self.ui.set_loading(1) | |
| # Update the synthesizer random seed | |
| if self.ui.random_seed_checkbox.isChecked(): | |
| seed = int(self.ui.seed_textbox.text()) | |
| self.ui.populate_gen_options(seed, self.trim_silences) | |
| else: | |
| seed = None | |
| if seed is not None: | |
| torch.manual_seed(seed) | |
| # Synthesize the spectrogram | |
| if self.synthesizer is None or seed is not None: | |
| self.init_synthesizer() | |
| embed = self.ui.selected_utterance.embed | |
| def preprocess_text(text): | |
| text = add_breaks(text) | |
| text = english_cleaners_predict(text) | |
| texts = [i.text.strip() for i in self.nlp(text).sents] # split paragraph to sentences | |
| return texts | |
| texts = preprocess_text(self.ui.text_prompt.toPlainText()) | |
| print(f"the list of inputs texts:\n{texts}") | |
| embeds = [embed] * len(texts) | |
| specs, alignments, stop_tokens = self.synthesizer.synthesize_spectrograms(texts, embeds, require_visualization=True) | |
| breaks = [spec.shape[1] for spec in specs] | |
| spec = np.concatenate(specs, axis=1) | |
| save_attention_multiple(alignments, "toolbox_results/attention") | |
| save_stop_tokens(stop_tokens, "toolbox_results/stop_tokens") | |
| self.ui.draw_spec(spec, "generated") | |
| self.current_generated = (self.ui.selected_utterance.speaker_name, spec, breaks, None) | |
| self.ui.set_loading(0) | |
| def vocode(self): | |
| speaker_name, spec, breaks, _ = self.current_generated | |
| assert spec is not None | |
| # Initialize the vocoder model and make it determinstic, if user provides a seed | |
| if self.ui.random_seed_checkbox.isChecked(): | |
| seed = int(self.ui.seed_textbox.text()) | |
| self.ui.populate_gen_options(seed, self.trim_silences) | |
| else: | |
| seed = None | |
| if seed is not None: | |
| torch.manual_seed(seed) | |
| # Synthesize the waveform | |
| if not vocoder.is_loaded() or seed is not None: | |
| self.init_vocoder() | |
| def vocoder_progress(i, seq_len, b_size, gen_rate): | |
| real_time_factor = (gen_rate / Synthesizer_infer.sample_rate) * 1000 | |
| line = "Waveform generation: %d/%d (batch size: %d, rate: %.1fkHz - %.2fx real time)" \ | |
| % (i * b_size, seq_len * b_size, b_size, gen_rate, real_time_factor) | |
| self.ui.log(line, "overwrite") | |
| self.ui.set_loading(i, seq_len) | |
| if self.ui.current_vocoder_fpath is not None and not self.ui.griffin_lim_checkbox.isChecked(): | |
| self.ui.log("") | |
| wav = vocoder.infer_waveform(spec, target=vocoder.hp.voc_target, overlap=vocoder.hp.voc_overlap, crossfade=vocoder.hp.is_crossfade, progress_callback=vocoder_progress) | |
| else: | |
| self.ui.log("Waveform generation with Griffin-Lim... ") | |
| wav = Synthesizer_infer.griffin_lim(spec) | |
| self.ui.set_loading(0) | |
| self.ui.log(" Done!", "append") | |
| self.ui.log(f"Generate time: {time.time() - self.start_generate_time}s") | |
| # Add breaks | |
| b_ends = np.cumsum(np.array(breaks) * Synthesizer_infer.hparams.hop_size) | |
| b_starts = np.concatenate(([0], b_ends[:-1])) | |
| wavs = [wav[start:end] for start, end, in zip(b_starts, b_ends)] | |
| breaks = [np.zeros(int(0.15 * Synthesizer_infer.sample_rate))] * len(breaks) | |
| wav = np.concatenate([i for w, b in zip(wavs, breaks) for i in (w, b)]) | |
| # Trim excessive silences | |
| if self.ui.trim_silences_checkbox.isChecked(): | |
| wav = encoder_infer.preprocess_wav(wav) | |
| path_ori = os.getcwd() | |
| file_ori = 'temp.wav' | |
| filename = os.path.join(path_ori, file_ori) | |
| sf.write(filename, wav.astype(np.float32), syn_hparams.sample_rate) | |
| self.ui.log("\nSaved output (haven't change speed) as %s\n\n" % filename) | |
| # Fix Speed(generate new audio) | |
| fix_file, speed_factor = work(*self.wav_ori_info, filename) | |
| self.ui.log(f"\nSaved output (fixed speed) as {fix_file}\n\n") | |
| wav, _ = librosa.load(fix_file, syn_hparams.sample_rate) | |
| os.remove(fix_file) | |
| # Play it | |
| wav = wav / np.abs(wav).max() * 4 | |
| self.ui.play(wav, Synthesizer_infer.sample_rate) | |
| # Name it (history displayed in combobox) | |
| # TODO better naming for the combobox items? | |
| wav_name = str(self.waves_count + 1) | |
| #Update waves combobox | |
| self.waves_count += 1 | |
| if self.waves_count > MAX_WAVS: | |
| self.waves_list.pop() | |
| self.waves_namelist.pop() | |
| self.waves_list.insert(0, wav) | |
| self.waves_namelist.insert(0, wav_name) | |
| self.ui.waves_cb.disconnect() | |
| self.ui.waves_cb_model.setStringList(self.waves_namelist) | |
| self.ui.waves_cb.setCurrentIndex(0) | |
| self.ui.waves_cb.currentIndexChanged.connect(self.set_current_wav) | |
| # Update current wav | |
| self.set_current_wav(0) | |
| #Enable replay and save buttons: | |
| self.ui.replay_wav_button.setDisabled(False) | |
| self.ui.export_wav_button.setDisabled(False) | |
| # Compute the embedding | |
| # TODO: this is problematic with different sampling rates, gotta fix it | |
| if not encoder_infer.is_loaded(): | |
| self.init_encoder() | |
| encoder_wav = encoder_infer.preprocess_wav(wav) | |
| embed, partial_embeds, _ = encoder_infer.embed_utterance(encoder_wav, return_partials=True) | |
| # Add the utterance | |
| name = speaker_name + "_gen_%05d_" % np.random.randint(100000) + str(speed_factor) | |
| utterance = Utterance(name, speaker_name, wav, spec, embed, partial_embeds, True) | |
| self.utterances.add(utterance) | |
| # Plot it | |
| self.ui.draw_embed(embed, name, "generated") | |
| self.ui.draw_umap_projections(self.utterances) | |
| self.ui.wav_gen_fig.savefig(f"toolbox_results/{name}_info.png", dpi=500) | |
| if len(self.utterances) >= self.ui.min_umap_points: | |
| self.ui.umap_fig.savefig(f"toolbox_results/umap_{len(self.utterances)}.png", dpi=500) | |
| def init_encoder(self): | |
| model_fpath = self.ui.current_encoder_fpath | |
| self.ui.log("Loading the encoder %s... " % model_fpath) | |
| self.ui.set_loading(1) | |
| start = timer() | |
| encoder_infer.load_model(model_fpath) | |
| self.ui.log("Done (%dms)." % int(1000 * (timer() - start)), "append") | |
| self.ui.set_loading(0) | |
| def init_synthesizer(self): | |
| model_fpath = self.ui.current_synthesizer_fpath | |
| self.ui.log("Loading the synthesizer %s... " % model_fpath) | |
| self.ui.set_loading(1) | |
| start = timer() | |
| self.synthesizer = Synthesizer_infer(model_fpath) | |
| self.ui.log("Done (%dms)." % int(1000 * (timer() - start)), "append") | |
| self.ui.set_loading(0) | |
| def init_vocoder(self): | |
| model_fpath = self.ui.current_vocoder_fpath | |
| # Case of Griffin-lim | |
| if model_fpath is None: | |
| return | |
| self.ui.log("Loading the vocoder %s... " % model_fpath) | |
| self.ui.set_loading(1) | |
| start = timer() | |
| vocoder.load_model(model_fpath) | |
| self.ui.log("Done (%dms)." % int(1000 * (timer() - start)), "append") | |
| self.ui.set_loading(0) | |
| def update_seed_textbox(self): | |
| self.ui.update_seed_textbox() | |