Modernize: Gradio 6.29, Python 3.12, models from CorentinJ/SV2TTS, librosa/numpy/torch compat

README.md

@@ -3,9 +3,9 @@ title: Clone Your Voice
 emoji: 📚
 colorFrom: blue
 colorTo: yellow
-python_version: 3.8.4
+python_version: "3.12"
 sdk: gradio
-sdk_version: 3.0.4
+sdk_version: 6.29.0
 app_file: app.py
 pinned: false
 ---

app.py

@@ -1,380 +1,374 @@
-import gradio as gr
-import os
-from utils.default_models import ensure_default_models
-import sys
-import traceback
-from pathlib import Path
-from time import perf_counter as timer
-import numpy as np
-import torch
-from encoder import inference as encoder
-from synthesizer.inference import Synthesizer
-#from toolbox.utterance import Utterance
-from vocoder import inference as vocoder
-import time
-import librosa
-import numpy as np
-#import sounddevice as sd
-import soundfile as sf
-import argparse
-from utils.argutils import print_args
-
-parser = argparse.ArgumentParser(
-    formatter_class=argparse.ArgumentDefaultsHelpFormatter
-)
-parser.add_argument("-e", "--enc_model_fpath", type=Path,
-                    default="saved_models/default/encoder.pt",
-                    help="Path to a saved encoder")
-parser.add_argument("-s", "--syn_model_fpath", type=Path,
-                    default="saved_models/default/synthesizer.pt",
-                    help="Path to a saved synthesizer")
-parser.add_argument("-v", "--voc_model_fpath", type=Path,
-                    default="saved_models/default/vocoder.pt",
-                    help="Path to a saved vocoder")
-parser.add_argument("--cpu", action="store_true", help=\
-    "If True, processing is done on CPU, even when a GPU is available.")
-parser.add_argument("--no_sound", action="store_true", help=\
-    "If True, audio won't be played.")
-parser.add_argument("--seed", type=int, default=None, help=\
-    "Optional random number seed value to make toolbox deterministic.")
-args = parser.parse_args()
-arg_dict = vars(args)
-print_args(args, parser)
-
-# Maximum of generated wavs to keep on memory
-MAX_WAVS = 15
-utterances = set()
-current_generated = (None, None, None, None) # speaker_name, spec, breaks, wav
-synthesizer = None # type: Synthesizer
-current_wav = None
-waves_list = []
-waves_count = 0
-waves_namelist = []
-
-# Hide GPUs from Pytorch to force CPU processing
-if arg_dict.pop("cpu"):
-    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
-
-print("Running a test of your configuration...\n")
-
-if torch.cuda.is_available():
-    device_id = torch.cuda.current_device()
-    gpu_properties = torch.cuda.get_device_properties(device_id)
-    ## Print some environment information (for debugging purposes)
-    print("Found %d GPUs available. Using GPU %d (%s) of compute capability %d.%d with "
-        "%.1fGb total memory.\n" %
-        (torch.cuda.device_count(),
-        device_id,
-        gpu_properties.name,
-        gpu_properties.major,
-        gpu_properties.minor,
-        gpu_properties.total_memory / 1e9))
-else:
-    print("Using CPU for inference.\n")
-
-## Load the models one by one.
-print("Preparing the encoder, the synthesizer and the vocoder...")
-ensure_default_models(Path("saved_models"))
-#encoder.load_model(args.enc_model_fpath)
-#synthesizer = Synthesizer(args.syn_model_fpath)
-#vocoder.load_model(args.voc_model_fpath)
-
-def compute_embedding(in_fpath):
-
-    if not encoder.is_loaded():
-        model_fpath = args.enc_model_fpath
-        print("Loading the encoder %s... " % model_fpath)
-        start = time.time() 
-        encoder.load_model(model_fpath)
-        print("Done (%dms)." % int(1000 * (time.time() - start)), "append")
-
-
-    ## Computing the embedding
-    # First, we load the wav using the function that the speaker encoder provides. This is
-    
-    # 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
-    print("Step 1- load_preprocess_wav",in_fpath)
-    wav = Synthesizer.load_preprocess_wav(in_fpath)
-    
-    # important: there is preprocessing that must be applied.
-
-    # The following two methods are equivalent:
-    # - Directly load from the filepath:
-    print("Step 2- preprocess_wav")
-    preprocessed_wav = encoder.preprocess_wav(wav)
-
-    # - If the wav is already loaded:
-    #original_wav, sampling_rate = librosa.load(str(in_fpath))
-    #preprocessed_wav = encoder.preprocess_wav(original_wav, sampling_rate)
-
-    # Compute the embedding
-    print("Step 3- embed_utterance")
-    embed, partial_embeds, _ = encoder.embed_utterance(preprocessed_wav, return_partials=True)
-
-
-    print("Loaded file succesfully")
-
-    # Then we derive the embedding. There are many functions and parameters that the
-    # speaker encoder interfaces. These are mostly for in-depth research. You will typically
-    # only use this function (with its default parameters):
-    #embed = encoder.embed_utterance(preprocessed_wav)
-    
-    return embed 
-def create_spectrogram(text,embed):
-        # If seed is specified, reset torch seed and force synthesizer reload
-        if args.seed is not None:
-            torch.manual_seed(args.seed)
-            synthesizer = Synthesizer(args.syn_model_fpath)
-        
-        
-        # Synthesize the spectrogram
-        model_fpath = args.syn_model_fpath
-        print("Loading the synthesizer %s... " % model_fpath)
-        start = time.time()
-        synthesizer = Synthesizer(model_fpath)
-        print("Done (%dms)." % int(1000 * (time.time()- start)), "append")          
-        
-
-        # The synthesizer works in batch, so you need to put your data in a list or numpy array
-        texts = [text]
-        embeds = [embed]
-        # If you know what the attention layer alignments are, you can retrieve them here by
-        # passing return_alignments=True
-        specs = synthesizer.synthesize_spectrograms(texts, embeds)
-        breaks = [spec.shape[1] for spec in specs]
-        spec = np.concatenate(specs, axis=1)
-        sample_rate=synthesizer.sample_rate
-        return spec, breaks , sample_rate
-
-
-def generate_waveform(current_generated):
-
-        speaker_name, spec, breaks = current_generated
-        assert spec is not None
-
-        ## Generating the waveform
-        print("Synthesizing the waveform:")
-        # If seed is specified, reset torch seed and reload vocoder
-        if args.seed is not None:
-            torch.manual_seed(args.seed)
-            vocoder.load_model(args.voc_model_fpath)
-
-        model_fpath = args.voc_model_fpath
-        # Synthesize the waveform
-        if not vocoder.is_loaded():
-            print("Loading the vocoder %s... " % model_fpath)
-            start = time.time()
-            vocoder.load_model(model_fpath)
-            print("Done (%dms)." % int(1000 * (time.time()- start)), "append")    
-
-        current_vocoder_fpath= model_fpath
-        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)
-            print(line, "overwrite")       
-
-
-        # Synthesizing the waveform is fairly straightforward. Remember that the longer the
-        # spectrogram, the more time-efficient the vocoder.
-        if  current_vocoder_fpath is not None:
-            print("")
-            generated_wav = vocoder.infer_waveform(spec, progress_callback=vocoder_progress)
-        else:
-            print("Waveform generation with Griffin-Lim... ")
-            generated_wav = Synthesizer.griffin_lim(spec)
-
-        print(" Done!", "append")
-
-
-        ## Post-generation
-        # There's a bug with sounddevice that makes the audio cut one second earlier, so we
-        # pad it.
-        generated_wav = np.pad(generated_wav, (0, Synthesizer.sample_rate), mode="constant")
-
-        # Add breaks
-        b_ends = np.cumsum(np.array(breaks) * Synthesizer.hparams.hop_size)
-        b_starts = np.concatenate(([0], b_ends[:-1]))
-        wavs = [generated_wav[start:end] for start, end, in zip(b_starts, b_ends)]
-        breaks = [np.zeros(int(0.15 * Synthesizer.sample_rate))] * len(breaks)
-        generated_wav = np.concatenate([i for w, b in zip(wavs, breaks) for i in (w, b)])
-
-
-        # Trim excess silences to compensate for gaps in spectrograms (issue #53)
-        generated_wav = encoder.preprocess_wav(generated_wav)
-
-
-        return generated_wav
-
-
-def save_on_disk(generated_wav,sample_rate):
-        # Save it on the disk
-        filename = "cloned_voice.wav"
-        print(generated_wav.dtype)
-        #OUT=os.environ['OUT_PATH']
-        # Returns `None` if key doesn't exist
-        #OUT=os.environ.get('OUT_PATH')
-        #result = os.path.join(OUT, filename)
-        result = filename
-        print(" > Saving output to {}".format(result))
-        sf.write(result, generated_wav.astype(np.float32), sample_rate)
-        print("\nSaved output as %s\n\n" % result) 
-      
-        return  result     
-def play_audio(generated_wav,sample_rate):
-        # Play the audio (non-blocking)
-        if not args.no_sound:
-          
-            try:
-                sd.stop()
-                sd.play(generated_wav, sample_rate)
-            except sd.PortAudioError as e:
-                print("\nCaught exception: %s" % repr(e))
-                print("Continuing without audio playback. Suppress this message with the \"--no_sound\" flag.\n")
-            except:
-                raise
-         
-
-def clean_memory():
-    import gc
-    #import GPUtil
-    # To see memory usage
-    print('Before clean ')
-    #GPUtil.showUtilization()
-    #cleaning memory 1
-    gc.collect()
-    torch.cuda.empty_cache()
-    time.sleep(2)
-    print('After Clean GPU')
-    #GPUtil.showUtilization()
-
-def clone_voice(in_fpath, text):
-    try:       
-            speaker_name = "output"
-            # Compute embedding
-            embed=compute_embedding(in_fpath)
-            print("Created the embedding")
-            # Generating the spectrogram
-            spec, breaks, sample_rate = create_spectrogram(text,embed)
-            current_generated = (speaker_name, spec, breaks)
-            print("Created the mel spectrogram")
-
-            # Create waveform
-            generated_wav=generate_waveform(current_generated)
-            print("Created the the waveform ")
-
-            # Save it on the disk
-            save_on_disk(generated_wav,sample_rate)
-
-            #Play the audio 
-            #play_audio(generated_wav,sample_rate)
-
-            return        
-    except Exception as e:
-        print("Caught exception: %s" % repr(e))
-        print("Restarting\n")
-
-# Set environment variables
-home_dir = os.getcwd()
-OUT_PATH=os.path.join(home_dir, "out/")
-os.environ['OUT_PATH'] = OUT_PATH
-
-# create output path
-os.makedirs(OUT_PATH, exist_ok=True)
-
-USE_CUDA = torch.cuda.is_available()  
-
-os.system('pip install -q pydub ffmpeg-normalize')
-CONFIG_SE_PATH = "config_se.json"
-CHECKPOINT_SE_PATH = "SE_checkpoint.pth.tar"
-def greet(Text,Voicetoclone ,input_mic=None):
-    text= "%s" % (Text)
-    #reference_files= "%s" % (Voicetoclone)
-
-    clean_memory()
-    print(text,len(text),type(text))
-    print(Voicetoclone,type(Voicetoclone))
-
-    if  len(text) == 0 : 
-        print("Please add text to the program")
-        Text="Please add text to the program, thank you."
-        is_no_text=True
-    else:
-        is_no_text=False
-
-    
-    if Voicetoclone==None and input_mic==None:
-        print("There is no input audio")
-        Text="Please add audio input, to the program, thank you."
-        Voicetoclone='trump.mp3'
-        if  is_no_text:
-            Text="Please add text and audio, to the program, thank you."
-
-    if  input_mic != "" and input_mic != None :
-        # Get the wav file from the microphone
-        print('The value of MIC IS :',input_mic,type(input_mic))
-        Voicetoclone= input_mic
-
-    text= "%s" % (Text)
-    reference_files= Voicetoclone
-    print("path url")
-    print(Voicetoclone)
-    sample= str(Voicetoclone)
-    os.environ['sample'] = sample
-    size= len(reference_files)*sys.getsizeof(reference_files)
-    size2= size / 1000000
-    if (size2 > 0.012) or len(text)>2000:
-      message="File is greater than 30mb or Text inserted is longer than 2000 characters. Please re-try with smaller sizes."
-      print(message)
-      raise SystemExit("File is greater than 30mb. Please re-try or Text inserted is longer than 2000 characters. Please re-try with smaller sizes.")
-    else:
-
-      env_var = 'sample'
-      if env_var in os.environ:
-            print(f'{env_var} value is {os.environ[env_var]}')
-      else:
-            print(f'{env_var} does not exist')
-      #os.system(f'ffmpeg-normalize {os.environ[env_var]} -nt rms -t=-27 -o {os.environ[env_var]} -ar 16000 -f')
-      in_fpath = Path(Voicetoclone)
-      #in_fpath= in_fpath.replace("\"", "").replace("\'", "")
-      
-      out_path=clone_voice(in_fpath, text)
-
-      print(" > text: {}".format(text))
-
-      print("Generated Audio")
-      return "cloned_voice.wav"
-
-demo = gr.Interface(
-    fn=greet, 
-    inputs=[gr.inputs.Textbox(label='What would you like the voice to say? (max. 2000 characters per request)'),
-            gr.Audio(
-            type="filepath",         
-            source="upload",
-            label='Please upload a voice to clone (max. 30mb)'),
-            gr.inputs.Audio(
-            source="microphone", 
-            label='or record',
-            type="filepath", 
-            optional=True)
-            ],
-    outputs="audio",
-
-    title = 'Clone Your Voice',
-            description = 'A simple application that Clone Your Voice.  Wait one minute to process.',
-            article = 
-                        '''<div>
-                            <p style="text-align: center"> All you need to do is record your voice, type what you want be say
-                            ,then wait for compiling. After that click on Play/Pause for listen the audio. The audio is saved in an wav format.
-                            For more information visit <a href="https://ruslanmv.com/">ruslanmv.com</a>
-                            </p>
-                        </div>''',
-
-           examples = [["I am the cloned version of Donald Trump. Well. I think what's happening to this country is unbelievably bad. We're no longer a respected country","trump.mp3","trump.mp3"],
-                        ["I am the cloned version of Elon Musk. Persistence is very important. You should not give up unless you are forced to give up.","musk.mp3","musk.mp3"] #,
-                       # ["I am the cloned version of Elizabeth. It has always been easy to hate and destroy. To build and to cherish is much more difficult." ,"queen.mp3","queen.mp3"]                    
-                      ]      
-    
-    )
+import gradio as gr
+import os
+from utils.default_models import ensure_default_models
+import sys
+import traceback
+from pathlib import Path
+from time import perf_counter as timer
+import numpy as np
+import torch
+from encoder import inference as encoder
+from synthesizer.inference import Synthesizer
+#from toolbox.utterance import Utterance
+from vocoder import inference as vocoder
+import time
+import librosa
+import numpy as np
+#import sounddevice as sd
+import soundfile as sf
+import argparse
+from utils.argutils import print_args
+
+parser = argparse.ArgumentParser(
+    formatter_class=argparse.ArgumentDefaultsHelpFormatter
+)
+parser.add_argument("-e", "--enc_model_fpath", type=Path,
+                    default="saved_models/default/encoder.pt",
+                    help="Path to a saved encoder")
+parser.add_argument("-s", "--syn_model_fpath", type=Path,
+                    default="saved_models/default/synthesizer.pt",
+                    help="Path to a saved synthesizer")
+parser.add_argument("-v", "--voc_model_fpath", type=Path,
+                    default="saved_models/default/vocoder.pt",
+                    help="Path to a saved vocoder")
+parser.add_argument("--cpu", action="store_true", help=\
+    "If True, processing is done on CPU, even when a GPU is available.")
+parser.add_argument("--no_sound", action="store_true", help=\
+    "If True, audio won't be played.")
+parser.add_argument("--seed", type=int, default=None, help=\
+    "Optional random number seed value to make toolbox deterministic.")
+args, _unknown = parser.parse_known_args()
+arg_dict = vars(args)
+print_args(args, parser)
+
+# Maximum of generated wavs to keep on memory
+MAX_WAVS = 15
+utterances = set()
+current_generated = (None, None, None, None) # speaker_name, spec, breaks, wav
+synthesizer = None # type: Synthesizer
+current_wav = None
+waves_list = []
+waves_count = 0
+waves_namelist = []
+
+# Hide GPUs from Pytorch to force CPU processing
+if arg_dict.pop("cpu"):
+    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+print("Running a test of your configuration...\n")
+
+if torch.cuda.is_available():
+    device_id = torch.cuda.current_device()
+    gpu_properties = torch.cuda.get_device_properties(device_id)
+    ## Print some environment information (for debugging purposes)
+    print("Found %d GPUs available. Using GPU %d (%s) of compute capability %d.%d with "
+        "%.1fGb total memory.\n" %
+        (torch.cuda.device_count(),
+        device_id,
+        gpu_properties.name,
+        gpu_properties.major,
+        gpu_properties.minor,
+        gpu_properties.total_memory / 1e9))
+else:
+    print("Using CPU for inference.\n")
+
+## Load the models one by one.
+print("Preparing the encoder, the synthesizer and the vocoder...")
+ensure_default_models(Path("saved_models"))
+#encoder.load_model(args.enc_model_fpath)
+#synthesizer = Synthesizer(args.syn_model_fpath)
+#vocoder.load_model(args.voc_model_fpath)
+
+def compute_embedding(in_fpath):
+
+    if not encoder.is_loaded():
+        model_fpath = args.enc_model_fpath
+        print("Loading the encoder %s... " % model_fpath)
+        start = time.time() 
+        encoder.load_model(model_fpath)
+        print("Done (%dms)." % int(1000 * (time.time() - start)), "append")
+
+
+    ## Computing the embedding
+    # First, we load the wav using the function that the speaker encoder provides. This is
+    
+    # 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
+    print("Step 1- load_preprocess_wav",in_fpath)
+    wav = Synthesizer.load_preprocess_wav(in_fpath)
+    
+    # important: there is preprocessing that must be applied.
+
+    # The following two methods are equivalent:
+    # - Directly load from the filepath:
+    print("Step 2- preprocess_wav")
+    preprocessed_wav = encoder.preprocess_wav(wav)
+
+    # - If the wav is already loaded:
+    #original_wav, sampling_rate = librosa.load(str(in_fpath))
+    #preprocessed_wav = encoder.preprocess_wav(original_wav, sampling_rate)
+
+    # Compute the embedding
+    print("Step 3- embed_utterance")
+    embed, partial_embeds, _ = encoder.embed_utterance(preprocessed_wav, return_partials=True)
+
+
+    print("Loaded file succesfully")
+
+    # Then we derive the embedding. There are many functions and parameters that the
+    # speaker encoder interfaces. These are mostly for in-depth research. You will typically
+    # only use this function (with its default parameters):
+    #embed = encoder.embed_utterance(preprocessed_wav)
+    
+    return embed 
+def create_spectrogram(text,embed):
+        # If seed is specified, reset torch seed and force synthesizer reload
+        if args.seed is not None:
+            torch.manual_seed(args.seed)
+            synthesizer = Synthesizer(args.syn_model_fpath)
+        
+        
+        # Synthesize the spectrogram
+        model_fpath = args.syn_model_fpath
+        print("Loading the synthesizer %s... " % model_fpath)
+        start = time.time()
+        synthesizer = Synthesizer(model_fpath)
+        print("Done (%dms)." % int(1000 * (time.time()- start)), "append")          
+        
+
+        # The synthesizer works in batch, so you need to put your data in a list or numpy array
+        texts = [text]
+        embeds = [embed]
+        # If you know what the attention layer alignments are, you can retrieve them here by
+        # passing return_alignments=True
+        specs = synthesizer.synthesize_spectrograms(texts, embeds)
+        breaks = [spec.shape[1] for spec in specs]
+        spec = np.concatenate(specs, axis=1)
+        sample_rate=synthesizer.sample_rate
+        return spec, breaks , sample_rate
+
+
+def generate_waveform(current_generated):
+
+        speaker_name, spec, breaks = current_generated
+        assert spec is not None
+
+        ## Generating the waveform
+        print("Synthesizing the waveform:")
+        # If seed is specified, reset torch seed and reload vocoder
+        if args.seed is not None:
+            torch.manual_seed(args.seed)
+            vocoder.load_model(args.voc_model_fpath)
+
+        model_fpath = args.voc_model_fpath
+        # Synthesize the waveform
+        if not vocoder.is_loaded():
+            print("Loading the vocoder %s... " % model_fpath)
+            start = time.time()
+            vocoder.load_model(model_fpath)
+            print("Done (%dms)." % int(1000 * (time.time()- start)), "append")    
+
+        current_vocoder_fpath= model_fpath
+        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)
+            print(line, "overwrite")       
+
+
+        # Synthesizing the waveform is fairly straightforward. Remember that the longer the
+        # spectrogram, the more time-efficient the vocoder.
+        if  current_vocoder_fpath is not None:
+            print("")
+            generated_wav = vocoder.infer_waveform(spec, progress_callback=vocoder_progress)
+        else:
+            print("Waveform generation with Griffin-Lim... ")
+            generated_wav = Synthesizer.griffin_lim(spec)
+
+        print(" Done!", "append")
+
+
+        ## Post-generation
+        # There's a bug with sounddevice that makes the audio cut one second earlier, so we
+        # pad it.
+        generated_wav = np.pad(generated_wav, (0, Synthesizer.sample_rate), mode="constant")
+
+        # Add breaks
+        b_ends = np.cumsum(np.array(breaks) * Synthesizer.hparams.hop_size)
+        b_starts = np.concatenate(([0], b_ends[:-1]))
+        wavs = [generated_wav[start:end] for start, end, in zip(b_starts, b_ends)]
+        breaks = [np.zeros(int(0.15 * Synthesizer.sample_rate))] * len(breaks)
+        generated_wav = np.concatenate([i for w, b in zip(wavs, breaks) for i in (w, b)])
+
+
+        # Trim excess silences to compensate for gaps in spectrograms (issue #53)
+        generated_wav = encoder.preprocess_wav(generated_wav)
+
+
+        return generated_wav
+
+
+def save_on_disk(generated_wav,sample_rate):
+        # Save it on the disk
+        import tempfile
+        filename = tempfile.NamedTemporaryFile(suffix="_cloned_voice.wav", delete=False).name
+        print(generated_wav.dtype)
+        #OUT=os.environ['OUT_PATH']
+        # Returns `None` if key doesn't exist
+        #OUT=os.environ.get('OUT_PATH')
+        #result = os.path.join(OUT, filename)
+        result = filename
+        print(" > Saving output to {}".format(result))
+        sf.write(result, generated_wav.astype(np.float32), sample_rate)
+        print("\nSaved output as %s\n\n" % result) 
+      
+        return  result     
+def play_audio(generated_wav,sample_rate):
+        # Play the audio (non-blocking)
+        if not args.no_sound:
+          
+            try:
+                sd.stop()
+                sd.play(generated_wav, sample_rate)
+            except sd.PortAudioError as e:
+                print("\nCaught exception: %s" % repr(e))
+                print("Continuing without audio playback. Suppress this message with the \"--no_sound\" flag.\n")
+            except:
+                raise
+         
+
+def clean_memory():
+    import gc
+    #import GPUtil
+    # To see memory usage
+    print('Before clean ')
+    #GPUtil.showUtilization()
+    #cleaning memory 1
+    gc.collect()
+    torch.cuda.empty_cache()
+    time.sleep(2)
+    print('After Clean GPU')
+    #GPUtil.showUtilization()
+
+def clone_voice(in_fpath, text):
+    try:       
+            speaker_name = "output"
+            # Compute embedding
+            embed=compute_embedding(in_fpath)
+            print("Created the embedding")
+            # Generating the spectrogram
+            spec, breaks, sample_rate = create_spectrogram(text,embed)
+            current_generated = (speaker_name, spec, breaks)
+            print("Created the mel spectrogram")
+
+            # Create waveform
+            generated_wav=generate_waveform(current_generated)
+            print("Created the the waveform ")
+
+            # Save it on the disk
+            return save_on_disk(generated_wav,sample_rate)
+    except Exception as e:
+        print("Caught exception: %s" % repr(e))
+        raise gr.Error("Voice cloning failed: %s" % e)
+
+# Set environment variables
+home_dir = os.getcwd()
+OUT_PATH=os.path.join(home_dir, "out/")
+os.environ['OUT_PATH'] = OUT_PATH
+
+# create output path
+os.makedirs(OUT_PATH, exist_ok=True)
+
+USE_CUDA = torch.cuda.is_available()  
+
+CONFIG_SE_PATH = "config_se.json"
+CHECKPOINT_SE_PATH = "SE_checkpoint.pth.tar"
+def greet(Text,Voicetoclone ,input_mic=None):
+    text= "%s" % (Text)
+    #reference_files= "%s" % (Voicetoclone)
+
+    clean_memory()
+    print(text,len(text),type(text))
+    print(Voicetoclone,type(Voicetoclone))
+
+    if  len(text) == 0 : 
+        print("Please add text to the program")
+        Text="Please add text to the program, thank you."
+        is_no_text=True
+    else:
+        is_no_text=False
+
+    
+    if Voicetoclone==None and input_mic==None:
+        print("There is no input audio")
+        Text="Please add audio input, to the program, thank you."
+        Voicetoclone='trump.mp3'
+        if  is_no_text:
+            Text="Please add text and audio, to the program, thank you."
+
+    if  input_mic != "" and input_mic != None :
+        # Get the wav file from the microphone
+        print('The value of MIC IS :',input_mic,type(input_mic))
+        Voicetoclone= input_mic
+
+    text= "%s" % (Text)
+    reference_files= Voicetoclone
+    print("path url")
+    print(Voicetoclone)
+    sample= str(Voicetoclone)
+    os.environ['sample'] = sample
+    size2= os.path.getsize(str(reference_files)) / 1000000 if os.path.exists(str(reference_files)) else 0
+    if (size2 > 30) or len(text)>2000:
+      message="File is greater than 30mb or Text inserted is longer than 2000 characters. Please re-try with smaller sizes."
+      print(message)
+      raise gr.Error(message)
+    else:
+
+      env_var = 'sample'
+      if env_var in os.environ:
+            print(f'{env_var} value is {os.environ[env_var]}')
+      else:
+            print(f'{env_var} does not exist')
+      #os.system(f'ffmpeg-normalize {os.environ[env_var]} -nt rms -t=-27 -o {os.environ[env_var]} -ar 16000 -f')
+      in_fpath = Path(Voicetoclone)
+      #in_fpath= in_fpath.replace("\"", "").replace("\'", "")
+      
+      out_path=clone_voice(in_fpath, text)
+
+      print(" > text: {}".format(text))
+
+      print("Generated Audio")
+      return out_path
+
+demo = gr.Interface(
+    fn=greet, 
+    inputs=[gr.Textbox(label='What would you like the voice to say? (max. 2000 characters per request)'),
+            gr.Audio(
+            type="filepath",
+            sources=["upload"],
+            label='Please upload a voice to clone (max. 30mb)'),
+            gr.Audio(
+            sources=["microphone"],
+            label='or record',
+            type="filepath")
+            ],
+    outputs=gr.Audio(type="filepath"),
+    cache_examples=False,
+
+    title = 'Clone Your Voice',
+            description = 'A simple application that Clone Your Voice.  Wait one minute to process.',
+            article = 
+                        '''<div>
+                            <p style="text-align: center"> All you need to do is record your voice, type what you want be say
+                            ,then wait for compiling. After that click on Play/Pause for listen the audio. The audio is saved in an wav format.
+                            For more information visit <a href="https://ruslanmv.com/">ruslanmv.com</a>
+                            </p>
+                        </div>''',
+
+           examples = [["I am the cloned version of Donald Trump. Well. I think what's happening to this country is unbelievably bad. We're no longer a respected country","trump.mp3","trump.mp3"],
+                        ["I am the cloned version of Elon Musk. Persistence is very important. You should not give up unless you are forced to give up.","musk.mp3","musk.mp3"] #,
+                       # ["I am the cloned version of Elizabeth. It has always been easy to hate and destroy. To build and to cherish is much more difficult." ,"queen.mp3","queen.mp3"]                    
+                      ]      
+    
+    )
 demo.launch()

