From 71c382be14e48361b25d8dcb492560f188aabb63 Mon Sep 17 00:00:00 2001
From: erogol <erogol@hotmail.com>
Date: Wed, 30 Dec 2020 14:23:08 +0100
Subject: [PATCH] copy model scale stats file with config.json to the trianing
 folder, fixed for model inits

---
 TTS/tts/utils/generic_utils.py |  3 +++
 TTS/utils/io.py                | 13 ++++++++++---
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/TTS/tts/utils/generic_utils.py b/TTS/tts/utils/generic_utils.py
index d4e321e4..568fc8c6 100644
--- a/TTS/tts/utils/generic_utils.py
+++ b/TTS/tts/utils/generic_utils.py
@@ -110,6 +110,7 @@ def setup_model(num_chars, num_speakers, c, speaker_embedding_dim=None):
                         num_heads=2,
                         num_layers_enc=6,
                         encoder_type=c.encoder_type,
+                        rel_attn_window_size=4,
                         dropout_p=0.1,
                         num_flow_blocks_dec=12,
                         kernel_size_dec=5,
@@ -132,6 +133,8 @@ def setup_model(num_chars, num_speakers, c, speaker_embedding_dim=None):
                         hidden_channels=128,
                         positional_encoding=c['positional_encoding'],
                         encoder_type=c['encoder_type'],
+                        encoder_params=c['encoder_params'],
+                        decoder_type=c['decoder_type'],
                         decoder_residual_conv_bn_params=c['decoder_residual_conv_bn_params'],
                         c_in_channels=0)
     return model
diff --git a/TTS/utils/io.py b/TTS/utils/io.py
index 3cc36e95..3fb7a92d 100644
--- a/TTS/utils/io.py
+++ b/TTS/utils/io.py
@@ -3,6 +3,7 @@ import re
 import json
 import yaml
 import pickle as pickle_tts
+from shutil import copyfile
 
 
 class RenamingUnpickler(pickle_tts.Unpickler):
@@ -44,16 +45,19 @@ def load_config(config_path: str) -> AttrDict:
     return config
 
 
-def copy_config_file(config_file, out_path, new_fields):
-    """Copy config.json to training folder and add
+def copy_model_files(c, config_file, out_path, new_fields):
+    """Copy config.json and other model files to training folder and add
     new fields.
 
     Args:
+        c (dict): model config from config.json.
         config_file (str): path to config file.
         out_path (str): output path to copy the file.
         new_fields (dict): new fileds to be added or edited
             in the config file.
     """
+    # copy config.json
+    copy_config_path = OUT_PATH
     config_lines = open(config_file, "r").readlines()
     # add extra information fields
     for key, value in new_fields.items():
@@ -62,6 +66,9 @@ def copy_config_file(config_file, out_path, new_fields):
         else:
             new_line = '"{}":{},\n'.format(key, value)
         config_lines.insert(1, new_line)
-    config_out_file = open(out_path, "w")
+    config_out_file = open(copy_config_path, "w")
     config_out_file.writelines(config_lines)
     config_out_file.close()
+    # copy model stats file if available
+    copy_stats_path = os.path.join(out_path, 'scale_stats.npy')
+    copyfile(c.audio['stats_path'], copy_stats_path)