Upload toolbox/__init__.py with huggingface_hub

toolbox/__init__.py

@@ -0,0 +1,357 @@
+from toolbox.ui import UI
+from encoder import inference as encoder
+from synthesizer.inference import Synthesizer
+from vocoder import inference as vocoder
+from pathlib import Path
+from time import perf_counter as timer
+from toolbox.utterance import Utterance
+import numpy as np
+import traceback
+import sys
+import torch
+import librosa
+from audioread.exceptions import NoBackendError
+
+# 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_WAVES = 15
+
+class Toolbox:
+    def __init__(self, datasets_root, enc_models_dir, syn_models_dir, voc_models_dir, seed, no_mp3_support):
+        if not no_mp3_support:
+            try:
+                librosa.load("samples/6829_00000.mp3")
+            except NoBackendError:
+                print("Librosa will be unable to open mp3 files if additional software is not installed.\n"
+                  "Please install ffmpeg or add the '--no_mp3_support' option to proceed without support for mp3 files.")
+                exit(-1)
+        self.no_mp3_support = no_mp3_support
+        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
+        self.current_wav = None
+        self.waves_list = []
+        self.waves_count = 0
+        self.waves_namelist = []
+
+        # 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(enc_models_dir, syn_models_dir, voc_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.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.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.sample_rate)
+
+    def replay_last_wav(self):
+        self.ui.play(self.current_wav, Synthesizer.sample_rate)
+
+    def reset_ui(self, encoder_models_dir, synthesizer_models_dir, vocoder_models_dir, seed):
+        self.ui.populate_browser(self.datasets_root, recognized_datasets, 0, True)
+        self.ui.populate_models(encoder_models_dir, synthesizer_models_dir, vocoder_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
+
+        if fpath.suffix.lower() == ".mp3" and self.no_mp3_support:
+                self.ui.log("Error: No mp3 file argument was passed but an mp3 file was used")
+                return
+
+        # 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.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.sampling_rate, 5)
+        if wav is None:
+            return 
+        self.ui.play(wav, encoder.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.make_spectrogram(wav)
+        self.ui.draw_spec(spec, "current")
+
+        # Compute the embedding
+        if not encoder.is_loaded():
+            self.init_encoder()
+        encoder_wav = encoder.preprocess_wav(wav)
+        embed, partial_embeds, _ = encoder.embed_utterance(encoder_wav, return_partials=True)
+
+        # 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)
+        
+    def clear_utterances(self):
+        self.utterances.clear()
+        self.ui.draw_umap_projections(self.utterances)
+        
+    def synthesize(self):
+        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()
+
+        texts = self.ui.text_prompt.toPlainText().split("\n")
+        embed = self.ui.selected_utterance.embed
+        embeds = [embed] * len(texts)
+        specs = self.synthesizer.synthesize_spectrograms(texts, embeds)
+        breaks = [spec.shape[1] for spec in specs]
+        spec = np.concatenate(specs, axis=1)
+        
+        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.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:
+            self.ui.log("")
+            wav = vocoder.infer_waveform(spec, progress_callback=vocoder_progress)
+        else:
+            self.ui.log("Waveform generation with Griffin-Lim... ")
+            wav = Synthesizer.griffin_lim(spec)
+        self.ui.set_loading(0)
+        self.ui.log(" Done!", "append")
+        
+        # Add breaks
+        b_ends = np.cumsum(np.array(breaks) * Synthesizer.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.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.preprocess_wav(wav)
+
+        # Play it
+        wav = wav / np.abs(wav).max() * 0.97
+        self.ui.play(wav, Synthesizer.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_WAVES:
+          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.is_loaded():
+            self.init_encoder()
+        encoder_wav = encoder.preprocess_wav(wav)
+        embed, partial_embeds, _ = encoder.embed_utterance(encoder_wav, return_partials=True)
+        
+        # Add the utterance
+        name = speaker_name + "_gen_%05d" % np.random.randint(100000)
+        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)
+        
+    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.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(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()