encoder/audio.py

@@ -1,117 +1,117 @@
-from scipy.ndimage.morphology import binary_dilation
-from encoder.params_data import *
-from pathlib import Path
-from typing import Optional, Union
-from warnings import warn
-import numpy as np
-import librosa
-import struct
-
-try:
-    import webrtcvad
-except:
-    warn("Unable to import 'webrtcvad'. This package enables noise removal and is recommended.")
-    webrtcvad=None
-
-int16_max = (2 ** 15) - 1
-
-
-def preprocess_wav(fpath_or_wav: Union[str, Path, np.ndarray],
-                   source_sr: Optional[int] = None,
-                   normalize: Optional[bool] = True,
-                   trim_silence: Optional[bool] = True):
-    """
-    Applies the preprocessing operations used in training the Speaker Encoder to a waveform 
-    either on disk or in memory. The waveform will be resampled to match the data hyperparameters.
-
-    :param fpath_or_wav: either a filepath to an audio file (many extensions are supported, not 
-    just .wav), either the waveform as a numpy array of floats.
-    :param source_sr: if passing an audio waveform, the sampling rate of the waveform before 
-    preprocessing. After preprocessing, the waveform's sampling rate will match the data 
-    hyperparameters. If passing a filepath, the sampling rate will be automatically detected and 
-    this argument will be ignored.
-    """
-    # Load the wav from disk if needed
-    if isinstance(fpath_or_wav, str) or isinstance(fpath_or_wav, Path):
-        wav, source_sr = librosa.load(str(fpath_or_wav), sr=None)
-    else:
-        wav = fpath_or_wav
-    
-    # Resample the wav if needed
-    if source_sr is not None and source_sr != sampling_rate:
-        wav = librosa.resample(wav, source_sr, sampling_rate)
-        
-    # Apply the preprocessing: normalize volume and shorten long silences 
-    if normalize:
-        wav = normalize_volume(wav, audio_norm_target_dBFS, increase_only=True)
-    if webrtcvad and trim_silence:
-        wav = trim_long_silences(wav)
-    
-    return wav
-
-
-def wav_to_mel_spectrogram(wav):
-    """
-    Derives a mel spectrogram ready to be used by the encoder from a preprocessed audio waveform.
-    Note: this not a log-mel spectrogram.
-    """
-    frames = librosa.feature.melspectrogram(
-        wav,
-        sampling_rate,
-        n_fft=int(sampling_rate * mel_window_length / 1000),
-        hop_length=int(sampling_rate * mel_window_step / 1000),
-        n_mels=mel_n_channels
-    )
-    return frames.astype(np.float32).T
-
-
-def trim_long_silences(wav):
-    """
-    Ensures that segments without voice in the waveform remain no longer than a 
-    threshold determined by the VAD parameters in params.py.
-
-    :param wav: the raw waveform as a numpy array of floats 
-    :return: the same waveform with silences trimmed away (length <= original wav length)
-    """
-    # Compute the voice detection window size
-    samples_per_window = (vad_window_length * sampling_rate) // 1000
-    
-    # Trim the end of the audio to have a multiple of the window size
-    wav = wav[:len(wav) - (len(wav) % samples_per_window)]
-    
-    # Convert the float waveform to 16-bit mono PCM
-    pcm_wave = struct.pack("%dh" % len(wav), *(np.round(wav * int16_max)).astype(np.int16))
-    
-    # Perform voice activation detection
-    voice_flags = []
-    vad = webrtcvad.Vad(mode=3)
-    for window_start in range(0, len(wav), samples_per_window):
-        window_end = window_start + samples_per_window
-        voice_flags.append(vad.is_speech(pcm_wave[window_start * 2:window_end * 2],
-                                         sample_rate=sampling_rate))
-    voice_flags = np.array(voice_flags)
-    
-    # Smooth the voice detection with a moving average
-    def moving_average(array, width):
-        array_padded = np.concatenate((np.zeros((width - 1) // 2), array, np.zeros(width // 2)))
-        ret = np.cumsum(array_padded, dtype=float)
-        ret[width:] = ret[width:] - ret[:-width]
-        return ret[width - 1:] / width
-    
-    audio_mask = moving_average(voice_flags, vad_moving_average_width)
-    audio_mask = np.round(audio_mask).astype(np.bool)
-    
-    # Dilate the voiced regions
-    audio_mask = binary_dilation(audio_mask, np.ones(vad_max_silence_length + 1))
-    audio_mask = np.repeat(audio_mask, samples_per_window)
-    
-    return wav[audio_mask == True]
-
-
-def normalize_volume(wav, target_dBFS, increase_only=False, decrease_only=False):
-    if increase_only and decrease_only:
-        raise ValueError("Both increase only and decrease only are set")
-    dBFS_change = target_dBFS - 10 * np.log10(np.mean(wav ** 2))
-    if (dBFS_change < 0 and increase_only) or (dBFS_change > 0 and decrease_only):
-        return wav
-    return wav * (10 ** (dBFS_change / 20))
+from scipy.ndimage import binary_dilation
+from encoder.params_data import *
+from pathlib import Path
+from typing import Optional, Union
+from warnings import warn
+import numpy as np
+import librosa
+import struct
+
+try:
+    import webrtcvad
+except:
+    warn("Unable to import 'webrtcvad'. This package enables noise removal and is recommended.")
+    webrtcvad=None
+
+int16_max = (2 ** 15) - 1
+
+
+def preprocess_wav(fpath_or_wav: Union[str, Path, np.ndarray],
+                   source_sr: Optional[int] = None,
+                   normalize: Optional[bool] = True,
+                   trim_silence: Optional[bool] = True):
+    """
+    Applies the preprocessing operations used in training the Speaker Encoder to a waveform 
+    either on disk or in memory. The waveform will be resampled to match the data hyperparameters.
+
+    :param fpath_or_wav: either a filepath to an audio file (many extensions are supported, not 
+    just .wav), either the waveform as a numpy array of floats.
+    :param source_sr: if passing an audio waveform, the sampling rate of the waveform before 
+    preprocessing. After preprocessing, the waveform's sampling rate will match the data 
+    hyperparameters. If passing a filepath, the sampling rate will be automatically detected and 
+    this argument will be ignored.
+    """
+    # Load the wav from disk if needed
+    if isinstance(fpath_or_wav, str) or isinstance(fpath_or_wav, Path):
+        wav, source_sr = librosa.load(str(fpath_or_wav), sr=None)
+    else:
+        wav = fpath_or_wav
+    
+    # Resample the wav if needed
+    if source_sr is not None and source_sr != sampling_rate:
+        wav = librosa.resample(wav, orig_sr=source_sr, target_sr=sampling_rate)
+        
+    # Apply the preprocessing: normalize volume and shorten long silences 
+    if normalize:
+        wav = normalize_volume(wav, audio_norm_target_dBFS, increase_only=True)
+    if webrtcvad and trim_silence:
+        wav = trim_long_silences(wav)
+    
+    return wav
+
+
+def wav_to_mel_spectrogram(wav):
+    """
+    Derives a mel spectrogram ready to be used by the encoder from a preprocessed audio waveform.
+    Note: this not a log-mel spectrogram.
+    """
+    frames = librosa.feature.melspectrogram(
+        y=wav,
+        sr=sampling_rate,
+        n_fft=int(sampling_rate * mel_window_length / 1000),
+        hop_length=int(sampling_rate * mel_window_step / 1000),
+        n_mels=mel_n_channels
+    )
+    return frames.astype(np.float32).T
+
+
+def trim_long_silences(wav):
+    """
+    Ensures that segments without voice in the waveform remain no longer than a 
+    threshold determined by the VAD parameters in params.py.
+
+    :param wav: the raw waveform as a numpy array of floats 
+    :return: the same waveform with silences trimmed away (length <= original wav length)
+    """
+    # Compute the voice detection window size
+    samples_per_window = (vad_window_length * sampling_rate) // 1000
+    
+    # Trim the end of the audio to have a multiple of the window size
+    wav = wav[:len(wav) - (len(wav) % samples_per_window)]
+    
+    # Convert the float waveform to 16-bit mono PCM
+    pcm_wave = struct.pack("%dh" % len(wav), *(np.round(wav * int16_max)).astype(np.int16))
+    
+    # Perform voice activation detection
+    voice_flags = []
+    vad = webrtcvad.Vad(mode=3)
+    for window_start in range(0, len(wav), samples_per_window):
+        window_end = window_start + samples_per_window
+        voice_flags.append(vad.is_speech(pcm_wave[window_start * 2:window_end * 2],
+                                         sample_rate=sampling_rate))
+    voice_flags = np.array(voice_flags)
+    
+    # Smooth the voice detection with a moving average
+    def moving_average(array, width):
+        array_padded = np.concatenate((np.zeros((width - 1) // 2), array, np.zeros(width // 2)))
+        ret = np.cumsum(array_padded, dtype=float)
+        ret[width:] = ret[width:] - ret[:-width]
+        return ret[width - 1:] / width
+    
+    audio_mask = moving_average(voice_flags, vad_moving_average_width)
+    audio_mask = np.round(audio_mask).astype(bool)
+    
+    # Dilate the voiced regions
+    audio_mask = binary_dilation(audio_mask, np.ones(vad_max_silence_length + 1))
+    audio_mask = np.repeat(audio_mask, samples_per_window)
+    
+    return wav[audio_mask == True]
+
+
+def normalize_volume(wav, target_dBFS, increase_only=False, decrease_only=False):
+    if increase_only and decrease_only:
+        raise ValueError("Both increase only and decrease only are set")
+    dBFS_change = target_dBFS - 10 * np.log10(np.mean(wav ** 2))
+    if (dBFS_change < 0 and increase_only) or (dBFS_change > 0 and decrease_only):
+        return wav
+    return wav * (10 ** (dBFS_change / 20))

encoder/inference.py

@@ -1,178 +1,178 @@
-from encoder.params_data import *
-from encoder.model import SpeakerEncoder
-from encoder.audio import preprocess_wav   # We want to expose this function from here
-from matplotlib import cm
-from encoder import audio
-from pathlib import Path
-import numpy as np
-import torch
-
-_model = None # type: SpeakerEncoder
-_device = None # type: torch.device
-
-
-def load_model(weights_fpath: Path, device=None):
-    """
-    Loads the model in memory. If this function is not explicitely called, it will be run on the
-    first call to embed_frames() with the default weights file.
-
-    :param weights_fpath: the path to saved model weights.
-    :param device: either a torch device or the name of a torch device (e.g. "cpu", "cuda"). The
-    model will be loaded and will run on this device. Outputs will however always be on the cpu.
-    If None, will default to your GPU if it"s available, otherwise your CPU.
-    """
-    # TODO: I think the slow loading of the encoder might have something to do with the device it
-    #   was saved on. Worth investigating.
-    global _model, _device
-    if device is None:
-        _device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    elif isinstance(device, str):
-        _device = torch.device(device)
-    _model = SpeakerEncoder(_device, torch.device("cpu"))
-    checkpoint = torch.load(weights_fpath, _device)
-    _model.load_state_dict(checkpoint["model_state"])
-    _model.eval()
-    print("Loaded encoder \"%s\" trained to step %d" % (weights_fpath.name, checkpoint["step"]))
-
-
-def is_loaded():
-    return _model is not None
-
-
-def embed_frames_batch(frames_batch):
-    """
-    Computes embeddings for a batch of mel spectrogram.
-
-    :param frames_batch: a batch mel of spectrogram as a numpy array of float32 of shape
-    (batch_size, n_frames, n_channels)
-    :return: the embeddings as a numpy array of float32 of shape (batch_size, model_embedding_size)
-    """
-    if _model is None:
-        raise Exception("Model was not loaded. Call load_model() before inference.")
-
-    frames = torch.from_numpy(frames_batch).to(_device)
-    embed = _model.forward(frames).detach().cpu().numpy()
-    return embed
-
-
-def compute_partial_slices(n_samples, partial_utterance_n_frames=partials_n_frames,
-                           min_pad_coverage=0.75, overlap=0.5):
-    """
-    Computes where to split an utterance waveform and its corresponding mel spectrogram to obtain
-    partial utterances of <partial_utterance_n_frames> each. Both the waveform and the mel
-    spectrogram slices are returned, so as to make each partial utterance waveform correspond to
-    its spectrogram. This function assumes that the mel spectrogram parameters used are those
-    defined in params_data.py.
-
-    The returned ranges may be indexing further than the length of the waveform. It is
-    recommended that you pad the waveform with zeros up to wave_slices[-1].stop.
-
-    :param n_samples: the number of samples in the waveform
-    :param partial_utterance_n_frames: the number of mel spectrogram frames in each partial
-    utterance
-    :param min_pad_coverage: when reaching the last partial utterance, it may or may not have
-    enough frames. If at least <min_pad_coverage> of <partial_utterance_n_frames> are present,
-    then the last partial utterance will be considered, as if we padded the audio. Otherwise,
-    it will be discarded, as if we trimmed the audio. If there aren't enough frames for 1 partial
-    utterance, this parameter is ignored so that the function always returns at least 1 slice.
-    :param overlap: by how much the partial utterance should overlap. If set to 0, the partial
-    utterances are entirely disjoint.
-    :return: the waveform slices and mel spectrogram slices as lists of array slices. Index
-    respectively the waveform and the mel spectrogram with these slices to obtain the partial
-    utterances.
-    """
-    assert 0 <= overlap < 1
-    assert 0 < min_pad_coverage <= 1
-
-    samples_per_frame = int((sampling_rate * mel_window_step / 1000))
-    n_frames = int(np.ceil((n_samples + 1) / samples_per_frame))
-    frame_step = max(int(np.round(partial_utterance_n_frames * (1 - overlap))), 1)
-
-    # Compute the slices
-    wav_slices, mel_slices = [], []
-    steps = max(1, n_frames - partial_utterance_n_frames + frame_step + 1)
-    for i in range(0, steps, frame_step):
-        mel_range = np.array([i, i + partial_utterance_n_frames])
-        wav_range = mel_range * samples_per_frame
-        mel_slices.append(slice(*mel_range))
-        wav_slices.append(slice(*wav_range))
-
-    # Evaluate whether extra padding is warranted or not
-    last_wav_range = wav_slices[-1]
-    coverage = (n_samples - last_wav_range.start) / (last_wav_range.stop - last_wav_range.start)
-    if coverage < min_pad_coverage and len(mel_slices) > 1:
-        mel_slices = mel_slices[:-1]
-        wav_slices = wav_slices[:-1]
-
-    return wav_slices, mel_slices
-
-
-def embed_utterance(wav, using_partials=True, return_partials=False, **kwargs):
-    """
-    Computes an embedding for a single utterance.
-
-    # TODO: handle multiple wavs to benefit from batching on GPU
-    :param wav: a preprocessed (see audio.py) utterance waveform as a numpy array of float32
-    :param using_partials: if True, then the utterance is split in partial utterances of
-    <partial_utterance_n_frames> frames and the utterance embedding is computed from their
-    normalized average. If False, the utterance is instead computed from feeding the entire
-    spectogram to the network.
-    :param return_partials: if True, the partial embeddings will also be returned along with the
-    wav slices that correspond to the partial embeddings.
-    :param kwargs: additional arguments to compute_partial_splits()
-    :return: the embedding as a numpy array of float32 of shape (model_embedding_size,). If
-    <return_partials> is True, the partial utterances as a numpy array of float32 of shape
-    (n_partials, model_embedding_size) and the wav partials as a list of slices will also be
-    returned. If <using_partials> is simultaneously set to False, both these values will be None
-    instead.
-    """
-    # Process the entire utterance if not using partials
-    if not using_partials:
-        frames = audio.wav_to_mel_spectrogram(wav)
-        embed = embed_frames_batch(frames[None, ...])[0]
-        if return_partials:
-            return embed, None, None
-        return embed
-
-    # Compute where to split the utterance into partials and pad if necessary
-    wave_slices, mel_slices = compute_partial_slices(len(wav), **kwargs)
-    max_wave_length = wave_slices[-1].stop
-    if max_wave_length >= len(wav):
-        wav = np.pad(wav, (0, max_wave_length - len(wav)), "constant")
-
-    # Split the utterance into partials
-    frames = audio.wav_to_mel_spectrogram(wav)
-    frames_batch = np.array([frames[s] for s in mel_slices])
-    partial_embeds = embed_frames_batch(frames_batch)
-
-    # Compute the utterance embedding from the partial embeddings
-    raw_embed = np.mean(partial_embeds, axis=0)
-    embed = raw_embed / np.linalg.norm(raw_embed, 2)
-
-    if return_partials:
-        return embed, partial_embeds, wave_slices
-    return embed
-
-
-def embed_speaker(wavs, **kwargs):
-    raise NotImplemented()
-
-
-def plot_embedding_as_heatmap(embed, ax=None, title="", shape=None, color_range=(0, 0.30)):
-    import matplotlib.pyplot as plt
-    if ax is None:
-        ax = plt.gca()
-
-    if shape is None:
-        height = int(np.sqrt(len(embed)))
-        shape = (height, -1)
-    embed = embed.reshape(shape)
-
-    cmap = cm.get_cmap()
-    mappable = ax.imshow(embed, cmap=cmap)
-    cbar = plt.colorbar(mappable, ax=ax, fraction=0.046, pad=0.04)
-    sm = cm.ScalarMappable(cmap=cmap)
-    sm.set_clim(*color_range)
-
-    ax.set_xticks([]), ax.set_yticks([])
-    ax.set_title(title)
+from encoder.params_data import *
+from encoder.model import SpeakerEncoder
+from encoder.audio import preprocess_wav   # We want to expose this function from here
+from matplotlib import cm
+from encoder import audio
+from pathlib import Path
+import numpy as np
+import torch
+
+_model = None # type: SpeakerEncoder
+_device = None # type: torch.device
+
+
+def load_model(weights_fpath: Path, device=None):
+    """
+    Loads the model in memory. If this function is not explicitely called, it will be run on the
+    first call to embed_frames() with the default weights file.
+
+    :param weights_fpath: the path to saved model weights.
+    :param device: either a torch device or the name of a torch device (e.g. "cpu", "cuda"). The
+    model will be loaded and will run on this device. Outputs will however always be on the cpu.
+    If None, will default to your GPU if it"s available, otherwise your CPU.
+    """
+    # TODO: I think the slow loading of the encoder might have something to do with the device it
+    #   was saved on. Worth investigating.
+    global _model, _device
+    if device is None:
+        _device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    elif isinstance(device, str):
+        _device = torch.device(device)
+    _model = SpeakerEncoder(_device, torch.device("cpu"))
+    checkpoint = torch.load(weights_fpath, map_location=_device, weights_only=False)
+    _model.load_state_dict(checkpoint["model_state"])
+    _model.eval()
+    print("Loaded encoder \"%s\" trained to step %d" % (weights_fpath.name, checkpoint["step"]))
+
+
+def is_loaded():
+    return _model is not None
+
+
+def embed_frames_batch(frames_batch):
+    """
+    Computes embeddings for a batch of mel spectrogram.
+
+    :param frames_batch: a batch mel of spectrogram as a numpy array of float32 of shape
+    (batch_size, n_frames, n_channels)
+    :return: the embeddings as a numpy array of float32 of shape (batch_size, model_embedding_size)
+    """
+    if _model is None:
+        raise Exception("Model was not loaded. Call load_model() before inference.")
+
+    frames = torch.from_numpy(frames_batch).to(_device)
+    embed = _model.forward(frames).detach().cpu().numpy()
+    return embed
+
+
+def compute_partial_slices(n_samples, partial_utterance_n_frames=partials_n_frames,
+                           min_pad_coverage=0.75, overlap=0.5):
+    """
+    Computes where to split an utterance waveform and its corresponding mel spectrogram to obtain
+    partial utterances of <partial_utterance_n_frames> each. Both the waveform and the mel
+    spectrogram slices are returned, so as to make each partial utterance waveform correspond to
+    its spectrogram. This function assumes that the mel spectrogram parameters used are those
+    defined in params_data.py.
+
+    The returned ranges may be indexing further than the length of the waveform. It is
+    recommended that you pad the waveform with zeros up to wave_slices[-1].stop.
+
+    :param n_samples: the number of samples in the waveform
+    :param partial_utterance_n_frames: the number of mel spectrogram frames in each partial
+    utterance
+    :param min_pad_coverage: when reaching the last partial utterance, it may or may not have
+    enough frames. If at least <min_pad_coverage> of <partial_utterance_n_frames> are present,
+    then the last partial utterance will be considered, as if we padded the audio. Otherwise,
+    it will be discarded, as if we trimmed the audio. If there aren't enough frames for 1 partial
+    utterance, this parameter is ignored so that the function always returns at least 1 slice.
+    :param overlap: by how much the partial utterance should overlap. If set to 0, the partial
+    utterances are entirely disjoint.
+    :return: the waveform slices and mel spectrogram slices as lists of array slices. Index
+    respectively the waveform and the mel spectrogram with these slices to obtain the partial
+    utterances.
+    """
+    assert 0 <= overlap < 1
+    assert 0 < min_pad_coverage <= 1
+
+    samples_per_frame = int((sampling_rate * mel_window_step / 1000))
+    n_frames = int(np.ceil((n_samples + 1) / samples_per_frame))
+    frame_step = max(int(np.round(partial_utterance_n_frames * (1 - overlap))), 1)
+
+    # Compute the slices
+    wav_slices, mel_slices = [], []
+    steps = max(1, n_frames - partial_utterance_n_frames + frame_step + 1)
+    for i in range(0, steps, frame_step):
+        mel_range = np.array([i, i + partial_utterance_n_frames])
+        wav_range = mel_range * samples_per_frame
+        mel_slices.append(slice(*mel_range))
+        wav_slices.append(slice(*wav_range))
+
+    # Evaluate whether extra padding is warranted or not
+    last_wav_range = wav_slices[-1]
+    coverage = (n_samples - last_wav_range.start) / (last_wav_range.stop - last_wav_range.start)
+    if coverage < min_pad_coverage and len(mel_slices) > 1:
+        mel_slices = mel_slices[:-1]
+        wav_slices = wav_slices[:-1]
+
+    return wav_slices, mel_slices
+
+
+def embed_utterance(wav, using_partials=True, return_partials=False, **kwargs):
+    """
+    Computes an embedding for a single utterance.
+
+    # TODO: handle multiple wavs to benefit from batching on GPU
+    :param wav: a preprocessed (see audio.py) utterance waveform as a numpy array of float32
+    :param using_partials: if True, then the utterance is split in partial utterances of
+    <partial_utterance_n_frames> frames and the utterance embedding is computed from their
+    normalized average. If False, the utterance is instead computed from feeding the entire
+    spectogram to the network.
+    :param return_partials: if True, the partial embeddings will also be returned along with the
+    wav slices that correspond to the partial embeddings.
+    :param kwargs: additional arguments to compute_partial_splits()
+    :return: the embedding as a numpy array of float32 of shape (model_embedding_size,). If
+    <return_partials> is True, the partial utterances as a numpy array of float32 of shape
+    (n_partials, model_embedding_size) and the wav partials as a list of slices will also be
+    returned. If <using_partials> is simultaneously set to False, both these values will be None
+    instead.
+    """
+    # Process the entire utterance if not using partials
+    if not using_partials:
+        frames = audio.wav_to_mel_spectrogram(wav)
+        embed = embed_frames_batch(frames[None, ...])[0]
+        if return_partials:
+            return embed, None, None
+        return embed
+
+    # Compute where to split the utterance into partials and pad if necessary
+    wave_slices, mel_slices = compute_partial_slices(len(wav), **kwargs)
+    max_wave_length = wave_slices[-1].stop
+    if max_wave_length >= len(wav):
+        wav = np.pad(wav, (0, max_wave_length - len(wav)), "constant")
+
+    # Split the utterance into partials
+    frames = audio.wav_to_mel_spectrogram(wav)
+    frames_batch = np.array([frames[s] for s in mel_slices])
+    partial_embeds = embed_frames_batch(frames_batch)
+
+    # Compute the utterance embedding from the partial embeddings
+    raw_embed = np.mean(partial_embeds, axis=0)
+    embed = raw_embed / np.linalg.norm(raw_embed, 2)
+
+    if return_partials:
+        return embed, partial_embeds, wave_slices
+    return embed
+
+
+def embed_speaker(wavs, **kwargs):
+    raise NotImplemented()
+
+
+def plot_embedding_as_heatmap(embed, ax=None, title="", shape=None, color_range=(0, 0.30)):
+    import matplotlib.pyplot as plt
+    if ax is None:
+        ax = plt.gca()
+
+    if shape is None:
+        height = int(np.sqrt(len(embed)))
+        shape = (height, -1)
+    embed = embed.reshape(shape)
+
+    cmap = cm.get_cmap()
+    mappable = ax.imshow(embed, cmap=cmap)
+    cbar = plt.colorbar(mappable, ax=ax, fraction=0.046, pad=0.04)
+    sm = cm.ScalarMappable(cmap=cmap)
+    sm.set_clim(*color_range)
+
+    ax.set_xticks([]), ax.set_yticks([])
+    ax.set_title(title)

encoder/model.py

@@ -1,135 +1,135 @@
-from encoder.params_model import *
-from encoder.params_data import *
-from scipy.interpolate import interp1d
-from sklearn.metrics import roc_curve
-from torch.nn.utils import clip_grad_norm_
-from scipy.optimize import brentq
-from torch import nn
-import numpy as np
-import torch
-
-
-class SpeakerEncoder(nn.Module):
-    def __init__(self, device, loss_device):
-        super().__init__()
-        self.loss_device = loss_device
-        
-        # Network defition
-        self.lstm = nn.LSTM(input_size=mel_n_channels,
-                            hidden_size=model_hidden_size, 
-                            num_layers=model_num_layers, 
-                            batch_first=True).to(device)
-        self.linear = nn.Linear(in_features=model_hidden_size, 
-                                out_features=model_embedding_size).to(device)
-        self.relu = torch.nn.ReLU().to(device)
-        
-        # Cosine similarity scaling (with fixed initial parameter values)
-        self.similarity_weight = nn.Parameter(torch.tensor([10.])).to(loss_device)
-        self.similarity_bias = nn.Parameter(torch.tensor([-5.])).to(loss_device)
-
-        # Loss
-        self.loss_fn = nn.CrossEntropyLoss().to(loss_device)
-        
-    def do_gradient_ops(self):
-        # Gradient scale
-        self.similarity_weight.grad *= 0.01
-        self.similarity_bias.grad *= 0.01
-            
-        # Gradient clipping
-        clip_grad_norm_(self.parameters(), 3, norm_type=2)
-    
-    def forward(self, utterances, hidden_init=None):
-        """
-        Computes the embeddings of a batch of utterance spectrograms.
-        
-        :param utterances: batch of mel-scale filterbanks of same duration as a tensor of shape 
-        (batch_size, n_frames, n_channels) 
-        :param hidden_init: initial hidden state of the LSTM as a tensor of shape (num_layers, 
-        batch_size, hidden_size). Will default to a tensor of zeros if None.
-        :return: the embeddings as a tensor of shape (batch_size, embedding_size)
-        """
-        # Pass the input through the LSTM layers and retrieve all outputs, the final hidden state
-        # and the final cell state.
-        out, (hidden, cell) = self.lstm(utterances, hidden_init)
-        
-        # We take only the hidden state of the last layer
-        embeds_raw = self.relu(self.linear(hidden[-1]))
-        
-        # L2-normalize it
-        embeds = embeds_raw / (torch.norm(embeds_raw, dim=1, keepdim=True) + 1e-5)        
-
-        return embeds
-    
-    def similarity_matrix(self, embeds):
-        """
-        Computes the similarity matrix according the section 2.1 of GE2E.
-
-        :param embeds: the embeddings as a tensor of shape (speakers_per_batch, 
-        utterances_per_speaker, embedding_size)
-        :return: the similarity matrix as a tensor of shape (speakers_per_batch,
-        utterances_per_speaker, speakers_per_batch)
-        """
-        speakers_per_batch, utterances_per_speaker = embeds.shape[:2]
-        
-        # Inclusive centroids (1 per speaker). Cloning is needed for reverse differentiation
-        centroids_incl = torch.mean(embeds, dim=1, keepdim=True)
-        centroids_incl = centroids_incl.clone() / (torch.norm(centroids_incl, dim=2, keepdim=True) + 1e-5)
-
-        # Exclusive centroids (1 per utterance)
-        centroids_excl = (torch.sum(embeds, dim=1, keepdim=True) - embeds)
-        centroids_excl /= (utterances_per_speaker - 1)
-        centroids_excl = centroids_excl.clone() / (torch.norm(centroids_excl, dim=2, keepdim=True) + 1e-5)
-
-        # Similarity matrix. The cosine similarity of already 2-normed vectors is simply the dot
-        # product of these vectors (which is just an element-wise multiplication reduced by a sum).
-        # We vectorize the computation for efficiency.
-        sim_matrix = torch.zeros(speakers_per_batch, utterances_per_speaker,
-                                 speakers_per_batch).to(self.loss_device)
-        mask_matrix = 1 - np.eye(speakers_per_batch, dtype=np.int)
-        for j in range(speakers_per_batch):
-            mask = np.where(mask_matrix[j])[0]
-            sim_matrix[mask, :, j] = (embeds[mask] * centroids_incl[j]).sum(dim=2)
-            sim_matrix[j, :, j] = (embeds[j] * centroids_excl[j]).sum(dim=1)
-        
-        ## Even more vectorized version (slower maybe because of transpose)
-        # sim_matrix2 = torch.zeros(speakers_per_batch, speakers_per_batch, utterances_per_speaker
-        #                           ).to(self.loss_device)
-        # eye = np.eye(speakers_per_batch, dtype=np.int)
-        # mask = np.where(1 - eye)
-        # sim_matrix2[mask] = (embeds[mask[0]] * centroids_incl[mask[1]]).sum(dim=2)
-        # mask = np.where(eye)
-        # sim_matrix2[mask] = (embeds * centroids_excl).sum(dim=2)
-        # sim_matrix2 = sim_matrix2.transpose(1, 2)
-        
-        sim_matrix = sim_matrix * self.similarity_weight + self.similarity_bias
-        return sim_matrix
-    
-    def loss(self, embeds):
-        """
-        Computes the softmax loss according the section 2.1 of GE2E.
-        
-        :param embeds: the embeddings as a tensor of shape (speakers_per_batch, 
-        utterances_per_speaker, embedding_size)
-        :return: the loss and the EER for this batch of embeddings.
-        """
-        speakers_per_batch, utterances_per_speaker = embeds.shape[:2]
-        
-        # Loss
-        sim_matrix = self.similarity_matrix(embeds)
-        sim_matrix = sim_matrix.reshape((speakers_per_batch * utterances_per_speaker, 
-                                         speakers_per_batch))
-        ground_truth = np.repeat(np.arange(speakers_per_batch), utterances_per_speaker)
-        target = torch.from_numpy(ground_truth).long().to(self.loss_device)
-        loss = self.loss_fn(sim_matrix, target)
-        
-        # EER (not backpropagated)
-        with torch.no_grad():
-            inv_argmax = lambda i: np.eye(1, speakers_per_batch, i, dtype=np.int)[0]
-            labels = np.array([inv_argmax(i) for i in ground_truth])
-            preds = sim_matrix.detach().cpu().numpy()
-
-            # Snippet from https://yangcha.github.io/EER-ROC/
-            fpr, tpr, thresholds = roc_curve(labels.flatten(), preds.flatten())           
-            eer = brentq(lambda x: 1. - x - interp1d(fpr, tpr)(x), 0., 1.)
-            
-        return loss, eer
+from encoder.params_model import *
+from encoder.params_data import *
+from scipy.interpolate import interp1d
+from sklearn.metrics import roc_curve
+from torch.nn.utils import clip_grad_norm_
+from scipy.optimize import brentq
+from torch import nn
+import numpy as np
+import torch
+
+
+class SpeakerEncoder(nn.Module):
+    def __init__(self, device, loss_device):
+        super().__init__()
+        self.loss_device = loss_device
+        
+        # Network defition
+        self.lstm = nn.LSTM(input_size=mel_n_channels,
+                            hidden_size=model_hidden_size, 
+                            num_layers=model_num_layers, 
+                            batch_first=True).to(device)
+        self.linear = nn.Linear(in_features=model_hidden_size, 
+                                out_features=model_embedding_size).to(device)
+        self.relu = torch.nn.ReLU().to(device)
+        
+        # Cosine similarity scaling (with fixed initial parameter values)
+        self.similarity_weight = nn.Parameter(torch.tensor([10.])).to(loss_device)
+        self.similarity_bias = nn.Parameter(torch.tensor([-5.])).to(loss_device)
+
+        # Loss
+        self.loss_fn = nn.CrossEntropyLoss().to(loss_device)
+        
+    def do_gradient_ops(self):
+        # Gradient scale
+        self.similarity_weight.grad *= 0.01
+        self.similarity_bias.grad *= 0.01
+            
+        # Gradient clipping
+        clip_grad_norm_(self.parameters(), 3, norm_type=2)
+    
+    def forward(self, utterances, hidden_init=None):
+        """
+        Computes the embeddings of a batch of utterance spectrograms.
+        
+        :param utterances: batch of mel-scale filterbanks of same duration as a tensor of shape 
+        (batch_size, n_frames, n_channels) 
+        :param hidden_init: initial hidden state of the LSTM as a tensor of shape (num_layers, 
+        batch_size, hidden_size). Will default to a tensor of zeros if None.
+        :return: the embeddings as a tensor of shape (batch_size, embedding_size)
+        """
+        # Pass the input through the LSTM layers and retrieve all outputs, the final hidden state
+        # and the final cell state.
+        out, (hidden, cell) = self.lstm(utterances, hidden_init)
+        
+        # We take only the hidden state of the last layer
+        embeds_raw = self.relu(self.linear(hidden[-1]))
+        
+        # L2-normalize it
+        embeds = embeds_raw / (torch.norm(embeds_raw, dim=1, keepdim=True) + 1e-5)        
+
+        return embeds
+    
+    def similarity_matrix(self, embeds):
+        """
+        Computes the similarity matrix according the section 2.1 of GE2E.
+
+        :param embeds: the embeddings as a tensor of shape (speakers_per_batch, 
+        utterances_per_speaker, embedding_size)
+        :return: the similarity matrix as a tensor of shape (speakers_per_batch,
+        utterances_per_speaker, speakers_per_batch)
+        """
+        speakers_per_batch, utterances_per_speaker = embeds.shape[:2]
+        
+        # Inclusive centroids (1 per speaker). Cloning is needed for reverse differentiation
+        centroids_incl = torch.mean(embeds, dim=1, keepdim=True)
+        centroids_incl = centroids_incl.clone() / (torch.norm(centroids_incl, dim=2, keepdim=True) + 1e-5)
+
+        # Exclusive centroids (1 per utterance)
+        centroids_excl = (torch.sum(embeds, dim=1, keepdim=True) - embeds)
+        centroids_excl /= (utterances_per_speaker - 1)
+        centroids_excl = centroids_excl.clone() / (torch.norm(centroids_excl, dim=2, keepdim=True) + 1e-5)
+
+        # Similarity matrix. The cosine similarity of already 2-normed vectors is simply the dot
+        # product of these vectors (which is just an element-wise multiplication reduced by a sum).
+        # We vectorize the computation for efficiency.
+        sim_matrix = torch.zeros(speakers_per_batch, utterances_per_speaker,
+                                 speakers_per_batch).to(self.loss_device)
+        mask_matrix = 1 - np.eye(speakers_per_batch, dtype=int)
+        for j in range(speakers_per_batch):
+            mask = np.where(mask_matrix[j])[0]
+            sim_matrix[mask, :, j] = (embeds[mask] * centroids_incl[j]).sum(dim=2)
+            sim_matrix[j, :, j] = (embeds[j] * centroids_excl[j]).sum(dim=1)
+        
+        ## Even more vectorized version (slower maybe because of transpose)
+        # sim_matrix2 = torch.zeros(speakers_per_batch, speakers_per_batch, utterances_per_speaker
+        #                           ).to(self.loss_device)
+        # eye = np.eye(speakers_per_batch, dtype=int)
+        # mask = np.where(1 - eye)
+        # sim_matrix2[mask] = (embeds[mask[0]] * centroids_incl[mask[1]]).sum(dim=2)
+        # mask = np.where(eye)
+        # sim_matrix2[mask] = (embeds * centroids_excl).sum(dim=2)
+        # sim_matrix2 = sim_matrix2.transpose(1, 2)
+        
+        sim_matrix = sim_matrix * self.similarity_weight + self.similarity_bias
+        return sim_matrix
+    
+    def loss(self, embeds):
+        """
+        Computes the softmax loss according the section 2.1 of GE2E.
+        
+        :param embeds: the embeddings as a tensor of shape (speakers_per_batch, 
+        utterances_per_speaker, embedding_size)
+        :return: the loss and the EER for this batch of embeddings.
+        """
+        speakers_per_batch, utterances_per_speaker = embeds.shape[:2]
+        
+        # Loss
+        sim_matrix = self.similarity_matrix(embeds)
+        sim_matrix = sim_matrix.reshape((speakers_per_batch * utterances_per_speaker, 
+                                         speakers_per_batch))
+        ground_truth = np.repeat(np.arange(speakers_per_batch), utterances_per_speaker)
+        target = torch.from_numpy(ground_truth).long().to(self.loss_device)
+        loss = self.loss_fn(sim_matrix, target)
+        
+        # EER (not backpropagated)
+        with torch.no_grad():
+            inv_argmax = lambda i: np.eye(1, speakers_per_batch, i, dtype=int)[0]
+            labels = np.array([inv_argmax(i) for i in ground_truth])
+            preds = sim_matrix.detach().cpu().numpy()
+
+            # Snippet from https://yangcha.github.io/EER-ROC/
+            fpr, tpr, thresholds = roc_curve(labels.flatten(), preds.flatten())           
+            eer = brentq(lambda x: 1. - x - interp1d(fpr, tpr)(x), 0., 1.)
+            
+        return loss, eer

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -1,14 +1,12 @@
-#torch==1.12.1+cu116 --find-links https://download.pytorch.org/whl/torch/
-https://download.pytorch.org/whl/cu116/torch-1.12.1%2Bcu116-cp38-cp38-linux_x86_64.whl#sha256=dda312901220895087cc83d3665464a3dc171d04460c61c31af463efbfb54896
-gradio==3.0.4
-PyQt5==5.15.6
-sounddevice==0.4.3
-SoundFile==0.10.3.post1
-tqdm==4.62.3
-Unidecode==1.3.2
-webrtcvad==2.0.10
-librosa==0.8.1
-inflect==5.3.0
-umap-learn==0.5.2
-numpy<1.24
-IPython
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch
+numpy
+scipy
+scikit-learn
+matplotlib
+librosa>=0.10
+soundfile
+webrtcvad-wheels
+inflect
+Unidecode
+tqdm

synthesizer/audio.py

@@ -1,206 +1,206 @@
-import librosa
-import librosa.filters
-import numpy as np
-from scipy import signal
-from scipy.io import wavfile
-import soundfile as sf
-
-
-def load_wav(path, sr):
-    return librosa.core.load(path, sr=sr)[0]
-
-def save_wav(wav, path, sr):
-    wav *= 32767 / max(0.01, np.max(np.abs(wav)))
-    #proposed by @dsmiller
-    wavfile.write(path, sr, wav.astype(np.int16))
-
-def save_wavenet_wav(wav, path, sr):
-    sf.write(path, wav.astype(np.float32), sr)
-
-def preemphasis(wav, k, preemphasize=True):
-    if preemphasize:
-        return signal.lfilter([1, -k], [1], wav)
-    return wav
-
-def inv_preemphasis(wav, k, inv_preemphasize=True):
-    if inv_preemphasize:
-        return signal.lfilter([1], [1, -k], wav)
-    return wav
-
-#From https://github.com/r9y9/wavenet_vocoder/blob/master/audio.py
-def start_and_end_indices(quantized, silence_threshold=2):
-    for start in range(quantized.size):
-        if abs(quantized[start] - 127) > silence_threshold:
-            break
-    for end in range(quantized.size - 1, 1, -1):
-        if abs(quantized[end] - 127) > silence_threshold:
-            break
-    
-    assert abs(quantized[start] - 127) > silence_threshold
-    assert abs(quantized[end] - 127) > silence_threshold
-    
-    return start, end
-
-def get_hop_size(hparams):
-    hop_size = hparams.hop_size
-    if hop_size is None:
-        assert hparams.frame_shift_ms is not None
-        hop_size = int(hparams.frame_shift_ms / 1000 * hparams.sample_rate)
-    return hop_size
-
-def linearspectrogram(wav, hparams):
-    D = _stft(preemphasis(wav, hparams.preemphasis, hparams.preemphasize), hparams)
-    S = _amp_to_db(np.abs(D), hparams) - hparams.ref_level_db
-    
-    if hparams.signal_normalization:
-        return _normalize(S, hparams)
-    return S
-
-def melspectrogram(wav, hparams):
-    D = _stft(preemphasis(wav, hparams.preemphasis, hparams.preemphasize), hparams)
-    S = _amp_to_db(_linear_to_mel(np.abs(D), hparams), hparams) - hparams.ref_level_db
-    
-    if hparams.signal_normalization:
-        return _normalize(S, hparams)
-    return S
-
-def inv_linear_spectrogram(linear_spectrogram, hparams):
-    """Converts linear spectrogram to waveform using librosa"""
-    if hparams.signal_normalization:
-        D = _denormalize(linear_spectrogram, hparams)
-    else:
-        D = linear_spectrogram
-    
-    S = _db_to_amp(D + hparams.ref_level_db) #Convert back to linear
-    
-    if hparams.use_lws:
-        processor = _lws_processor(hparams)
-        D = processor.run_lws(S.astype(np.float64).T ** hparams.power)
-        y = processor.istft(D).astype(np.float32)
-        return inv_preemphasis(y, hparams.preemphasis, hparams.preemphasize)
-    else:
-        return inv_preemphasis(_griffin_lim(S ** hparams.power, hparams), hparams.preemphasis, hparams.preemphasize)
-
-def inv_mel_spectrogram(mel_spectrogram, hparams):
-    """Converts mel spectrogram to waveform using librosa"""
-    if hparams.signal_normalization:
-        D = _denormalize(mel_spectrogram, hparams)
-    else:
-        D = mel_spectrogram
-    
-    S = _mel_to_linear(_db_to_amp(D + hparams.ref_level_db), hparams)  # Convert back to linear
-    
-    if hparams.use_lws:
-        processor = _lws_processor(hparams)
-        D = processor.run_lws(S.astype(np.float64).T ** hparams.power)
-        y = processor.istft(D).astype(np.float32)
-        return inv_preemphasis(y, hparams.preemphasis, hparams.preemphasize)
-    else:
-        return inv_preemphasis(_griffin_lim(S ** hparams.power, hparams), hparams.preemphasis, hparams.preemphasize)
-
-def _lws_processor(hparams):
-    import lws
-    return lws.lws(hparams.n_fft, get_hop_size(hparams), fftsize=hparams.win_size, mode="speech")
-
-def _griffin_lim(S, hparams):
-    """librosa implementation of Griffin-Lim
-    Based on https://github.com/librosa/librosa/issues/434
-    """
-    angles = np.exp(2j * np.pi * np.random.rand(*S.shape))
-    S_complex = np.abs(S).astype(np.complex)
-    y = _istft(S_complex * angles, hparams)
-    for i in range(hparams.griffin_lim_iters):
-        angles = np.exp(1j * np.angle(_stft(y, hparams)))
-        y = _istft(S_complex * angles, hparams)
-    return y
-
-def _stft(y, hparams):
-    if hparams.use_lws:
-        return _lws_processor(hparams).stft(y).T
-    else:
-        return librosa.stft(y=y, n_fft=hparams.n_fft, hop_length=get_hop_size(hparams), win_length=hparams.win_size)
-
-def _istft(y, hparams):
-    return librosa.istft(y, hop_length=get_hop_size(hparams), win_length=hparams.win_size)
-
-##########################################################
-#Those are only correct when using lws!!! (This was messing with Wavenet quality for a long time!)
-def num_frames(length, fsize, fshift):
-    """Compute number of time frames of spectrogram
-    """
-    pad = (fsize - fshift)
-    if length % fshift == 0:
-        M = (length + pad * 2 - fsize) // fshift + 1
-    else:
-        M = (length + pad * 2 - fsize) // fshift + 2
-    return M
-
-
-def pad_lr(x, fsize, fshift):
-    """Compute left and right padding
-    """
-    M = num_frames(len(x), fsize, fshift)
-    pad = (fsize - fshift)
-    T = len(x) + 2 * pad
-    r = (M - 1) * fshift + fsize - T
-    return pad, pad + r
-##########################################################
-#Librosa correct padding
-def librosa_pad_lr(x, fsize, fshift):
-    return 0, (x.shape[0] // fshift + 1) * fshift - x.shape[0]
-
-# Conversions
-_mel_basis = None
-_inv_mel_basis = None
-
-def _linear_to_mel(spectogram, hparams):
-    global _mel_basis
-    if _mel_basis is None:
-        _mel_basis = _build_mel_basis(hparams)
-    return np.dot(_mel_basis, spectogram)
-
-def _mel_to_linear(mel_spectrogram, hparams):
-    global _inv_mel_basis
-    if _inv_mel_basis is None:
-        _inv_mel_basis = np.linalg.pinv(_build_mel_basis(hparams))
-    return np.maximum(1e-10, np.dot(_inv_mel_basis, mel_spectrogram))
-
-def _build_mel_basis(hparams):
-    assert hparams.fmax <= hparams.sample_rate // 2
-    return librosa.filters.mel(hparams.sample_rate, hparams.n_fft, n_mels=hparams.num_mels,
-                               fmin=hparams.fmin, fmax=hparams.fmax)
-
-def _amp_to_db(x, hparams):
-    min_level = np.exp(hparams.min_level_db / 20 * np.log(10))
-    return 20 * np.log10(np.maximum(min_level, x))
-
-def _db_to_amp(x):
-    return np.power(10.0, (x) * 0.05)
-
-def _normalize(S, hparams):
-    if hparams.allow_clipping_in_normalization:
-        if hparams.symmetric_mels:
-            return np.clip((2 * hparams.max_abs_value) * ((S - hparams.min_level_db) / (-hparams.min_level_db)) - hparams.max_abs_value,
-                           -hparams.max_abs_value, hparams.max_abs_value)
-        else:
-            return np.clip(hparams.max_abs_value * ((S - hparams.min_level_db) / (-hparams.min_level_db)), 0, hparams.max_abs_value)
-    
-    assert S.max() <= 0 and S.min() - hparams.min_level_db >= 0
-    if hparams.symmetric_mels:
-        return (2 * hparams.max_abs_value) * ((S - hparams.min_level_db) / (-hparams.min_level_db)) - hparams.max_abs_value
-    else:
-        return hparams.max_abs_value * ((S - hparams.min_level_db) / (-hparams.min_level_db))
-
-def _denormalize(D, hparams):
-    if hparams.allow_clipping_in_normalization:
-        if hparams.symmetric_mels:
-            return (((np.clip(D, -hparams.max_abs_value,
-                              hparams.max_abs_value) + hparams.max_abs_value) * -hparams.min_level_db / (2 * hparams.max_abs_value))
-                    + hparams.min_level_db)
-        else:
-            return ((np.clip(D, 0, hparams.max_abs_value) * -hparams.min_level_db / hparams.max_abs_value) + hparams.min_level_db)
-    
-    if hparams.symmetric_mels:
-        return (((D + hparams.max_abs_value) * -hparams.min_level_db / (2 * hparams.max_abs_value)) + hparams.min_level_db)
-    else:
-        return ((D * -hparams.min_level_db / hparams.max_abs_value) + hparams.min_level_db)
+import librosa
+import librosa.filters
+import numpy as np
+from scipy import signal
+from scipy.io import wavfile
+import soundfile as sf
+
+
+def load_wav(path, sr):
+    return librosa.core.load(path, sr=sr)[0]
+
+def save_wav(wav, path, sr):
+    wav *= 32767 / max(0.01, np.max(np.abs(wav)))
+    #proposed by @dsmiller
+    wavfile.write(path, sr, wav.astype(np.int16))
+
+def save_wavenet_wav(wav, path, sr):
+    sf.write(path, wav.astype(np.float32), sr)
+
+def preemphasis(wav, k, preemphasize=True):
+    if preemphasize:
+        return signal.lfilter([1, -k], [1], wav)
+    return wav
+
+def inv_preemphasis(wav, k, inv_preemphasize=True):
+    if inv_preemphasize:
+        return signal.lfilter([1], [1, -k], wav)
+    return wav
+
+#From https://github.com/r9y9/wavenet_vocoder/blob/master/audio.py
+def start_and_end_indices(quantized, silence_threshold=2):
+    for start in range(quantized.size):
+        if abs(quantized[start] - 127) > silence_threshold:
+            break
+    for end in range(quantized.size - 1, 1, -1):
+        if abs(quantized[end] - 127) > silence_threshold:
+            break
+    
+    assert abs(quantized[start] - 127) > silence_threshold
+    assert abs(quantized[end] - 127) > silence_threshold
+    
+    return start, end
+
+def get_hop_size(hparams):
+    hop_size = hparams.hop_size
+    if hop_size is None:
+        assert hparams.frame_shift_ms is not None
+        hop_size = int(hparams.frame_shift_ms / 1000 * hparams.sample_rate)
+    return hop_size
+
+def linearspectrogram(wav, hparams):
+    D = _stft(preemphasis(wav, hparams.preemphasis, hparams.preemphasize), hparams)
+    S = _amp_to_db(np.abs(D), hparams) - hparams.ref_level_db
+    
+    if hparams.signal_normalization:
+        return _normalize(S, hparams)
+    return S
+
+def melspectrogram(wav, hparams):
+    D = _stft(preemphasis(wav, hparams.preemphasis, hparams.preemphasize), hparams)
+    S = _amp_to_db(_linear_to_mel(np.abs(D), hparams), hparams) - hparams.ref_level_db
+    
+    if hparams.signal_normalization:
+        return _normalize(S, hparams)
+    return S
+
+def inv_linear_spectrogram(linear_spectrogram, hparams):
+    """Converts linear spectrogram to waveform using librosa"""
+    if hparams.signal_normalization:
+        D = _denormalize(linear_spectrogram, hparams)
+    else:
+        D = linear_spectrogram
+    
+    S = _db_to_amp(D + hparams.ref_level_db) #Convert back to linear
+    
+    if hparams.use_lws:
+        processor = _lws_processor(hparams)
+        D = processor.run_lws(S.astype(np.float64).T ** hparams.power)
+        y = processor.istft(D).astype(np.float32)
+        return inv_preemphasis(y, hparams.preemphasis, hparams.preemphasize)
+    else:
+        return inv_preemphasis(_griffin_lim(S ** hparams.power, hparams), hparams.preemphasis, hparams.preemphasize)
+
+def inv_mel_spectrogram(mel_spectrogram, hparams):
+    """Converts mel spectrogram to waveform using librosa"""
+    if hparams.signal_normalization:
+        D = _denormalize(mel_spectrogram, hparams)
+    else:
+        D = mel_spectrogram
+    
+    S = _mel_to_linear(_db_to_amp(D + hparams.ref_level_db), hparams)  # Convert back to linear
+    
+    if hparams.use_lws:
+        processor = _lws_processor(hparams)
+        D = processor.run_lws(S.astype(np.float64).T ** hparams.power)
+        y = processor.istft(D).astype(np.float32)
+        return inv_preemphasis(y, hparams.preemphasis, hparams.preemphasize)
+    else:
+        return inv_preemphasis(_griffin_lim(S ** hparams.power, hparams), hparams.preemphasis, hparams.preemphasize)
+
+def _lws_processor(hparams):
+    import lws
+    return lws.lws(hparams.n_fft, get_hop_size(hparams), fftsize=hparams.win_size, mode="speech")
+
+def _griffin_lim(S, hparams):
+    """librosa implementation of Griffin-Lim
+    Based on https://github.com/librosa/librosa/issues/434
+    """
+    angles = np.exp(2j * np.pi * np.random.rand(*S.shape))
+    S_complex = np.abs(S).astype(complex)
+    y = _istft(S_complex * angles, hparams)
+    for i in range(hparams.griffin_lim_iters):
+        angles = np.exp(1j * np.angle(_stft(y, hparams)))
+        y = _istft(S_complex * angles, hparams)
+    return y
+
+def _stft(y, hparams):
+    if hparams.use_lws:
+        return _lws_processor(hparams).stft(y).T
+    else:
+        return librosa.stft(y=y, n_fft=hparams.n_fft, hop_length=get_hop_size(hparams), win_length=hparams.win_size)
+
+def _istft(y, hparams):
+    return librosa.istft(y, hop_length=get_hop_size(hparams), win_length=hparams.win_size)
+
+##########################################################
+#Those are only correct when using lws!!! (This was messing with Wavenet quality for a long time!)
+def num_frames(length, fsize, fshift):
+    """Compute number of time frames of spectrogram
+    """
+    pad = (fsize - fshift)
+    if length % fshift == 0:
+        M = (length + pad * 2 - fsize) // fshift + 1
+    else:
+        M = (length + pad * 2 - fsize) // fshift + 2
+    return M
+
+
+def pad_lr(x, fsize, fshift):
+    """Compute left and right padding
+    """
+    M = num_frames(len(x), fsize, fshift)
+    pad = (fsize - fshift)
+    T = len(x) + 2 * pad
+    r = (M - 1) * fshift + fsize - T
+    return pad, pad + r
+##########################################################
+#Librosa correct padding
+def librosa_pad_lr(x, fsize, fshift):
+    return 0, (x.shape[0] // fshift + 1) * fshift - x.shape[0]
+
+# Conversions
+_mel_basis = None
+_inv_mel_basis = None
+
+def _linear_to_mel(spectogram, hparams):
+    global _mel_basis
+    if _mel_basis is None:
+        _mel_basis = _build_mel_basis(hparams)
+    return np.dot(_mel_basis, spectogram)
+
+def _mel_to_linear(mel_spectrogram, hparams):
+    global _inv_mel_basis
+    if _inv_mel_basis is None:
+        _inv_mel_basis = np.linalg.pinv(_build_mel_basis(hparams))
+    return np.maximum(1e-10, np.dot(_inv_mel_basis, mel_spectrogram))
+
+def _build_mel_basis(hparams):
+    assert hparams.fmax <= hparams.sample_rate // 2
+    return librosa.filters.mel(sr=hparams.sample_rate, n_fft=hparams.n_fft, n_mels=hparams.num_mels,
+                               fmin=hparams.fmin, fmax=hparams.fmax)
+
+def _amp_to_db(x, hparams):
+    min_level = np.exp(hparams.min_level_db / 20 * np.log(10))
+    return 20 * np.log10(np.maximum(min_level, x))
+
+def _db_to_amp(x):
+    return np.power(10.0, (x) * 0.05)
+
+def _normalize(S, hparams):
+    if hparams.allow_clipping_in_normalization:
+        if hparams.symmetric_mels:
+            return np.clip((2 * hparams.max_abs_value) * ((S - hparams.min_level_db) / (-hparams.min_level_db)) - hparams.max_abs_value,
+                           -hparams.max_abs_value, hparams.max_abs_value)
+        else:
+            return np.clip(hparams.max_abs_value * ((S - hparams.min_level_db) / (-hparams.min_level_db)), 0, hparams.max_abs_value)
+    
+    assert S.max() <= 0 and S.min() - hparams.min_level_db >= 0
+    if hparams.symmetric_mels:
+        return (2 * hparams.max_abs_value) * ((S - hparams.min_level_db) / (-hparams.min_level_db)) - hparams.max_abs_value
+    else:
+        return hparams.max_abs_value * ((S - hparams.min_level_db) / (-hparams.min_level_db))
+
+def _denormalize(D, hparams):
+    if hparams.allow_clipping_in_normalization:
+        if hparams.symmetric_mels:
+            return (((np.clip(D, -hparams.max_abs_value,
+                              hparams.max_abs_value) + hparams.max_abs_value) * -hparams.min_level_db / (2 * hparams.max_abs_value))
+                    + hparams.min_level_db)
+        else:
+            return ((np.clip(D, 0, hparams.max_abs_value) * -hparams.min_level_db / hparams.max_abs_value) + hparams.min_level_db)
+    
+    if hparams.symmetric_mels:
+        return (((D + hparams.max_abs_value) * -hparams.min_level_db / (2 * hparams.max_abs_value)) + hparams.min_level_db)
+    else:
+        return ((D * -hparams.min_level_db / hparams.max_abs_value) + hparams.min_level_db)

synthesizer/inference.py

@@ -134,7 +134,7 @@ class Synthesizer:
         train the synthesizer.
         """
         print("Loading fpath and hparams.sample_rate :",str(fpath), hparams.sample_rate)
-        wav = librosa.load(str(fpath), hparams.sample_rate)[0]
+        wav = librosa.load(str(fpath), sr=hparams.sample_rate)[0]
         if hparams.rescale:
             wav = wav / np.abs(wav).max() * hparams.rescaling_max
         return wav

synthesizer/models/tacotron.py

@@ -1,519 +1,519 @@
-import os
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from pathlib import Path
-from typing import Union
-
-
-class HighwayNetwork(nn.Module):
-    def __init__(self, size):
-        super().__init__()
-        self.W1 = nn.Linear(size, size)
-        self.W2 = nn.Linear(size, size)
-        self.W1.bias.data.fill_(0.)
-
-    def forward(self, x):
-        x1 = self.W1(x)
-        x2 = self.W2(x)
-        g = torch.sigmoid(x2)
-        y = g * F.relu(x1) + (1. - g) * x
-        return y
-
-
-class Encoder(nn.Module):
-    def __init__(self, embed_dims, num_chars, encoder_dims, K, num_highways, dropout):
-        super().__init__()
-        prenet_dims = (encoder_dims, encoder_dims)
-        cbhg_channels = encoder_dims
-        self.embedding = nn.Embedding(num_chars, embed_dims)
-        self.pre_net = PreNet(embed_dims, fc1_dims=prenet_dims[0], fc2_dims=prenet_dims[1],
-                              dropout=dropout)
-        self.cbhg = CBHG(K=K, in_channels=cbhg_channels, channels=cbhg_channels,
-                         proj_channels=[cbhg_channels, cbhg_channels],
-                         num_highways=num_highways)
-
-    def forward(self, x, speaker_embedding=None):
-        x = self.embedding(x)
-        x = self.pre_net(x)
-        x.transpose_(1, 2)
-        x = self.cbhg(x)
-        if speaker_embedding is not None:
-            x = self.add_speaker_embedding(x, speaker_embedding)
-        return x
-
-    def add_speaker_embedding(self, x, speaker_embedding):
-        # SV2TTS
-        # The input x is the encoder output and is a 3D tensor with size (batch_size, num_chars, tts_embed_dims)
-        # When training, speaker_embedding is also a 2D tensor with size (batch_size, speaker_embedding_size)
-        #     (for inference, speaker_embedding is a 1D tensor with size (speaker_embedding_size))
-        # This concats the speaker embedding for each char in the encoder output
-
-        # Save the dimensions as human-readable names
-        batch_size = x.size()[0]
-        num_chars = x.size()[1]
-
-        if speaker_embedding.dim() == 1:
-            idx = 0
-        else:
-            idx = 1
-
-        # Start by making a copy of each speaker embedding to match the input text length
-        # The output of this has size (batch_size, num_chars * tts_embed_dims)
-        speaker_embedding_size = speaker_embedding.size()[idx]
-        e = speaker_embedding.repeat_interleave(num_chars, dim=idx)
-
-        # Reshape it and transpose
-        e = e.reshape(batch_size, speaker_embedding_size, num_chars)
-        e = e.transpose(1, 2)
-
-        # Concatenate the tiled speaker embedding with the encoder output
-        x = torch.cat((x, e), 2)
-        return x
-
-
-class BatchNormConv(nn.Module):
-    def __init__(self, in_channels, out_channels, kernel, relu=True):
-        super().__init__()
-        self.conv = nn.Conv1d(in_channels, out_channels, kernel, stride=1, padding=kernel // 2, bias=False)
-        self.bnorm = nn.BatchNorm1d(out_channels)
-        self.relu = relu
-
-    def forward(self, x):
-        x = self.conv(x)
-        x = F.relu(x) if self.relu is True else x
-        return self.bnorm(x)
-
-
-class CBHG(nn.Module):
-    def __init__(self, K, in_channels, channels, proj_channels, num_highways):
-        super().__init__()
-
-        # List of all rnns to call `flatten_parameters()` on
-        self._to_flatten = []
-
-        self.bank_kernels = [i for i in range(1, K + 1)]
-        self.conv1d_bank = nn.ModuleList()
-        for k in self.bank_kernels:
-            conv = BatchNormConv(in_channels, channels, k)
-            self.conv1d_bank.append(conv)
-
-        self.maxpool = nn.MaxPool1d(kernel_size=2, stride=1, padding=1)
-
-        self.conv_project1 = BatchNormConv(len(self.bank_kernels) * channels, proj_channels[0], 3)
-        self.conv_project2 = BatchNormConv(proj_channels[0], proj_channels[1], 3, relu=False)
-
-        # Fix the highway input if necessary
-        if proj_channels[-1] != channels:
-            self.highway_mismatch = True
-            self.pre_highway = nn.Linear(proj_channels[-1], channels, bias=False)
-        else:
-            self.highway_mismatch = False
-
-        self.highways = nn.ModuleList()
-        for i in range(num_highways):
-            hn = HighwayNetwork(channels)
-            self.highways.append(hn)
-
-        self.rnn = nn.GRU(channels, channels // 2, batch_first=True, bidirectional=True)
-        self._to_flatten.append(self.rnn)
-
-        # Avoid fragmentation of RNN parameters and associated warning
-        self._flatten_parameters()
-
-    def forward(self, x):
-        # Although we `_flatten_parameters()` on init, when using DataParallel
-        # the model gets replicated, making it no longer guaranteed that the
-        # weights are contiguous in GPU memory. Hence, we must call it again
-        self._flatten_parameters()
-
-        # Save these for later
-        residual = x
-        seq_len = x.size(-1)
-        conv_bank = []
-
-        # Convolution Bank
-        for conv in self.conv1d_bank:
-            c = conv(x) # Convolution
-            conv_bank.append(c[:, :, :seq_len])
-
-        # Stack along the channel axis
-        conv_bank = torch.cat(conv_bank, dim=1)
-
-        # dump the last padding to fit residual
-        x = self.maxpool(conv_bank)[:, :, :seq_len]
-
-        # Conv1d projections
-        x = self.conv_project1(x)
-        x = self.conv_project2(x)
-
-        # Residual Connect
-        x = x + residual
-
-        # Through the highways
-        x = x.transpose(1, 2)
-        if self.highway_mismatch is True:
-            x = self.pre_highway(x)
-        for h in self.highways: x = h(x)
-
-        # And then the RNN
-        x, _ = self.rnn(x)
-        return x
-
-    def _flatten_parameters(self):
-        """Calls `flatten_parameters` on all the rnns used by the WaveRNN. Used
-        to improve efficiency and avoid PyTorch yelling at us."""
-        [m.flatten_parameters() for m in self._to_flatten]
-
-class PreNet(nn.Module):
-    def __init__(self, in_dims, fc1_dims=256, fc2_dims=128, dropout=0.5):
-        super().__init__()
-        self.fc1 = nn.Linear(in_dims, fc1_dims)
-        self.fc2 = nn.Linear(fc1_dims, fc2_dims)
-        self.p = dropout
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = F.relu(x)
-        x = F.dropout(x, self.p, training=True)
-        x = self.fc2(x)
-        x = F.relu(x)
-        x = F.dropout(x, self.p, training=True)
-        return x
-
-
-class Attention(nn.Module):
-    def __init__(self, attn_dims):
-        super().__init__()
-        self.W = nn.Linear(attn_dims, attn_dims, bias=False)
-        self.v = nn.Linear(attn_dims, 1, bias=False)
-
-    def forward(self, encoder_seq_proj, query, t):
-
-        # print(encoder_seq_proj.shape)
-        # Transform the query vector
-        query_proj = self.W(query).unsqueeze(1)
-
-        # Compute the scores
-        u = self.v(torch.tanh(encoder_seq_proj + query_proj))
-        scores = F.softmax(u, dim=1)
-
-        return scores.transpose(1, 2)
-
-
-class LSA(nn.Module):
-    def __init__(self, attn_dim, kernel_size=31, filters=32):
-        super().__init__()
-        self.conv = nn.Conv1d(1, filters, padding=(kernel_size - 1) // 2, kernel_size=kernel_size, bias=True)
-        self.L = nn.Linear(filters, attn_dim, bias=False)
-        self.W = nn.Linear(attn_dim, attn_dim, bias=True) # Include the attention bias in this term
-        self.v = nn.Linear(attn_dim, 1, bias=False)
-        self.cumulative = None
-        self.attention = None
-
-    def init_attention(self, encoder_seq_proj):
-        device = next(self.parameters()).device  # use same device as parameters
-        b, t, c = encoder_seq_proj.size()
-        self.cumulative = torch.zeros(b, t, device=device)
-        self.attention = torch.zeros(b, t, device=device)
-
-    def forward(self, encoder_seq_proj, query, t, chars):
-
-        if t == 0: self.init_attention(encoder_seq_proj)
-
-        processed_query = self.W(query).unsqueeze(1)
-
-        location = self.cumulative.unsqueeze(1)
-        processed_loc = self.L(self.conv(location).transpose(1, 2))
-
-        u = self.v(torch.tanh(processed_query + encoder_seq_proj + processed_loc))
-        u = u.squeeze(-1)
-
-        # Mask zero padding chars
-        u = u * (chars != 0).float()
-
-        # Smooth Attention
-        # scores = torch.sigmoid(u) / torch.sigmoid(u).sum(dim=1, keepdim=True)
-        scores = F.softmax(u, dim=1)
-        self.attention = scores
-        self.cumulative = self.cumulative + self.attention
-
-        return scores.unsqueeze(-1).transpose(1, 2)
-
-
-class Decoder(nn.Module):
-    # Class variable because its value doesn't change between classes
-    # yet ought to be scoped by class because its a property of a Decoder
-    max_r = 20
-    def __init__(self, n_mels, encoder_dims, decoder_dims, lstm_dims,
-                 dropout, speaker_embedding_size):
-        super().__init__()
-        self.register_buffer("r", torch.tensor(1, dtype=torch.int))
-        self.n_mels = n_mels
-        prenet_dims = (decoder_dims * 2, decoder_dims * 2)
-        self.prenet = PreNet(n_mels, fc1_dims=prenet_dims[0], fc2_dims=prenet_dims[1],
-                             dropout=dropout)
-        self.attn_net = LSA(decoder_dims)
-        self.attn_rnn = nn.GRUCell(encoder_dims + prenet_dims[1] + speaker_embedding_size, decoder_dims)
-        self.rnn_input = nn.Linear(encoder_dims + decoder_dims + speaker_embedding_size, lstm_dims)
-        self.res_rnn1 = nn.LSTMCell(lstm_dims, lstm_dims)
-        self.res_rnn2 = nn.LSTMCell(lstm_dims, lstm_dims)
-        self.mel_proj = nn.Linear(lstm_dims, n_mels * self.max_r, bias=False)
-        self.stop_proj = nn.Linear(encoder_dims + speaker_embedding_size + lstm_dims, 1)
-
-    def zoneout(self, prev, current, p=0.1):
-        device = next(self.parameters()).device  # Use same device as parameters
-        mask = torch.zeros(prev.size(), device=device).bernoulli_(p)
-        return prev * mask + current * (1 - mask)
-
-    def forward(self, encoder_seq, encoder_seq_proj, prenet_in,
-                hidden_states, cell_states, context_vec, t, chars):
-
-        # Need this for reshaping mels
-        batch_size = encoder_seq.size(0)
-
-        # Unpack the hidden and cell states
-        attn_hidden, rnn1_hidden, rnn2_hidden = hidden_states
-        rnn1_cell, rnn2_cell = cell_states
-
-        # PreNet for the Attention RNN
-        prenet_out = self.prenet(prenet_in)
-
-        # Compute the Attention RNN hidden state
-        attn_rnn_in = torch.cat([context_vec, prenet_out], dim=-1)
-        attn_hidden = self.attn_rnn(attn_rnn_in.squeeze(1), attn_hidden)
-
-        # Compute the attention scores
-        scores = self.attn_net(encoder_seq_proj, attn_hidden, t, chars)
-
-        # Dot product to create the context vector
-        context_vec = scores @ encoder_seq
-        context_vec = context_vec.squeeze(1)
-
-        # Concat Attention RNN output w. Context Vector & project
-        x = torch.cat([context_vec, attn_hidden], dim=1)
-        x = self.rnn_input(x)
-
-        # Compute first Residual RNN
-        rnn1_hidden_next, rnn1_cell = self.res_rnn1(x, (rnn1_hidden, rnn1_cell))
-        if self.training:
-            rnn1_hidden = self.zoneout(rnn1_hidden, rnn1_hidden_next)
-        else:
-            rnn1_hidden = rnn1_hidden_next
-        x = x + rnn1_hidden
-
-        # Compute second Residual RNN
-        rnn2_hidden_next, rnn2_cell = self.res_rnn2(x, (rnn2_hidden, rnn2_cell))
-        if self.training:
-            rnn2_hidden = self.zoneout(rnn2_hidden, rnn2_hidden_next)
-        else:
-            rnn2_hidden = rnn2_hidden_next
-        x = x + rnn2_hidden
-
-        # Project Mels
-        mels = self.mel_proj(x)
-        mels = mels.view(batch_size, self.n_mels, self.max_r)[:, :, :self.r]
-        hidden_states = (attn_hidden, rnn1_hidden, rnn2_hidden)
-        cell_states = (rnn1_cell, rnn2_cell)
-
-        # Stop token prediction
-        s = torch.cat((x, context_vec), dim=1)
-        s = self.stop_proj(s)
-        stop_tokens = torch.sigmoid(s)
-
-        return mels, scores, hidden_states, cell_states, context_vec, stop_tokens
-
-
-class Tacotron(nn.Module):
-    def __init__(self, embed_dims, num_chars, encoder_dims, decoder_dims, n_mels, 
-                 fft_bins, postnet_dims, encoder_K, lstm_dims, postnet_K, num_highways,
-                 dropout, stop_threshold, speaker_embedding_size):
-        super().__init__()
-        self.n_mels = n_mels
-        self.lstm_dims = lstm_dims
-        self.encoder_dims = encoder_dims
-        self.decoder_dims = decoder_dims
-        self.speaker_embedding_size = speaker_embedding_size
-        self.encoder = Encoder(embed_dims, num_chars, encoder_dims,
-                               encoder_K, num_highways, dropout)
-        self.encoder_proj = nn.Linear(encoder_dims + speaker_embedding_size, decoder_dims, bias=False)
-        self.decoder = Decoder(n_mels, encoder_dims, decoder_dims, lstm_dims,
-                               dropout, speaker_embedding_size)
-        self.postnet = CBHG(postnet_K, n_mels, postnet_dims,
-                            [postnet_dims, fft_bins], num_highways)
-        self.post_proj = nn.Linear(postnet_dims, fft_bins, bias=False)
-
-        self.init_model()
-        self.num_params()
-
-        self.register_buffer("step", torch.zeros(1, dtype=torch.long))
-        self.register_buffer("stop_threshold", torch.tensor(stop_threshold, dtype=torch.float32))
-
-    @property
-    def r(self):
-        return self.decoder.r.item()
-
-    @r.setter
-    def r(self, value):
-        self.decoder.r = self.decoder.r.new_tensor(value, requires_grad=False)
-
-    def forward(self, x, m, speaker_embedding):
-        device = next(self.parameters()).device  # use same device as parameters
-
-        self.step += 1
-        batch_size, _, steps  = m.size()
-
-        # Initialise all hidden states and pack into tuple
-        attn_hidden = torch.zeros(batch_size, self.decoder_dims, device=device)
-        rnn1_hidden = torch.zeros(batch_size, self.lstm_dims, device=device)
-        rnn2_hidden = torch.zeros(batch_size, self.lstm_dims, device=device)
-        hidden_states = (attn_hidden, rnn1_hidden, rnn2_hidden)
-
-        # Initialise all lstm cell states and pack into tuple
-        rnn1_cell = torch.zeros(batch_size, self.lstm_dims, device=device)
-        rnn2_cell = torch.zeros(batch_size, self.lstm_dims, device=device)
-        cell_states = (rnn1_cell, rnn2_cell)
-
-        # <GO> Frame for start of decoder loop
-        go_frame = torch.zeros(batch_size, self.n_mels, device=device)
-
-        # Need an initial context vector
-        context_vec = torch.zeros(batch_size, self.encoder_dims + self.speaker_embedding_size, device=device)
-
-        # SV2TTS: Run the encoder with the speaker embedding
-        # The projection avoids unnecessary matmuls in the decoder loop
-        encoder_seq = self.encoder(x, speaker_embedding)
-        encoder_seq_proj = self.encoder_proj(encoder_seq)
-
-        # Need a couple of lists for outputs
-        mel_outputs, attn_scores, stop_outputs = [], [], []
-
-        # Run the decoder loop
-        for t in range(0, steps, self.r):
-            prenet_in = m[:, :, t - 1] if t > 0 else go_frame
-            mel_frames, scores, hidden_states, cell_states, context_vec, stop_tokens = \
-                self.decoder(encoder_seq, encoder_seq_proj, prenet_in,
-                             hidden_states, cell_states, context_vec, t, x)
-            mel_outputs.append(mel_frames)
-            attn_scores.append(scores)
-            stop_outputs.extend([stop_tokens] * self.r)
-
-        # Concat the mel outputs into sequence
-        mel_outputs = torch.cat(mel_outputs, dim=2)
-
-        # Post-Process for Linear Spectrograms
-        postnet_out = self.postnet(mel_outputs)
-        linear = self.post_proj(postnet_out)
-        linear = linear.transpose(1, 2)
-
-        # For easy visualisation
-        attn_scores = torch.cat(attn_scores, 1)
-        # attn_scores = attn_scores.cpu().data.numpy()
-        stop_outputs = torch.cat(stop_outputs, 1)
-
-        return mel_outputs, linear, attn_scores, stop_outputs
-
-    def generate(self, x, speaker_embedding=None, steps=2000):
-        self.eval()
-        device = next(self.parameters()).device  # use same device as parameters
-
-        batch_size, _  = x.size()
-
-        # Need to initialise all hidden states and pack into tuple for tidyness
-        attn_hidden = torch.zeros(batch_size, self.decoder_dims, device=device)
-        rnn1_hidden = torch.zeros(batch_size, self.lstm_dims, device=device)
-        rnn2_hidden = torch.zeros(batch_size, self.lstm_dims, device=device)
-        hidden_states = (attn_hidden, rnn1_hidden, rnn2_hidden)
-
-        # Need to initialise all lstm cell states and pack into tuple for tidyness
-        rnn1_cell = torch.zeros(batch_size, self.lstm_dims, device=device)
-        rnn2_cell = torch.zeros(batch_size, self.lstm_dims, device=device)
-        cell_states = (rnn1_cell, rnn2_cell)
-
-        # Need a <GO> Frame for start of decoder loop
-        go_frame = torch.zeros(batch_size, self.n_mels, device=device)
-
-        # Need an initial context vector
-        context_vec = torch.zeros(batch_size, self.encoder_dims + self.speaker_embedding_size, device=device)
-
-        # SV2TTS: Run the encoder with the speaker embedding
-        # The projection avoids unnecessary matmuls in the decoder loop
-        encoder_seq = self.encoder(x, speaker_embedding)
-        encoder_seq_proj = self.encoder_proj(encoder_seq)
-
-        # Need a couple of lists for outputs
-        mel_outputs, attn_scores, stop_outputs = [], [], []
-
-        # Run the decoder loop
-        for t in range(0, steps, self.r):
-            prenet_in = mel_outputs[-1][:, :, -1] if t > 0 else go_frame
-            mel_frames, scores, hidden_states, cell_states, context_vec, stop_tokens = \
-            self.decoder(encoder_seq, encoder_seq_proj, prenet_in,
-                         hidden_states, cell_states, context_vec, t, x)
-            mel_outputs.append(mel_frames)
-            attn_scores.append(scores)
-            stop_outputs.extend([stop_tokens] * self.r)
-            # Stop the loop when all stop tokens in batch exceed threshold
-            if (stop_tokens > 0.5).all() and t > 10: break
-
-        # Concat the mel outputs into sequence
-        mel_outputs = torch.cat(mel_outputs, dim=2)
-
-        # Post-Process for Linear Spectrograms
-        postnet_out = self.postnet(mel_outputs)
-        linear = self.post_proj(postnet_out)
-
-
-        linear = linear.transpose(1, 2)
-
-        # For easy visualisation
-        attn_scores = torch.cat(attn_scores, 1)
-        stop_outputs = torch.cat(stop_outputs, 1)
-
-        self.train()
-
-        return mel_outputs, linear, attn_scores
-
-    def init_model(self):
-        for p in self.parameters():
-            if p.dim() > 1: nn.init.xavier_uniform_(p)
-
-    def get_step(self):
-        return self.step.data.item()
-
-    def reset_step(self):
-        # assignment to parameters or buffers is overloaded, updates internal dict entry
-        self.step = self.step.data.new_tensor(1)
-
-    def log(self, path, msg):
-        with open(path, "a") as f:
-            print(msg, file=f)
-
-    def load(self, path, optimizer=None):
-        # Use device of model params as location for loaded state
-        device = next(self.parameters()).device
-        checkpoint = torch.load(str(path), map_location=device)
-        self.load_state_dict(checkpoint["model_state"])
-
-        if "optimizer_state" in checkpoint and optimizer is not None:
-            optimizer.load_state_dict(checkpoint["optimizer_state"])
-
-    def save(self, path, optimizer=None):
-        if optimizer is not None:
-            torch.save({
-                "model_state": self.state_dict(),
-                "optimizer_state": optimizer.state_dict(),
-            }, str(path))
-        else:
-            torch.save({
-                "model_state": self.state_dict(),
-            }, str(path))
-
-
-    def num_params(self, print_out=True):
-        parameters = filter(lambda p: p.requires_grad, self.parameters())
-        parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
-        if print_out:
-            print("Trainable Parameters: %.3fM" % parameters)
-        return parameters
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from pathlib import Path
+from typing import Union
+
+
+class HighwayNetwork(nn.Module):
+    def __init__(self, size):
+        super().__init__()
+        self.W1 = nn.Linear(size, size)
+        self.W2 = nn.Linear(size, size)
+        self.W1.bias.data.fill_(0.)
+
+    def forward(self, x):
+        x1 = self.W1(x)
+        x2 = self.W2(x)
+        g = torch.sigmoid(x2)
+        y = g * F.relu(x1) + (1. - g) * x
+        return y
+
+
+class Encoder(nn.Module):
+    def __init__(self, embed_dims, num_chars, encoder_dims, K, num_highways, dropout):
+        super().__init__()
+        prenet_dims = (encoder_dims, encoder_dims)
+        cbhg_channels = encoder_dims
+        self.embedding = nn.Embedding(num_chars, embed_dims)
+        self.pre_net = PreNet(embed_dims, fc1_dims=prenet_dims[0], fc2_dims=prenet_dims[1],
+                              dropout=dropout)
+        self.cbhg = CBHG(K=K, in_channels=cbhg_channels, channels=cbhg_channels,
+                         proj_channels=[cbhg_channels, cbhg_channels],
+                         num_highways=num_highways)
+
+    def forward(self, x, speaker_embedding=None):
+        x = self.embedding(x)
+        x = self.pre_net(x)
+        x.transpose_(1, 2)
+        x = self.cbhg(x)
+        if speaker_embedding is not None:
+            x = self.add_speaker_embedding(x, speaker_embedding)
+        return x
+
+    def add_speaker_embedding(self, x, speaker_embedding):
+        # SV2TTS
+        # The input x is the encoder output and is a 3D tensor with size (batch_size, num_chars, tts_embed_dims)
+        # When training, speaker_embedding is also a 2D tensor with size (batch_size, speaker_embedding_size)
+        #     (for inference, speaker_embedding is a 1D tensor with size (speaker_embedding_size))
+        # This concats the speaker embedding for each char in the encoder output
+
+        # Save the dimensions as human-readable names
+        batch_size = x.size()[0]
+        num_chars = x.size()[1]
+
+        if speaker_embedding.dim() == 1:
+            idx = 0
+        else:
+            idx = 1
+
+        # Start by making a copy of each speaker embedding to match the input text length
+        # The output of this has size (batch_size, num_chars * tts_embed_dims)
+        speaker_embedding_size = speaker_embedding.size()[idx]
+        e = speaker_embedding.repeat_interleave(num_chars, dim=idx)
+
+        # Reshape it and transpose
+        e = e.reshape(batch_size, speaker_embedding_size, num_chars)
+        e = e.transpose(1, 2)
+
+        # Concatenate the tiled speaker embedding with the encoder output
+        x = torch.cat((x, e), 2)
+        return x
+
+
+class BatchNormConv(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel, relu=True):
+        super().__init__()
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel, stride=1, padding=kernel // 2, bias=False)
+        self.bnorm = nn.BatchNorm1d(out_channels)
+        self.relu = relu
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = F.relu(x) if self.relu is True else x
+        return self.bnorm(x)
+
+
+class CBHG(nn.Module):
+    def __init__(self, K, in_channels, channels, proj_channels, num_highways):
+        super().__init__()
+
+        # List of all rnns to call `flatten_parameters()` on
+        self._to_flatten = []
+
+        self.bank_kernels = [i for i in range(1, K + 1)]
+        self.conv1d_bank = nn.ModuleList()
+        for k in self.bank_kernels:
+            conv = BatchNormConv(in_channels, channels, k)
+            self.conv1d_bank.append(conv)
+
+        self.maxpool = nn.MaxPool1d(kernel_size=2, stride=1, padding=1)
+
+        self.conv_project1 = BatchNormConv(len(self.bank_kernels) * channels, proj_channels[0], 3)
+        self.conv_project2 = BatchNormConv(proj_channels[0], proj_channels[1], 3, relu=False)
+
+        # Fix the highway input if necessary
+        if proj_channels[-1] != channels:
+            self.highway_mismatch = True
+            self.pre_highway = nn.Linear(proj_channels[-1], channels, bias=False)
+        else:
+            self.highway_mismatch = False
+
+        self.highways = nn.ModuleList()
+        for i in range(num_highways):
+            hn = HighwayNetwork(channels)
+            self.highways.append(hn)
+
+        self.rnn = nn.GRU(channels, channels // 2, batch_first=True, bidirectional=True)
+        self._to_flatten.append(self.rnn)
+
+        # Avoid fragmentation of RNN parameters and associated warning
+        self._flatten_parameters()
+
+    def forward(self, x):
+        # Although we `_flatten_parameters()` on init, when using DataParallel
+        # the model gets replicated, making it no longer guaranteed that the
+        # weights are contiguous in GPU memory. Hence, we must call it again
+        self._flatten_parameters()
+
+        # Save these for later
+        residual = x
+        seq_len = x.size(-1)
+        conv_bank = []
+
+        # Convolution Bank
+        for conv in self.conv1d_bank:
+            c = conv(x) # Convolution
+            conv_bank.append(c[:, :, :seq_len])
+
+        # Stack along the channel axis
+        conv_bank = torch.cat(conv_bank, dim=1)
+
+        # dump the last padding to fit residual
+        x = self.maxpool(conv_bank)[:, :, :seq_len]
+
+        # Conv1d projections
+        x = self.conv_project1(x)
+        x = self.conv_project2(x)
+
+        # Residual Connect
+        x = x + residual
+
+        # Through the highways
+        x = x.transpose(1, 2)
+        if self.highway_mismatch is True:
+            x = self.pre_highway(x)
+        for h in self.highways: x = h(x)
+
+        # And then the RNN
+        x, _ = self.rnn(x)
+        return x
+
+    def _flatten_parameters(self):
+        """Calls `flatten_parameters` on all the rnns used by the WaveRNN. Used
+        to improve efficiency and avoid PyTorch yelling at us."""
+        [m.flatten_parameters() for m in self._to_flatten]
+
+class PreNet(nn.Module):
+    def __init__(self, in_dims, fc1_dims=256, fc2_dims=128, dropout=0.5):
+        super().__init__()
+        self.fc1 = nn.Linear(in_dims, fc1_dims)
+        self.fc2 = nn.Linear(fc1_dims, fc2_dims)
+        self.p = dropout
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = F.relu(x)
+        x = F.dropout(x, self.p, training=True)
+        x = self.fc2(x)
+        x = F.relu(x)
+        x = F.dropout(x, self.p, training=True)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, attn_dims):
+        super().__init__()
+        self.W = nn.Linear(attn_dims, attn_dims, bias=False)
+        self.v = nn.Linear(attn_dims, 1, bias=False)
+
+    def forward(self, encoder_seq_proj, query, t):
+
+        # print(encoder_seq_proj.shape)
+        # Transform the query vector
+        query_proj = self.W(query).unsqueeze(1)
+
+        # Compute the scores
+        u = self.v(torch.tanh(encoder_seq_proj + query_proj))
+        scores = F.softmax(u, dim=1)
+
+        return scores.transpose(1, 2)
+
+
+class LSA(nn.Module):
+    def __init__(self, attn_dim, kernel_size=31, filters=32):
+        super().__init__()
+        self.conv = nn.Conv1d(1, filters, padding=(kernel_size - 1) // 2, kernel_size=kernel_size, bias=True)
+        self.L = nn.Linear(filters, attn_dim, bias=False)
+        self.W = nn.Linear(attn_dim, attn_dim, bias=True) # Include the attention bias in this term
+        self.v = nn.Linear(attn_dim, 1, bias=False)
+        self.cumulative = None
+        self.attention = None
+
+    def init_attention(self, encoder_seq_proj):
+        device = next(self.parameters()).device  # use same device as parameters
+        b, t, c = encoder_seq_proj.size()
+        self.cumulative = torch.zeros(b, t, device=device)
+        self.attention = torch.zeros(b, t, device=device)
+
+    def forward(self, encoder_seq_proj, query, t, chars):
+
+        if t == 0: self.init_attention(encoder_seq_proj)
+
+        processed_query = self.W(query).unsqueeze(1)
+
+        location = self.cumulative.unsqueeze(1)
+        processed_loc = self.L(self.conv(location).transpose(1, 2))
+
+        u = self.v(torch.tanh(processed_query + encoder_seq_proj + processed_loc))
+        u = u.squeeze(-1)
+
+        # Mask zero padding chars
+        u = u * (chars != 0).float()
+
+        # Smooth Attention
+        # scores = torch.sigmoid(u) / torch.sigmoid(u).sum(dim=1, keepdim=True)
+        scores = F.softmax(u, dim=1)
+        self.attention = scores
+        self.cumulative = self.cumulative + self.attention
+
+        return scores.unsqueeze(-1).transpose(1, 2)
+
+
+class Decoder(nn.Module):
+    # Class variable because its value doesn't change between classes
+    # yet ought to be scoped by class because its a property of a Decoder
+    max_r = 20
+    def __init__(self, n_mels, encoder_dims, decoder_dims, lstm_dims,
+                 dropout, speaker_embedding_size):
+        super().__init__()
+        self.register_buffer("r", torch.tensor(1, dtype=torch.int))
+        self.n_mels = n_mels
+        prenet_dims = (decoder_dims * 2, decoder_dims * 2)
+        self.prenet = PreNet(n_mels, fc1_dims=prenet_dims[0], fc2_dims=prenet_dims[1],
+                             dropout=dropout)
+        self.attn_net = LSA(decoder_dims)
+        self.attn_rnn = nn.GRUCell(encoder_dims + prenet_dims[1] + speaker_embedding_size, decoder_dims)
+        self.rnn_input = nn.Linear(encoder_dims + decoder_dims + speaker_embedding_size, lstm_dims)
+        self.res_rnn1 = nn.LSTMCell(lstm_dims, lstm_dims)
+        self.res_rnn2 = nn.LSTMCell(lstm_dims, lstm_dims)
+        self.mel_proj = nn.Linear(lstm_dims, n_mels * self.max_r, bias=False)
+        self.stop_proj = nn.Linear(encoder_dims + speaker_embedding_size + lstm_dims, 1)
+
+    def zoneout(self, prev, current, p=0.1):
+        device = next(self.parameters()).device  # Use same device as parameters
+        mask = torch.zeros(prev.size(), device=device).bernoulli_(p)
+        return prev * mask + current * (1 - mask)
+
+    def forward(self, encoder_seq, encoder_seq_proj, prenet_in,
+                hidden_states, cell_states, context_vec, t, chars):
+
+        # Need this for reshaping mels
+        batch_size = encoder_seq.size(0)
+
+        # Unpack the hidden and cell states
+        attn_hidden, rnn1_hidden, rnn2_hidden = hidden_states
+        rnn1_cell, rnn2_cell = cell_states
+
+        # PreNet for the Attention RNN
+        prenet_out = self.prenet(prenet_in)
+
+        # Compute the Attention RNN hidden state
+        attn_rnn_in = torch.cat([context_vec, prenet_out], dim=-1)
+        attn_hidden = self.attn_rnn(attn_rnn_in.squeeze(1), attn_hidden)
+
+        # Compute the attention scores
+        scores = self.attn_net(encoder_seq_proj, attn_hidden, t, chars)
+
+        # Dot product to create the context vector
+        context_vec = scores @ encoder_seq
+        context_vec = context_vec.squeeze(1)
+
+        # Concat Attention RNN output w. Context Vector & project
+        x = torch.cat([context_vec, attn_hidden], dim=1)
+        x = self.rnn_input(x)
+
+        # Compute first Residual RNN
+        rnn1_hidden_next, rnn1_cell = self.res_rnn1(x, (rnn1_hidden, rnn1_cell))
+        if self.training:
+            rnn1_hidden = self.zoneout(rnn1_hidden, rnn1_hidden_next)
+        else:
+            rnn1_hidden = rnn1_hidden_next
+        x = x + rnn1_hidden
+
+        # Compute second Residual RNN
+        rnn2_hidden_next, rnn2_cell = self.res_rnn2(x, (rnn2_hidden, rnn2_cell))
+        if self.training:
+            rnn2_hidden = self.zoneout(rnn2_hidden, rnn2_hidden_next)
+        else:
+            rnn2_hidden = rnn2_hidden_next
+        x = x + rnn2_hidden
+
+        # Project Mels
+        mels = self.mel_proj(x)
+        mels = mels.view(batch_size, self.n_mels, self.max_r)[:, :, :self.r]
+        hidden_states = (attn_hidden, rnn1_hidden, rnn2_hidden)
+        cell_states = (rnn1_cell, rnn2_cell)
+
+        # Stop token prediction
+        s = torch.cat((x, context_vec), dim=1)
+        s = self.stop_proj(s)
+        stop_tokens = torch.sigmoid(s)
+
+        return mels, scores, hidden_states, cell_states, context_vec, stop_tokens
+
+
+class Tacotron(nn.Module):
+    def __init__(self, embed_dims, num_chars, encoder_dims, decoder_dims, n_mels, 
+                 fft_bins, postnet_dims, encoder_K, lstm_dims, postnet_K, num_highways,
+                 dropout, stop_threshold, speaker_embedding_size):
+        super().__init__()
+        self.n_mels = n_mels
+        self.lstm_dims = lstm_dims
+        self.encoder_dims = encoder_dims
+        self.decoder_dims = decoder_dims
+        self.speaker_embedding_size = speaker_embedding_size
+        self.encoder = Encoder(embed_dims, num_chars, encoder_dims,
+                               encoder_K, num_highways, dropout)
+        self.encoder_proj = nn.Linear(encoder_dims + speaker_embedding_size, decoder_dims, bias=False)
+        self.decoder = Decoder(n_mels, encoder_dims, decoder_dims, lstm_dims,
+                               dropout, speaker_embedding_size)
+        self.postnet = CBHG(postnet_K, n_mels, postnet_dims,
+                            [postnet_dims, fft_bins], num_highways)
+        self.post_proj = nn.Linear(postnet_dims, fft_bins, bias=False)
+
+        self.init_model()
+        self.num_params()
+
+        self.register_buffer("step", torch.zeros(1, dtype=torch.long))
+        self.register_buffer("stop_threshold", torch.tensor(stop_threshold, dtype=torch.float32))
+
+    @property
+    def r(self):
+        return self.decoder.r.item()
+
+    @r.setter
+    def r(self, value):
+        self.decoder.r = self.decoder.r.new_tensor(value, requires_grad=False)
+
+    def forward(self, x, m, speaker_embedding):
+        device = next(self.parameters()).device  # use same device as parameters
+
+        self.step += 1
+        batch_size, _, steps  = m.size()
+
+        # Initialise all hidden states and pack into tuple
+        attn_hidden = torch.zeros(batch_size, self.decoder_dims, device=device)
+        rnn1_hidden = torch.zeros(batch_size, self.lstm_dims, device=device)
+        rnn2_hidden = torch.zeros(batch_size, self.lstm_dims, device=device)
+        hidden_states = (attn_hidden, rnn1_hidden, rnn2_hidden)
+
+        # Initialise all lstm cell states and pack into tuple
+        rnn1_cell = torch.zeros(batch_size, self.lstm_dims, device=device)
+        rnn2_cell = torch.zeros(batch_size, self.lstm_dims, device=device)
+        cell_states = (rnn1_cell, rnn2_cell)
+
+        # <GO> Frame for start of decoder loop
+        go_frame = torch.zeros(batch_size, self.n_mels, device=device)
+
+        # Need an initial context vector
+        context_vec = torch.zeros(batch_size, self.encoder_dims + self.speaker_embedding_size, device=device)
+
+        # SV2TTS: Run the encoder with the speaker embedding
+        # The projection avoids unnecessary matmuls in the decoder loop
+        encoder_seq = self.encoder(x, speaker_embedding)
+        encoder_seq_proj = self.encoder_proj(encoder_seq)
+
+        # Need a couple of lists for outputs
+        mel_outputs, attn_scores, stop_outputs = [], [], []
+
+        # Run the decoder loop
+        for t in range(0, steps, self.r):
+            prenet_in = m[:, :, t - 1] if t > 0 else go_frame
+            mel_frames, scores, hidden_states, cell_states, context_vec, stop_tokens = \
+                self.decoder(encoder_seq, encoder_seq_proj, prenet_in,
+                             hidden_states, cell_states, context_vec, t, x)
+            mel_outputs.append(mel_frames)
+            attn_scores.append(scores)
+            stop_outputs.extend([stop_tokens] * self.r)
+
+        # Concat the mel outputs into sequence
+        mel_outputs = torch.cat(mel_outputs, dim=2)
+
+        # Post-Process for Linear Spectrograms
+        postnet_out = self.postnet(mel_outputs)
+        linear = self.post_proj(postnet_out)
+        linear = linear.transpose(1, 2)
+
+        # For easy visualisation
+        attn_scores = torch.cat(attn_scores, 1)
+        # attn_scores = attn_scores.cpu().data.numpy()
+        stop_outputs = torch.cat(stop_outputs, 1)
+
+        return mel_outputs, linear, attn_scores, stop_outputs
+
+    def generate(self, x, speaker_embedding=None, steps=2000):
+        self.eval()
+        device = next(self.parameters()).device  # use same device as parameters
+
+        batch_size, _  = x.size()
+
+        # Need to initialise all hidden states and pack into tuple for tidyness
+        attn_hidden = torch.zeros(batch_size, self.decoder_dims, device=device)
+        rnn1_hidden = torch.zeros(batch_size, self.lstm_dims, device=device)
+        rnn2_hidden = torch.zeros(batch_size, self.lstm_dims, device=device)
+        hidden_states = (attn_hidden, rnn1_hidden, rnn2_hidden)
+
+        # Need to initialise all lstm cell states and pack into tuple for tidyness
+        rnn1_cell = torch.zeros(batch_size, self.lstm_dims, device=device)
+        rnn2_cell = torch.zeros(batch_size, self.lstm_dims, device=device)
+        cell_states = (rnn1_cell, rnn2_cell)
+
+        # Need a <GO> Frame for start of decoder loop
+        go_frame = torch.zeros(batch_size, self.n_mels, device=device)
+
+        # Need an initial context vector
+        context_vec = torch.zeros(batch_size, self.encoder_dims + self.speaker_embedding_size, device=device)
+
+        # SV2TTS: Run the encoder with the speaker embedding
+        # The projection avoids unnecessary matmuls in the decoder loop
+        encoder_seq = self.encoder(x, speaker_embedding)
+        encoder_seq_proj = self.encoder_proj(encoder_seq)
+
+        # Need a couple of lists for outputs
+        mel_outputs, attn_scores, stop_outputs = [], [], []
+
+        # Run the decoder loop
+        for t in range(0, steps, self.r):
+            prenet_in = mel_outputs[-1][:, :, -1] if t > 0 else go_frame
+            mel_frames, scores, hidden_states, cell_states, context_vec, stop_tokens = \
+            self.decoder(encoder_seq, encoder_seq_proj, prenet_in,
+                         hidden_states, cell_states, context_vec, t, x)
+            mel_outputs.append(mel_frames)
+            attn_scores.append(scores)
+            stop_outputs.extend([stop_tokens] * self.r)
+            # Stop the loop when all stop tokens in batch exceed threshold
+            if (stop_tokens > 0.5).all() and t > 10: break
+
+        # Concat the mel outputs into sequence
+        mel_outputs = torch.cat(mel_outputs, dim=2)
+
+        # Post-Process for Linear Spectrograms
+        postnet_out = self.postnet(mel_outputs)
+        linear = self.post_proj(postnet_out)
+
+
+        linear = linear.transpose(1, 2)
+
+        # For easy visualisation
+        attn_scores = torch.cat(attn_scores, 1)
+        stop_outputs = torch.cat(stop_outputs, 1)
+
+        self.train()
+
+        return mel_outputs, linear, attn_scores
+
+    def init_model(self):
+        for p in self.parameters():
+            if p.dim() > 1: nn.init.xavier_uniform_(p)
+
+    def get_step(self):
+        return self.step.data.item()
+
+    def reset_step(self):
+        # assignment to parameters or buffers is overloaded, updates internal dict entry
+        self.step = self.step.data.new_tensor(1)
+
+    def log(self, path, msg):
+        with open(path, "a") as f:
+            print(msg, file=f)
+
+    def load(self, path, optimizer=None):
+        # Use device of model params as location for loaded state
+        device = next(self.parameters()).device
+        checkpoint = torch.load(str(path), map_location=device, weights_only=False)
+        self.load_state_dict(checkpoint["model_state"])
+
+        if "optimizer_state" in checkpoint and optimizer is not None:
+            optimizer.load_state_dict(checkpoint["optimizer_state"])
+
+    def save(self, path, optimizer=None):
+        if optimizer is not None:
+            torch.save({
+                "model_state": self.state_dict(),
+                "optimizer_state": optimizer.state_dict(),
+            }, str(path))
+        else:
+            torch.save({
+                "model_state": self.state_dict(),
+            }, str(path))
+
+
+    def num_params(self, print_out=True):
+        parameters = filter(lambda p: p.requires_grad, self.parameters())
+        parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
+        if print_out:
+            print("Trainable Parameters: %.3fM" % parameters)
+        return parameters

utils/default_models.py

@@ -1,57 +1,25 @@
-import urllib.request
 from pathlib import Path
-from threading import Thread
-from urllib.error import HTTPError
-
-from tqdm import tqdm
 
+from huggingface_hub import hf_hub_download
 
+# The original Google Drive links are dead; the same pretrained SV2TTS checkpoints
+# (identical byte sizes) are mirrored on the Hub.
+MODEL_REPO = "CorentinJ/SV2TTS"
 default_models = {
-    "encoder": ("https://drive.google.com/uc?export=download&id=1q8mEGwCkFy23KZsinbuvdKAQLqNKbYf1", 17090379),
-    #"synthesizer": ("https://drive.google.com/u/0/uc?id=1EqFMIbvxffxtjiVrtykroF6_mUh-5Z3s&export=download&confirm=t", 370554559),
-    "synthesizer": ("https://sourceforge.net/projects/encoders/files/synthesizer.pt/download", 370554559),
-    "vocoder": ("https://drive.google.com/uc?export=download&id=1cf2NO6FtI0jDuy8AV3Xgn6leO6dHjIgu", 53845290),
+    "encoder": 17090379,
+    "synthesizer": 370554559,
+    "vocoder": 53845290,
 }
 
 
-class DownloadProgressBar(tqdm):
-    def update_to(self, b=1, bsize=1, tsize=None):
-        if tsize is not None:
-            self.total = tsize
-        self.update(b * bsize - self.n)
-
-
-def download(url: str, target: Path, bar_pos=0):
-    # Ensure the directory exists
-    target.parent.mkdir(exist_ok=True, parents=True)
-
-    desc = f"Downloading {target.name}"
-    with DownloadProgressBar(unit="B", unit_scale=True, miniters=1, desc=desc, position=bar_pos, leave=False) as t:
-        try:
-            urllib.request.urlretrieve(url, filename=target, reporthook=t.update_to)
-        except HTTPError:
-            return
-
-
 def ensure_default_models(models_dir: Path):
-    # Define download tasks
-    jobs = []
-    for model_name, (url, size) in default_models.items():
-        target_path = models_dir / "default" / f"{model_name}.pt"
-        if target_path.exists():
-            if target_path.stat().st_size != size:
-                print(f"File {target_path} is not of expected size, redownloading...")
-            else:
-                continue
-
-        thread = Thread(target=download, args=(url, target_path, len(jobs)))
-        thread.start()
-        jobs.append((thread, target_path, size))
-
-    # Run and join threads
-    for thread, target_path, size in jobs:
-        thread.join()
-
+    target_dir = Path(models_dir) / "default"
+    target_dir.mkdir(exist_ok=True, parents=True)
+    for model_name, size in default_models.items():
+        target_path = target_dir / f"{model_name}.pt"
+        if target_path.exists() and target_path.stat().st_size == size:
+            continue
+        print(f"Downloading {model_name}.pt from {MODEL_REPO}...")
+        hf_hub_download(MODEL_REPO, f"{model_name}.pt", local_dir=str(target_dir))
         assert target_path.exists() and target_path.stat().st_size == size, \
-            f"Download for {target_path.name} failed. You may download models manually instead.\n" \
-            f"https://drive.google.com/drive/folders/1fU6umc5uQAVR2udZdHX-lDgXYzTyqG_j"
+            f"Download for {target_path.name} failed."

vocoder/audio.py

@@ -1,108 +1,108 @@
-import math
-import numpy as np
-import librosa
-import vocoder.hparams as hp
-from scipy.signal import lfilter
-import soundfile as sf
-
-
-def label_2_float(x, bits) :
-    return 2 * x / (2**bits - 1.) - 1.
-
-
-def float_2_label(x, bits) :
-    assert abs(x).max() <= 1.0
-    x = (x + 1.) * (2**bits - 1) / 2
-    return x.clip(0, 2**bits - 1)
-
-
-def load_wav(path) :
-    return librosa.load(str(path), sr=hp.sample_rate)[0]
-
-
-def save_wav(x, path) :
-    sf.write(path, x.astype(np.float32), hp.sample_rate)
-
-
-def split_signal(x) :
-    unsigned = x + 2**15
-    coarse = unsigned // 256
-    fine = unsigned % 256
-    return coarse, fine
-
-
-def combine_signal(coarse, fine) :
-    return coarse * 256 + fine - 2**15
-
-
-def encode_16bits(x) :
-    return np.clip(x * 2**15, -2**15, 2**15 - 1).astype(np.int16)
-
-
-mel_basis = None
-
-
-def linear_to_mel(spectrogram):
-    global mel_basis
-    if mel_basis is None:
-        mel_basis = build_mel_basis()
-    return np.dot(mel_basis, spectrogram)
-
-
-def build_mel_basis():
-    return librosa.filters.mel(hp.sample_rate, hp.n_fft, n_mels=hp.num_mels, fmin=hp.fmin)
-
-
-def normalize(S):
-    return np.clip((S - hp.min_level_db) / -hp.min_level_db, 0, 1)
-
-
-def denormalize(S):
-    return (np.clip(S, 0, 1) * -hp.min_level_db) + hp.min_level_db
-
-
-def amp_to_db(x):
-    return 20 * np.log10(np.maximum(1e-5, x))
-
-
-def db_to_amp(x):
-    return np.power(10.0, x * 0.05)
-
-
-def spectrogram(y):
-    D = stft(y)
-    S = amp_to_db(np.abs(D)) - hp.ref_level_db
-    return normalize(S)
-
-
-def melspectrogram(y):
-    D = stft(y)
-    S = amp_to_db(linear_to_mel(np.abs(D)))
-    return normalize(S)
-
-
-def stft(y):
-    return librosa.stft(y=y, n_fft=hp.n_fft, hop_length=hp.hop_length, win_length=hp.win_length)
-
-
-def pre_emphasis(x):
-    return lfilter([1, -hp.preemphasis], [1], x)
-
-
-def de_emphasis(x):
-    return lfilter([1], [1, -hp.preemphasis], x)
-
-
-def encode_mu_law(x, mu) :
-    mu = mu - 1
-    fx = np.sign(x) * np.log(1 + mu * np.abs(x)) / np.log(1 + mu)
-    return np.floor((fx + 1) / 2 * mu + 0.5)
-
-
-def decode_mu_law(y, mu, from_labels=True) :
-    if from_labels: 
-        y = label_2_float(y, math.log2(mu))
-    mu = mu - 1
-    x = np.sign(y) / mu * ((1 + mu) ** np.abs(y) - 1)
-    return x
-
+import math
+import numpy as np
+import librosa
+import vocoder.hparams as hp
+from scipy.signal import lfilter
+import soundfile as sf
+
+
+def label_2_float(x, bits) :
+    return 2 * x / (2**bits - 1.) - 1.
+
+
+def float_2_label(x, bits) :
+    assert abs(x).max() <= 1.0
+    x = (x + 1.) * (2**bits - 1) / 2
+    return x.clip(0, 2**bits - 1)
+
+
+def load_wav(path) :
+    return librosa.load(str(path), sr=hp.sample_rate)[0]
+
+
+def save_wav(x, path) :
+    sf.write(path, x.astype(np.float32), hp.sample_rate)
+
+
+def split_signal(x) :
+    unsigned = x + 2**15
+    coarse = unsigned // 256
+    fine = unsigned % 256
+    return coarse, fine
+
+
+def combine_signal(coarse, fine) :
+    return coarse * 256 + fine - 2**15
+
+
+def encode_16bits(x) :
+    return np.clip(x * 2**15, -2**15, 2**15 - 1).astype(np.int16)
+
+
+mel_basis = None
+
+
+def linear_to_mel(spectrogram):
+    global mel_basis
+    if mel_basis is None:
+        mel_basis = build_mel_basis()
+    return np.dot(mel_basis, spectrogram)
+
+
+def build_mel_basis():
+    return librosa.filters.mel(sr=hp.sample_rate, n_fft=hp.n_fft, n_mels=hp.num_mels, fmin=hp.fmin)
+
+
+def normalize(S):
+    return np.clip((S - hp.min_level_db) / -hp.min_level_db, 0, 1)
+
+
+def denormalize(S):
+    return (np.clip(S, 0, 1) * -hp.min_level_db) + hp.min_level_db
+
+
+def amp_to_db(x):
+    return 20 * np.log10(np.maximum(1e-5, x))
+
+
+def db_to_amp(x):
+    return np.power(10.0, x * 0.05)
+
+
+def spectrogram(y):
+    D = stft(y)
+    S = amp_to_db(np.abs(D)) - hp.ref_level_db
+    return normalize(S)
+
+
+def melspectrogram(y):
+    D = stft(y)
+    S = amp_to_db(linear_to_mel(np.abs(D)))
+    return normalize(S)
+
+
+def stft(y):
+    return librosa.stft(y=y, n_fft=hp.n_fft, hop_length=hp.hop_length, win_length=hp.win_length)
+
+
+def pre_emphasis(x):
+    return lfilter([1, -hp.preemphasis], [1], x)
+
+
+def de_emphasis(x):
+    return lfilter([1], [1, -hp.preemphasis], x)
+
+
+def encode_mu_law(x, mu) :
+    mu = mu - 1
+    fx = np.sign(x) * np.log(1 + mu * np.abs(x)) / np.log(1 + mu)
+    return np.floor((fx + 1) / 2 * mu + 0.5)
+
+
+def decode_mu_law(y, mu, from_labels=True) :
+    if from_labels: 
+        y = label_2_float(y, math.log2(mu))
+    mu = mu - 1
+    x = np.sign(y) / mu * ((1 + mu) ** np.abs(y) - 1)
+    return x
+

vocoder/inference.py

@@ -1,64 +1,64 @@
-from vocoder.models.fatchord_version import WaveRNN
-from vocoder import hparams as hp
-import torch
-
-
-_model = None   # type: WaveRNN
-
-def load_model(weights_fpath, verbose=True):
-    global _model, _device
-    
-    if verbose:
-        print("Building Wave-RNN")
-    _model = WaveRNN(
-        rnn_dims=hp.voc_rnn_dims,
-        fc_dims=hp.voc_fc_dims,
-        bits=hp.bits,
-        pad=hp.voc_pad,
-        upsample_factors=hp.voc_upsample_factors,
-        feat_dims=hp.num_mels,
-        compute_dims=hp.voc_compute_dims,
-        res_out_dims=hp.voc_res_out_dims,
-        res_blocks=hp.voc_res_blocks,
-        hop_length=hp.hop_length,
-        sample_rate=hp.sample_rate,
-        mode=hp.voc_mode
-    )
-
-    if torch.cuda.is_available():
-        _model = _model.cuda()
-        _device = torch.device('cuda')
-    else:
-        _device = torch.device('cpu')
-    
-    if verbose:
-        print("Loading model weights at %s" % weights_fpath)
-    checkpoint = torch.load(weights_fpath, _device)
-    _model.load_state_dict(checkpoint['model_state'])
-    _model.eval()
-
-
-def is_loaded():
-    return _model is not None
-
-
-def infer_waveform(mel, normalize=True,  batched=True, target=8000, overlap=800, 
-                   progress_callback=None):
-    """
-    Infers the waveform of a mel spectrogram output by the synthesizer (the format must match 
-    that of the synthesizer!)
-    
-    :param normalize:  
-    :param batched: 
-    :param target: 
-    :param overlap: 
-    :return: 
-    """
-    if _model is None:
-        raise Exception("Please load Wave-RNN in memory before using it")
-    
-    if normalize:
-        mel = mel / hp.mel_max_abs_value
-    mel = torch.from_numpy(mel[None, ...])
-    wav = _model.generate(mel, batched, target, overlap, hp.mu_law, progress_callback)
-    return wav
+from vocoder.models.fatchord_version import WaveRNN
+from vocoder import hparams as hp
+import torch
+
+
+_model = None   # type: WaveRNN
+
+def load_model(weights_fpath, verbose=True):
+    global _model, _device
+    
+    if verbose:
+        print("Building Wave-RNN")
+    _model = WaveRNN(
+        rnn_dims=hp.voc_rnn_dims,
+        fc_dims=hp.voc_fc_dims,
+        bits=hp.bits,
+        pad=hp.voc_pad,
+        upsample_factors=hp.voc_upsample_factors,
+        feat_dims=hp.num_mels,
+        compute_dims=hp.voc_compute_dims,
+        res_out_dims=hp.voc_res_out_dims,
+        res_blocks=hp.voc_res_blocks,
+        hop_length=hp.hop_length,
+        sample_rate=hp.sample_rate,
+        mode=hp.voc_mode
+    )
+
+    if torch.cuda.is_available():
+        _model = _model.cuda()
+        _device = torch.device('cuda')
+    else:
+        _device = torch.device('cpu')
+    
+    if verbose:
+        print("Loading model weights at %s" % weights_fpath)
+    checkpoint = torch.load(weights_fpath, map_location=_device, weights_only=False)
+    _model.load_state_dict(checkpoint['model_state'])
+    _model.eval()
+
+
+def is_loaded():
+    return _model is not None
+
+
+def infer_waveform(mel, normalize=True,  batched=True, target=8000, overlap=800, 
+                   progress_callback=None):
+    """
+    Infers the waveform of a mel spectrogram output by the synthesizer (the format must match 
+    that of the synthesizer!)
+    
+    :param normalize:  
+    :param batched: 
+    :param target: 
+    :param overlap: 
+    :return: 
+    """
+    if _model is None:
+        raise Exception("Please load Wave-RNN in memory before using it")
+    
+    if normalize:
+        mel = mel / hp.mel_max_abs_value
+    mel = torch.from_numpy(mel[None, ...])
+    wav = _model.generate(mel, batched, target, overlap, hp.mu_law, progress_callback)
+    return wav

vocoder/models/fatchord_version.py

@@ -1,434 +1,434 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from vocoder.distribution import sample_from_discretized_mix_logistic
-from vocoder.display import *
-from vocoder.audio import *
-
-
-class ResBlock(nn.Module):
-    def __init__(self, dims):
-        super().__init__()
-        self.conv1 = nn.Conv1d(dims, dims, kernel_size=1, bias=False)
-        self.conv2 = nn.Conv1d(dims, dims, kernel_size=1, bias=False)
-        self.batch_norm1 = nn.BatchNorm1d(dims)
-        self.batch_norm2 = nn.BatchNorm1d(dims)
-
-    def forward(self, x):
-        residual = x
-        x = self.conv1(x)
-        x = self.batch_norm1(x)
-        x = F.relu(x)
-        x = self.conv2(x)
-        x = self.batch_norm2(x)
-        return x + residual
-
-
-class MelResNet(nn.Module):
-    def __init__(self, res_blocks, in_dims, compute_dims, res_out_dims, pad):
-        super().__init__()
-        k_size = pad * 2 + 1
-        self.conv_in = nn.Conv1d(in_dims, compute_dims, kernel_size=k_size, bias=False)
-        self.batch_norm = nn.BatchNorm1d(compute_dims)
-        self.layers = nn.ModuleList()
-        for i in range(res_blocks):
-            self.layers.append(ResBlock(compute_dims))
-        self.conv_out = nn.Conv1d(compute_dims, res_out_dims, kernel_size=1)
-
-    def forward(self, x):
-        x = self.conv_in(x)
-        x = self.batch_norm(x)
-        x = F.relu(x)
-        for f in self.layers: x = f(x)
-        x = self.conv_out(x)
-        return x
-
-
-class Stretch2d(nn.Module):
-    def __init__(self, x_scale, y_scale):
-        super().__init__()
-        self.x_scale = x_scale
-        self.y_scale = y_scale
-
-    def forward(self, x):
-        b, c, h, w = x.size()
-        x = x.unsqueeze(-1).unsqueeze(3)
-        x = x.repeat(1, 1, 1, self.y_scale, 1, self.x_scale)
-        return x.view(b, c, h * self.y_scale, w * self.x_scale)
-
-
-class UpsampleNetwork(nn.Module):
-    def __init__(self, feat_dims, upsample_scales, compute_dims,
-                 res_blocks, res_out_dims, pad):
-        super().__init__()
-        total_scale = np.cumproduct(upsample_scales)[-1]
-        self.indent = pad * total_scale
-        self.resnet = MelResNet(res_blocks, feat_dims, compute_dims, res_out_dims, pad)
-        self.resnet_stretch = Stretch2d(total_scale, 1)
-        self.up_layers = nn.ModuleList()
-        for scale in upsample_scales:
-            k_size = (1, scale * 2 + 1)
-            padding = (0, scale)
-            stretch = Stretch2d(scale, 1)
-            conv = nn.Conv2d(1, 1, kernel_size=k_size, padding=padding, bias=False)
-            conv.weight.data.fill_(1. / k_size[1])
-            self.up_layers.append(stretch)
-            self.up_layers.append(conv)
-
-    def forward(self, m):
-        aux = self.resnet(m).unsqueeze(1)
-        aux = self.resnet_stretch(aux)
-        aux = aux.squeeze(1)
-        m = m.unsqueeze(1)
-        for f in self.up_layers: m = f(m)
-        m = m.squeeze(1)[:, :, self.indent:-self.indent]
-        return m.transpose(1, 2), aux.transpose(1, 2)
-
-
-class WaveRNN(nn.Module):
-    def __init__(self, rnn_dims, fc_dims, bits, pad, upsample_factors,
-                 feat_dims, compute_dims, res_out_dims, res_blocks,
-                 hop_length, sample_rate, mode='RAW'):
-        super().__init__()
-        self.mode = mode
-        self.pad = pad
-        if self.mode == 'RAW' :
-            self.n_classes = 2 ** bits
-        elif self.mode == 'MOL' :
-            self.n_classes = 30
-        else :
-            RuntimeError("Unknown model mode value - ", self.mode)
-
-        self.rnn_dims = rnn_dims
-        self.aux_dims = res_out_dims // 4
-        self.hop_length = hop_length
-        self.sample_rate = sample_rate
-
-        self.upsample = UpsampleNetwork(feat_dims, upsample_factors, compute_dims, res_blocks, res_out_dims, pad)
-        self.I = nn.Linear(feat_dims + self.aux_dims + 1, rnn_dims)
-        self.rnn1 = nn.GRU(rnn_dims, rnn_dims, batch_first=True)
-        self.rnn2 = nn.GRU(rnn_dims + self.aux_dims, rnn_dims, batch_first=True)
-        self.fc1 = nn.Linear(rnn_dims + self.aux_dims, fc_dims)
-        self.fc2 = nn.Linear(fc_dims + self.aux_dims, fc_dims)
-        self.fc3 = nn.Linear(fc_dims, self.n_classes)
-
-        self.step = nn.Parameter(torch.zeros(1).long(), requires_grad=False)
-        self.num_params()
-
-    def forward(self, x, mels):
-        self.step += 1
-        bsize = x.size(0)
-        if torch.cuda.is_available():
-            h1 = torch.zeros(1, bsize, self.rnn_dims).cuda()
-            h2 = torch.zeros(1, bsize, self.rnn_dims).cuda()
-        else:
-            h1 = torch.zeros(1, bsize, self.rnn_dims).cpu()
-            h2 = torch.zeros(1, bsize, self.rnn_dims).cpu()
-        mels, aux = self.upsample(mels)
-
-        aux_idx = [self.aux_dims * i for i in range(5)]
-        a1 = aux[:, :, aux_idx[0]:aux_idx[1]]
-        a2 = aux[:, :, aux_idx[1]:aux_idx[2]]
-        a3 = aux[:, :, aux_idx[2]:aux_idx[3]]
-        a4 = aux[:, :, aux_idx[3]:aux_idx[4]]
-
-        x = torch.cat([x.unsqueeze(-1), mels, a1], dim=2)
-        x = self.I(x)
-        res = x
-        x, _ = self.rnn1(x, h1)
-
-        x = x + res
-        res = x
-        x = torch.cat([x, a2], dim=2)
-        x, _ = self.rnn2(x, h2)
-
-        x = x + res
-        x = torch.cat([x, a3], dim=2)
-        x = F.relu(self.fc1(x))
-
-        x = torch.cat([x, a4], dim=2)
-        x = F.relu(self.fc2(x))
-        return self.fc3(x)
-
-    def generate(self, mels, batched, target, overlap, mu_law, progress_callback=None):
-        mu_law = mu_law if self.mode == 'RAW' else False
-        progress_callback = progress_callback or self.gen_display
-
-        self.eval()
-        output = []
-        start = time.time()
-        rnn1 = self.get_gru_cell(self.rnn1)
-        rnn2 = self.get_gru_cell(self.rnn2)
-
-        with torch.no_grad():
-            if torch.cuda.is_available():
-                mels = mels.cuda()
-            else:
-                mels = mels.cpu()
-            wave_len = (mels.size(-1) - 1) * self.hop_length
-            mels = self.pad_tensor(mels.transpose(1, 2), pad=self.pad, side='both')
-            mels, aux = self.upsample(mels.transpose(1, 2))
-
-            if batched:
-                mels = self.fold_with_overlap(mels, target, overlap)
-                aux = self.fold_with_overlap(aux, target, overlap)
-
-            b_size, seq_len, _ = mels.size()
-
-            if torch.cuda.is_available():
-                h1 = torch.zeros(b_size, self.rnn_dims).cuda()
-                h2 = torch.zeros(b_size, self.rnn_dims).cuda()
-                x = torch.zeros(b_size, 1).cuda()
-            else:
-                h1 = torch.zeros(b_size, self.rnn_dims).cpu()
-                h2 = torch.zeros(b_size, self.rnn_dims).cpu()
-                x = torch.zeros(b_size, 1).cpu()
-
-            d = self.aux_dims
-            aux_split = [aux[:, :, d * i:d * (i + 1)] for i in range(4)]
-
-            for i in range(seq_len):
-
-                m_t = mels[:, i, :]
-
-                a1_t, a2_t, a3_t, a4_t = (a[:, i, :] for a in aux_split)
-
-                x = torch.cat([x, m_t, a1_t], dim=1)
-                x = self.I(x)
-                h1 = rnn1(x, h1)
-
-                x = x + h1
-                inp = torch.cat([x, a2_t], dim=1)
-                h2 = rnn2(inp, h2)
-
-                x = x + h2
-                x = torch.cat([x, a3_t], dim=1)
-                x = F.relu(self.fc1(x))
-
-                x = torch.cat([x, a4_t], dim=1)
-                x = F.relu(self.fc2(x))
-
-                logits = self.fc3(x)
-
-                if self.mode == 'MOL':
-                    sample = sample_from_discretized_mix_logistic(logits.unsqueeze(0).transpose(1, 2))
-                    output.append(sample.view(-1))
-                    if torch.cuda.is_available():
-                        # x = torch.FloatTensor([[sample]]).cuda()
-                        x = sample.transpose(0, 1).cuda()
-                    else:
-                        x = sample.transpose(0, 1)
-
-                elif self.mode == 'RAW' :
-                    posterior = F.softmax(logits, dim=1)
-                    distrib = torch.distributions.Categorical(posterior)
-
-                    sample = 2 * distrib.sample().float() / (self.n_classes - 1.) - 1.
-                    output.append(sample)
-                    x = sample.unsqueeze(-1)
-                else:
-                    raise RuntimeError("Unknown model mode value - ", self.mode)
-
-                if i % 100 == 0:
-                    gen_rate = (i + 1) / (time.time() - start) * b_size / 1000
-                    progress_callback(i, seq_len, b_size, gen_rate)
-
-        output = torch.stack(output).transpose(0, 1)
-        output = output.cpu().numpy()
-        output = output.astype(np.float64)
-        
-        if batched:
-            output = self.xfade_and_unfold(output, target, overlap)
-        else:
-            output = output[0]
-
-        if mu_law:
-            output = decode_mu_law(output, self.n_classes, False)
-        if hp.apply_preemphasis:
-            output = de_emphasis(output)
-
-        # Fade-out at the end to avoid signal cutting out suddenly
-        fade_out = np.linspace(1, 0, 20 * self.hop_length)
-        output = output[:wave_len]
-        output[-20 * self.hop_length:] *= fade_out
-        
-        self.train()
-
-        return output
-
-
-    def gen_display(self, i, seq_len, b_size, gen_rate):
-        pbar = progbar(i, seq_len)
-        msg = f'| {pbar} {i*b_size}/{seq_len*b_size} | Batch Size: {b_size} | Gen Rate: {gen_rate:.1f}kHz | '
-        stream(msg)
-
-    def get_gru_cell(self, gru):
-        gru_cell = nn.GRUCell(gru.input_size, gru.hidden_size)
-        gru_cell.weight_hh.data = gru.weight_hh_l0.data
-        gru_cell.weight_ih.data = gru.weight_ih_l0.data
-        gru_cell.bias_hh.data = gru.bias_hh_l0.data
-        gru_cell.bias_ih.data = gru.bias_ih_l0.data
-        return gru_cell
-
-    def pad_tensor(self, x, pad, side='both'):
-        # NB - this is just a quick method i need right now
-        # i.e., it won't generalise to other shapes/dims
-        b, t, c = x.size()
-        total = t + 2 * pad if side == 'both' else t + pad
-        if torch.cuda.is_available():
-            padded = torch.zeros(b, total, c).cuda()
-        else:
-            padded = torch.zeros(b, total, c).cpu()
-        if side == 'before' or side == 'both':
-            padded[:, pad:pad + t, :] = x
-        elif side == 'after':
-            padded[:, :t, :] = x
-        return padded
-
-    def fold_with_overlap(self, x, target, overlap):
-
-        ''' Fold the tensor with overlap for quick batched inference.
-            Overlap will be used for crossfading in xfade_and_unfold()
-
-        Args:
-            x (tensor)    : Upsampled conditioning features.
-                            shape=(1, timesteps, features)
-            target (int)  : Target timesteps for each index of batch
-            overlap (int) : Timesteps for both xfade and rnn warmup
-
-        Return:
-            (tensor) : shape=(num_folds, target + 2 * overlap, features)
-
-        Details:
-            x = [[h1, h2, ... hn]]
-
-            Where each h is a vector of conditioning features
-
-            Eg: target=2, overlap=1 with x.size(1)=10
-
-            folded = [[h1, h2, h3, h4],
-                      [h4, h5, h6, h7],
-                      [h7, h8, h9, h10]]
-        '''
-
-        _, total_len, features = x.size()
-
-        # Calculate variables needed
-        num_folds = (total_len - overlap) // (target + overlap)
-        extended_len = num_folds * (overlap + target) + overlap
-        remaining = total_len - extended_len
-
-        # Pad if some time steps poking out
-        if remaining != 0:
-            num_folds += 1
-            padding = target + 2 * overlap - remaining
-            x = self.pad_tensor(x, padding, side='after')
-
-        if torch.cuda.is_available():
-            folded = torch.zeros(num_folds, target + 2 * overlap, features).cuda()
-        else:
-            folded = torch.zeros(num_folds, target + 2 * overlap, features).cpu()
-
-        # Get the values for the folded tensor
-        for i in range(num_folds):
-            start = i * (target + overlap)
-            end = start + target + 2 * overlap
-            folded[i] = x[:, start:end, :]
-
-        return folded
-
-    def xfade_and_unfold(self, y, target, overlap):
-
-        ''' Applies a crossfade and unfolds into a 1d array.
-
-        Args:
-            y (ndarry)    : Batched sequences of audio samples
-                            shape=(num_folds, target + 2 * overlap)
-                            dtype=np.float64
-            overlap (int) : Timesteps for both xfade and rnn warmup
-
-        Return:
-            (ndarry) : audio samples in a 1d array
-                       shape=(total_len)
-                       dtype=np.float64
-
-        Details:
-            y = [[seq1],
-                 [seq2],
-                 [seq3]]
-
-            Apply a gain envelope at both ends of the sequences
-
-            y = [[seq1_in, seq1_target, seq1_out],
-                 [seq2_in, seq2_target, seq2_out],
-                 [seq3_in, seq3_target, seq3_out]]
-
-            Stagger and add up the groups of samples:
-
-            [seq1_in, seq1_target, (seq1_out + seq2_in), seq2_target, ...]
-
-        '''
-
-        num_folds, length = y.shape
-        target = length - 2 * overlap
-        total_len = num_folds * (target + overlap) + overlap
-
-        # Need some silence for the rnn warmup
-        silence_len = overlap // 2
-        fade_len = overlap - silence_len
-        silence = np.zeros((silence_len), dtype=np.float64)
-
-        # Equal power crossfade
-        t = np.linspace(-1, 1, fade_len, dtype=np.float64)
-        fade_in = np.sqrt(0.5 * (1 + t))
-        fade_out = np.sqrt(0.5 * (1 - t))
-
-        # Concat the silence to the fades
-        fade_in = np.concatenate([silence, fade_in])
-        fade_out = np.concatenate([fade_out, silence])
-
-        # Apply the gain to the overlap samples
-        y[:, :overlap] *= fade_in
-        y[:, -overlap:] *= fade_out
-
-        unfolded = np.zeros((total_len), dtype=np.float64)
-
-        # Loop to add up all the samples
-        for i in range(num_folds):
-            start = i * (target + overlap)
-            end = start + target + 2 * overlap
-            unfolded[start:end] += y[i]
-
-        return unfolded
-
-    def get_step(self) :
-        return self.step.data.item()
-
-    def checkpoint(self, model_dir, optimizer) :
-        k_steps = self.get_step() // 1000
-        self.save(model_dir.joinpath("checkpoint_%dk_steps.pt" % k_steps), optimizer)
-
-    def log(self, path, msg) :
-        with open(path, 'a') as f:
-            print(msg, file=f)
-
-    def load(self, path, optimizer) :
-        checkpoint = torch.load(path)
-        if "optimizer_state" in checkpoint:
-            self.load_state_dict(checkpoint["model_state"])
-            optimizer.load_state_dict(checkpoint["optimizer_state"])
-        else:
-            # Backwards compatibility
-            self.load_state_dict(checkpoint)
-
-    def save(self, path, optimizer) :
-        torch.save({
-            "model_state": self.state_dict(),
-            "optimizer_state": optimizer.state_dict(),
-        }, path)
-
-    def num_params(self, print_out=True):
-        parameters = filter(lambda p: p.requires_grad, self.parameters())
-        parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
-        if print_out :
-            print('Trainable Parameters: %.3fM' % parameters)
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from vocoder.distribution import sample_from_discretized_mix_logistic
+from vocoder.display import *
+from vocoder.audio import *
+
+
+class ResBlock(nn.Module):
+    def __init__(self, dims):
+        super().__init__()
+        self.conv1 = nn.Conv1d(dims, dims, kernel_size=1, bias=False)
+        self.conv2 = nn.Conv1d(dims, dims, kernel_size=1, bias=False)
+        self.batch_norm1 = nn.BatchNorm1d(dims)
+        self.batch_norm2 = nn.BatchNorm1d(dims)
+
+    def forward(self, x):
+        residual = x
+        x = self.conv1(x)
+        x = self.batch_norm1(x)
+        x = F.relu(x)
+        x = self.conv2(x)
+        x = self.batch_norm2(x)
+        return x + residual
+
+
+class MelResNet(nn.Module):
+    def __init__(self, res_blocks, in_dims, compute_dims, res_out_dims, pad):
+        super().__init__()
+        k_size = pad * 2 + 1
+        self.conv_in = nn.Conv1d(in_dims, compute_dims, kernel_size=k_size, bias=False)
+        self.batch_norm = nn.BatchNorm1d(compute_dims)
+        self.layers = nn.ModuleList()
+        for i in range(res_blocks):
+            self.layers.append(ResBlock(compute_dims))
+        self.conv_out = nn.Conv1d(compute_dims, res_out_dims, kernel_size=1)
+
+    def forward(self, x):
+        x = self.conv_in(x)
+        x = self.batch_norm(x)
+        x = F.relu(x)
+        for f in self.layers: x = f(x)
+        x = self.conv_out(x)
+        return x
+
+
+class Stretch2d(nn.Module):
+    def __init__(self, x_scale, y_scale):
+        super().__init__()
+        self.x_scale = x_scale
+        self.y_scale = y_scale
+
+    def forward(self, x):
+        b, c, h, w = x.size()
+        x = x.unsqueeze(-1).unsqueeze(3)
+        x = x.repeat(1, 1, 1, self.y_scale, 1, self.x_scale)
+        return x.view(b, c, h * self.y_scale, w * self.x_scale)
+
+
+class UpsampleNetwork(nn.Module):
+    def __init__(self, feat_dims, upsample_scales, compute_dims,
+                 res_blocks, res_out_dims, pad):
+        super().__init__()
+        total_scale = np.cumproduct(upsample_scales)[-1]
+        self.indent = pad * total_scale
+        self.resnet = MelResNet(res_blocks, feat_dims, compute_dims, res_out_dims, pad)
+        self.resnet_stretch = Stretch2d(total_scale, 1)
+        self.up_layers = nn.ModuleList()
+        for scale in upsample_scales:
+            k_size = (1, scale * 2 + 1)
+            padding = (0, scale)
+            stretch = Stretch2d(scale, 1)
+            conv = nn.Conv2d(1, 1, kernel_size=k_size, padding=padding, bias=False)
+            conv.weight.data.fill_(1. / k_size[1])
+            self.up_layers.append(stretch)
+            self.up_layers.append(conv)
+
+    def forward(self, m):
+        aux = self.resnet(m).unsqueeze(1)
+        aux = self.resnet_stretch(aux)
+        aux = aux.squeeze(1)
+        m = m.unsqueeze(1)
+        for f in self.up_layers: m = f(m)
+        m = m.squeeze(1)[:, :, self.indent:-self.indent]
+        return m.transpose(1, 2), aux.transpose(1, 2)
+
+
+class WaveRNN(nn.Module):
+    def __init__(self, rnn_dims, fc_dims, bits, pad, upsample_factors,
+                 feat_dims, compute_dims, res_out_dims, res_blocks,
+                 hop_length, sample_rate, mode='RAW'):
+        super().__init__()
+        self.mode = mode
+        self.pad = pad
+        if self.mode == 'RAW' :
+            self.n_classes = 2 ** bits
+        elif self.mode == 'MOL' :
+            self.n_classes = 30
+        else :
+            RuntimeError("Unknown model mode value - ", self.mode)
+
+        self.rnn_dims = rnn_dims
+        self.aux_dims = res_out_dims // 4
+        self.hop_length = hop_length
+        self.sample_rate = sample_rate
+
+        self.upsample = UpsampleNetwork(feat_dims, upsample_factors, compute_dims, res_blocks, res_out_dims, pad)
+        self.I = nn.Linear(feat_dims + self.aux_dims + 1, rnn_dims)
+        self.rnn1 = nn.GRU(rnn_dims, rnn_dims, batch_first=True)
+        self.rnn2 = nn.GRU(rnn_dims + self.aux_dims, rnn_dims, batch_first=True)
+        self.fc1 = nn.Linear(rnn_dims + self.aux_dims, fc_dims)
+        self.fc2 = nn.Linear(fc_dims + self.aux_dims, fc_dims)
+        self.fc3 = nn.Linear(fc_dims, self.n_classes)
+
+        self.step = nn.Parameter(torch.zeros(1).long(), requires_grad=False)
+        self.num_params()
+
+    def forward(self, x, mels):
+        self.step += 1
+        bsize = x.size(0)
+        if torch.cuda.is_available():
+            h1 = torch.zeros(1, bsize, self.rnn_dims).cuda()
+            h2 = torch.zeros(1, bsize, self.rnn_dims).cuda()
+        else:
+            h1 = torch.zeros(1, bsize, self.rnn_dims).cpu()
+            h2 = torch.zeros(1, bsize, self.rnn_dims).cpu()
+        mels, aux = self.upsample(mels)
+
+        aux_idx = [self.aux_dims * i for i in range(5)]
+        a1 = aux[:, :, aux_idx[0]:aux_idx[1]]
+        a2 = aux[:, :, aux_idx[1]:aux_idx[2]]
+        a3 = aux[:, :, aux_idx[2]:aux_idx[3]]
+        a4 = aux[:, :, aux_idx[3]:aux_idx[4]]
+
+        x = torch.cat([x.unsqueeze(-1), mels, a1], dim=2)
+        x = self.I(x)
+        res = x
+        x, _ = self.rnn1(x, h1)
+
+        x = x + res
+        res = x
+        x = torch.cat([x, a2], dim=2)
+        x, _ = self.rnn2(x, h2)
+
+        x = x + res
+        x = torch.cat([x, a3], dim=2)
+        x = F.relu(self.fc1(x))
+
+        x = torch.cat([x, a4], dim=2)
+        x = F.relu(self.fc2(x))
+        return self.fc3(x)
+
+    def generate(self, mels, batched, target, overlap, mu_law, progress_callback=None):
+        mu_law = mu_law if self.mode == 'RAW' else False
+        progress_callback = progress_callback or self.gen_display
+
+        self.eval()
+        output = []
+        start = time.time()
+        rnn1 = self.get_gru_cell(self.rnn1)
+        rnn2 = self.get_gru_cell(self.rnn2)
+
+        with torch.no_grad():
+            if torch.cuda.is_available():
+                mels = mels.cuda()
+            else:
+                mels = mels.cpu()
+            wave_len = (mels.size(-1) - 1) * self.hop_length
+            mels = self.pad_tensor(mels.transpose(1, 2), pad=self.pad, side='both')
+            mels, aux = self.upsample(mels.transpose(1, 2))
+
+            if batched:
+                mels = self.fold_with_overlap(mels, target, overlap)
+                aux = self.fold_with_overlap(aux, target, overlap)
+
+            b_size, seq_len, _ = mels.size()
+
+            if torch.cuda.is_available():
+                h1 = torch.zeros(b_size, self.rnn_dims).cuda()
+                h2 = torch.zeros(b_size, self.rnn_dims).cuda()
+                x = torch.zeros(b_size, 1).cuda()
+            else:
+                h1 = torch.zeros(b_size, self.rnn_dims).cpu()
+                h2 = torch.zeros(b_size, self.rnn_dims).cpu()
+                x = torch.zeros(b_size, 1).cpu()
+
+            d = self.aux_dims
+            aux_split = [aux[:, :, d * i:d * (i + 1)] for i in range(4)]
+
+            for i in range(seq_len):
+
+                m_t = mels[:, i, :]
+
+                a1_t, a2_t, a3_t, a4_t = (a[:, i, :] for a in aux_split)
+
+                x = torch.cat([x, m_t, a1_t], dim=1)
+                x = self.I(x)
+                h1 = rnn1(x, h1)
+
+                x = x + h1
+                inp = torch.cat([x, a2_t], dim=1)
+                h2 = rnn2(inp, h2)
+
+                x = x + h2
+                x = torch.cat([x, a3_t], dim=1)
+                x = F.relu(self.fc1(x))
+
+                x = torch.cat([x, a4_t], dim=1)
+                x = F.relu(self.fc2(x))
+
+                logits = self.fc3(x)
+
+                if self.mode == 'MOL':
+                    sample = sample_from_discretized_mix_logistic(logits.unsqueeze(0).transpose(1, 2))
+                    output.append(sample.view(-1))
+                    if torch.cuda.is_available():
+                        # x = torch.FloatTensor([[sample]]).cuda()
+                        x = sample.transpose(0, 1).cuda()
+                    else:
+                        x = sample.transpose(0, 1)
+
+                elif self.mode == 'RAW' :
+                    posterior = F.softmax(logits, dim=1)
+                    distrib = torch.distributions.Categorical(posterior)
+
+                    sample = 2 * distrib.sample().float() / (self.n_classes - 1.) - 1.
+                    output.append(sample)
+                    x = sample.unsqueeze(-1)
+                else:
+                    raise RuntimeError("Unknown model mode value - ", self.mode)
+
+                if i % 100 == 0:
+                    gen_rate = (i + 1) / (time.time() - start) * b_size / 1000
+                    progress_callback(i, seq_len, b_size, gen_rate)
+
+        output = torch.stack(output).transpose(0, 1)
+        output = output.cpu().numpy()
+        output = output.astype(np.float64)
+        
+        if batched:
+            output = self.xfade_and_unfold(output, target, overlap)
+        else:
+            output = output[0]
+
+        if mu_law:
+            output = decode_mu_law(output, self.n_classes, False)
+        if hp.apply_preemphasis:
+            output = de_emphasis(output)
+
+        # Fade-out at the end to avoid signal cutting out suddenly
+        fade_out = np.linspace(1, 0, 20 * self.hop_length)
+        output = output[:wave_len]
+        output[-20 * self.hop_length:] *= fade_out
+        
+        self.train()
+
+        return output
+
+
+    def gen_display(self, i, seq_len, b_size, gen_rate):
+        pbar = progbar(i, seq_len)
+        msg = f'| {pbar} {i*b_size}/{seq_len*b_size} | Batch Size: {b_size} | Gen Rate: {gen_rate:.1f}kHz | '
+        stream(msg)
+
+    def get_gru_cell(self, gru):
+        gru_cell = nn.GRUCell(gru.input_size, gru.hidden_size)
+        gru_cell.weight_hh.data = gru.weight_hh_l0.data
+        gru_cell.weight_ih.data = gru.weight_ih_l0.data
+        gru_cell.bias_hh.data = gru.bias_hh_l0.data
+        gru_cell.bias_ih.data = gru.bias_ih_l0.data
+        return gru_cell
+
+    def pad_tensor(self, x, pad, side='both'):
+        # NB - this is just a quick method i need right now
+        # i.e., it won't generalise to other shapes/dims
+        b, t, c = x.size()
+        total = t + 2 * pad if side == 'both' else t + pad
+        if torch.cuda.is_available():
+            padded = torch.zeros(b, total, c).cuda()
+        else:
+            padded = torch.zeros(b, total, c).cpu()
+        if side == 'before' or side == 'both':
+            padded[:, pad:pad + t, :] = x
+        elif side == 'after':
+            padded[:, :t, :] = x
+        return padded
+
+    def fold_with_overlap(self, x, target, overlap):
+
+        ''' Fold the tensor with overlap for quick batched inference.
+            Overlap will be used for crossfading in xfade_and_unfold()
+
+        Args:
+            x (tensor)    : Upsampled conditioning features.
+                            shape=(1, timesteps, features)
+            target (int)  : Target timesteps for each index of batch
+            overlap (int) : Timesteps for both xfade and rnn warmup
+
+        Return:
+            (tensor) : shape=(num_folds, target + 2 * overlap, features)
+
+        Details:
+            x = [[h1, h2, ... hn]]
+
+            Where each h is a vector of conditioning features
+
+            Eg: target=2, overlap=1 with x.size(1)=10
+
+            folded = [[h1, h2, h3, h4],
+                      [h4, h5, h6, h7],
+                      [h7, h8, h9, h10]]
+        '''
+
+        _, total_len, features = x.size()
+
+        # Calculate variables needed
+        num_folds = (total_len - overlap) // (target + overlap)
+        extended_len = num_folds * (overlap + target) + overlap
+        remaining = total_len - extended_len
+
+        # Pad if some time steps poking out
+        if remaining != 0:
+            num_folds += 1
+            padding = target + 2 * overlap - remaining
+            x = self.pad_tensor(x, padding, side='after')
+
+        if torch.cuda.is_available():
+            folded = torch.zeros(num_folds, target + 2 * overlap, features).cuda()
+        else:
+            folded = torch.zeros(num_folds, target + 2 * overlap, features).cpu()
+
+        # Get the values for the folded tensor
+        for i in range(num_folds):
+            start = i * (target + overlap)
+            end = start + target + 2 * overlap
+            folded[i] = x[:, start:end, :]
+
+        return folded
+
+    def xfade_and_unfold(self, y, target, overlap):
+
+        ''' Applies a crossfade and unfolds into a 1d array.
+
+        Args:
+            y (ndarry)    : Batched sequences of audio samples
+                            shape=(num_folds, target + 2 * overlap)
+                            dtype=np.float64
+            overlap (int) : Timesteps for both xfade and rnn warmup
+
+        Return:
+            (ndarry) : audio samples in a 1d array
+                       shape=(total_len)
+                       dtype=np.float64
+
+        Details:
+            y = [[seq1],
+                 [seq2],
+                 [seq3]]
+
+            Apply a gain envelope at both ends of the sequences
+
+            y = [[seq1_in, seq1_target, seq1_out],
+                 [seq2_in, seq2_target, seq2_out],
+                 [seq3_in, seq3_target, seq3_out]]
+
+            Stagger and add up the groups of samples:
+
+            [seq1_in, seq1_target, (seq1_out + seq2_in), seq2_target, ...]
+
+        '''
+
+        num_folds, length = y.shape
+        target = length - 2 * overlap
+        total_len = num_folds * (target + overlap) + overlap
+
+        # Need some silence for the rnn warmup
+        silence_len = overlap // 2
+        fade_len = overlap - silence_len
+        silence = np.zeros((silence_len), dtype=np.float64)
+
+        # Equal power crossfade
+        t = np.linspace(-1, 1, fade_len, dtype=np.float64)
+        fade_in = np.sqrt(0.5 * (1 + t))
+        fade_out = np.sqrt(0.5 * (1 - t))
+
+        # Concat the silence to the fades
+        fade_in = np.concatenate([silence, fade_in])
+        fade_out = np.concatenate([fade_out, silence])
+
+        # Apply the gain to the overlap samples
+        y[:, :overlap] *= fade_in
+        y[:, -overlap:] *= fade_out
+
+        unfolded = np.zeros((total_len), dtype=np.float64)
+
+        # Loop to add up all the samples
+        for i in range(num_folds):
+            start = i * (target + overlap)
+            end = start + target + 2 * overlap
+            unfolded[start:end] += y[i]
+
+        return unfolded
+
+    def get_step(self) :
+        return self.step.data.item()
+
+    def checkpoint(self, model_dir, optimizer) :
+        k_steps = self.get_step() // 1000
+        self.save(model_dir.joinpath("checkpoint_%dk_steps.pt" % k_steps), optimizer)
+
+    def log(self, path, msg) :
+        with open(path, 'a') as f:
+            print(msg, file=f)
+
+    def load(self, path, optimizer) :
+        checkpoint = torch.load(path, weights_only=False)
+        if "optimizer_state" in checkpoint:
+            self.load_state_dict(checkpoint["model_state"])
+            optimizer.load_state_dict(checkpoint["optimizer_state"])
+        else:
+            # Backwards compatibility
+            self.load_state_dict(checkpoint)
+
+    def save(self, path, optimizer) :
+        torch.save({
+            "model_state": self.state_dict(),
+            "optimizer_state": optimizer.state_dict(),
+        }, path)
+
+    def num_params(self, print_out=True):
+        parameters = filter(lambda p: p.requires_grad, self.parameters())
+        parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
+        if print_out :
+            print('Trainable Parameters: %.3fM' % parameters)