Merge pull request #107 from TheButlah/master

Added CPU support, other fixes too

Author: fatchord

.gitignore

@@ -1 +1,47 @@
+# PyCharm files
 .idea
+
+# Mac files
+.DS_Store
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Jupyter Notebook
+.ipynb_checkpoints

gen_tacotron.py

@@ -8,7 +8,7 @@
 from utils.text import text_to_sequence
 from utils.display import save_attention, simple_table
 
-if __name__ == "__main__" :
+if __name__ == "__main__":
 
     # Parse Arguments
     parser = argparse.ArgumentParser(description='TTS Generator')
@@ -19,6 +19,7 @@
     parser.add_argument('--overlap', '-o', type=int, help='[int] number of crossover samples')
     parser.add_argument('--weights_path', '-w', type=str, help='[string/path] Load in different Tacotron Weights')
     parser.add_argument('--save_attention', '-a', dest='save_attn', action='store_true', help='Save Attention Plots')
+    parser.add_argument('--force_cpu', '-c', action='store_true', help='Forces CPU-only training, even when in CUDA capable environment')
     parser.set_defaults(batched=hp.voc_gen_batched)
     parser.set_defaults(target=hp.voc_target)
     parser.set_defaults(overlap=hp.voc_overlap)
@@ -36,6 +37,12 @@
 
     paths = Paths(hp.data_path, hp.voc_model_id, hp.tts_model_id)
 
+    if not args.force_cpu and torch.cuda.is_available():
+        device = torch.device('cuda')
+    else:
+        device = torch.device('cpu')
+    print('Using device:', device)
+
     print('\nInitialising WaveRNN Model...\n')
 
     # Instantiate WaveRNN Model
@@ -50,7 +57,7 @@
                         res_blocks=hp.voc_res_blocks,
                         hop_length=hp.hop_length,
                         sample_rate=hp.sample_rate,
-                        mode=hp.voc_mode).cuda()
+                        mode=hp.voc_mode).to(device)
 
     voc_model.restore(paths.voc_latest_weights)
 
@@ -68,15 +75,15 @@
                          lstm_dims=hp.tts_lstm_dims,
                          postnet_K=hp.tts_postnet_K,
                          num_highways=hp.tts_num_highways,
-                         dropout=hp.tts_dropout).cuda()
+                         dropout=hp.tts_dropout).to(device)
 
     tts_restore_path = weights_path if weights_path else paths.tts_latest_weights
     tts_model.restore(tts_restore_path)
 
-    if input_text :
+    if input_text:
         inputs = [text_to_sequence(input_text.strip(), hp.tts_cleaner_names)]
-    else :
-        with open('sentences.txt') as f :
+    else:
+        with open('sentences.txt') as f:
             inputs = [text_to_sequence(l.strip(), hp.tts_cleaner_names) for l in f]
 
     voc_k = voc_model.get_step() // 1000
@@ -89,21 +96,21 @@
                   ('Target Samples', target if batched else 'N/A'),
                   ('Overlap Samples', overlap if batched else 'N/A')])
 
-    for i, x in enumerate(inputs, 1) :
+    for i, x in enumerate(inputs, 1):
 
         print(f'\n| Generating {i}/{len(inputs)}')
         _, m, attention = tts_model.generate(x)
 
-        if input_text :
+        if input_text:
             save_path = f'{paths.tts_output}__input_{input_text[:10]}_{tts_k}k.wav'
-        else :
+        else:
             save_path = f'{paths.tts_output}{i}_batched{str(batched)}_{tts_k}k.wav'
 
-        if save_attn : save_attention(attention, save_path)
+        if save_attn: save_attention(attention, save_path)
 
         m = torch.tensor(m).unsqueeze(0)
         m = (m + 4) / 8
 
         voc_model.generate(m, save_path, batched, hp.voc_target, hp.voc_overlap, hp.mu_law)
 
-    print('\n\nDone.\n')
+    print('\n\nDone.\n')

gen_wavernn.py

@@ -7,23 +7,23 @@
 import argparse
 
 
-def gen_testset(model, test_set, samples, batched, target, overlap, save_path) :
+def gen_testset(model, test_set, samples, batched, target, overlap, save_path):
 
     k = model.get_step() // 1000
 
     for i, (m, x) in enumerate(test_set, 1):
 
-        if i > samples : break
+        if i > samples: break
 
         print('\n| Generating: %i/%i' % (i, samples))
 
         x = x[0].numpy()
 
         bits = 16 if hp.voc_mode == 'MOL' else hp.bits
 
-        if hp.mu_law and hp.voc_mode != 'MOL' :
+        if hp.mu_law and hp.voc_mode != 'MOL':
             x = decode_mu_law(x, 2**bits, from_labels=True)
-        else :
+        else:
             x = label_2_float(x, bits)
 
         save_wav(x, f'{save_path}{k}k_steps_{i}_target.wav')
@@ -34,7 +34,7 @@ def gen_testset(model, test_set, samples, batched, target, overlap, save_path) :
         _ = model.generate(m, save_str, batched, target, overlap, hp.mu_law)
 
 
-def gen_from_file(model, load_path, save_path, batched, target, overlap) :
+def gen_from_file(model, load_path, save_path, batched, target, overlap):
 
     k = model.get_step() // 1000
     file_name = load_path.split('/')[-1]
@@ -62,6 +62,7 @@ def gen_from_file(model, load_path, save_path, batched, target, overlap) :
     parser.add_argument('--file', '-f', type=str, help='[string/path] for testing a wav outside dataset')
     parser.add_argument('--weights', '-w', type=str, help='[string/path] checkpoint file to load weights from')
     parser.add_argument('--gta', '-g', dest='use_gta', action='store_true', help='Generate from GTA testset')
+    parser.add_argument('--force_cpu', '-c', action='store_true', help='Forces CPU-only training, even when in CUDA capable environment')
 
     parser.set_defaults(batched=hp.voc_gen_batched)
     parser.set_defaults(samples=hp.voc_gen_at_checkpoint)
@@ -80,6 +81,12 @@ def gen_from_file(model, load_path, save_path, batched, target, overlap) :
     file = args.file
     gta = args.gta
 
+    if not args.force_cpu and torch.cuda.is_available():
+        device = torch.device('cuda')
+    else:
+        device = torch.device('cpu')
+    print('Using device:', device)
+
     print('\nInitialising Model...\n')
 
     model = WaveRNN(rnn_dims=hp.voc_rnn_dims,
@@ -93,7 +100,7 @@ def gen_from_file(model, load_path, save_path, batched, target, overlap) :
                     res_blocks=hp.voc_res_blocks,
                     hop_length=hp.hop_length,
                     sample_rate=hp.sample_rate,
-                    mode=hp.voc_mode).cuda()
+                    mode=hp.voc_mode).to(device)
 
     paths = Paths(hp.data_path, hp.voc_model_id, hp.tts_model_id)
 
@@ -107,9 +114,9 @@ def gen_from_file(model, load_path, save_path, batched, target, overlap) :
 
     _, test_set = get_vocoder_datasets(paths.data, 1, gta)
 
-    if file :
+    if file:
         gen_from_file(model, file, paths.voc_output, batched, target, overlap)
-    else :
+    else:
         gen_testset(model, test_set, samples, batched, target, overlap, paths.voc_output)
 
     print('\n\nExiting...\n')

models/__init__.py

@@ -4,8 +4,8 @@
 from utils.display import *
 from utils.dsp import *
 
-class WaveRNN(nn.Module) :
-    def __init__(self, hidden_size=896, quantisation=256) :
+class WaveRNN(nn.Module):
+    def __init__(self, hidden_size=896, quantisation=256):
         super(WaveRNN, self).__init__()
 
         self.hidden_size = hidden_size
@@ -33,7 +33,7 @@ def __init__(self, hidden_size=896, quantisation=256) :
         self.num_params()
 
 
-    def forward(self, prev_y, prev_hidden, current_coarse) :
+    def forward(self, prev_y, prev_hidden, current_coarse):
 
         # Main matmul - the projection is split 3 ways
         R_hidden = self.R(prev_hidden)
@@ -71,9 +71,10 @@ def forward(self, prev_y, prev_hidden, current_coarse) :
         return out_coarse, out_fine, hidden
 
 
-    def generate(self, seq_len) :
-        
-        with torch.no_grad() :
+    def generate(self, seq_len):
+        device = next(self.parameters()).device  # use same device as parameters
+
+        with torch.no_grad():
 
             # First split up the biases for the gates 
             b_coarse_u, b_fine_u = torch.split(self.bias_u, self.split_size)
@@ -84,17 +85,17 @@ def generate(self, seq_len) :
             c_outputs, f_outputs = [], []
 
             # Some initial inputs
-            out_coarse = torch.LongTensor([0]).cuda()
-            out_fine = torch.LongTensor([0]).cuda()
+            out_coarse = torch.tensor([0], dtype=torch.long, device=device)
+            out_fine = torch.tensor([0], dtype=torch.long, device=device)
 
             # We'll meed a hidden state
-            hidden = self.init_hidden()
+            hidden = self.get_initial_hidden()
 
             # Need a clock for display
             start = time.time()
 
             # Loop for generation
-            for i in range(seq_len) :
+            for i in range(seq_len):
 
                 # Split into two hidden states
                 hidden_coarse, hidden_fine = \
@@ -162,10 +163,11 @@ def generate(self, seq_len) :
 
         return output, coarse, fine
 
-    def init_hidden(self, batch_size=1) :
-        return torch.zeros(batch_size, self.hidden_size).cuda()
+    def get_initial_hidden(self, batch_size=1):
+        device = next(self.parameters()).device  # use same device as parameters
+        return torch.zeros(batch_size, self.hidden_size, device=device)
 
-    def num_params(self) :
+    def num_params(self):
         parameters = filter(lambda p: p.requires_grad, self.parameters())
         parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
-        print('Trainable Parameters: %.3f million' % parameters)
+        print('Trainable Parameters: %.3f million' % parameters)