Spaces:
Runtime error
Runtime error
| # from https://github.com/jaywalnut310/vits | |
| import os | |
| import sys | |
| import logging | |
| import subprocess | |
| import torch | |
| import numpy as np | |
| from omegaconf import OmegaConf | |
| from scipy.io.wavfile import read | |
| MATPLOTLIB_FLAG = False | |
| logging.basicConfig( | |
| stream=sys.stdout, | |
| level=logging.INFO, | |
| format='[%(levelname)s|%(filename)s:%(lineno)s][%(asctime)s] >>> %(message)s' | |
| ) | |
| logger = logging | |
| def load_checkpoint(checkpoint_path, rank=0, model_g=None, model_d=None, optim_g=None, optim_d=None): | |
| assert os.path.isfile(checkpoint_path) | |
| checkpoint_dict = torch.load(checkpoint_path, map_location='cpu') | |
| iteration = checkpoint_dict['iteration'] | |
| learning_rate = checkpoint_dict['learning_rate'] | |
| config = checkpoint_dict['config'] | |
| if model_g is not None: | |
| model_g, optim_g = load_model( | |
| model_g, | |
| checkpoint_dict['model_g'], | |
| optim_g, | |
| checkpoint_dict['optimizer_g']) | |
| if model_d is not None: | |
| model_d, optim_d = load_model( | |
| model_d, | |
| checkpoint_dict['model_d'], | |
| optim_d, | |
| checkpoint_dict['optimizer_d']) | |
| if rank == 0: | |
| logger.info( | |
| "Loaded checkpoint '{}' (iteration {})".format( | |
| checkpoint_path, | |
| iteration | |
| ) | |
| ) | |
| return model_g, model_d, optim_g, optim_d, learning_rate, iteration, config | |
| def load_checkpoint_diffsize(checkpoint_path, rank=0, model_g=None, model_d=None): | |
| assert os.path.isfile(checkpoint_path) | |
| checkpoint_dict = torch.load(checkpoint_path, map_location='cpu') | |
| iteration = checkpoint_dict['iteration'] | |
| learning_rate = checkpoint_dict['learning_rate'] | |
| config = checkpoint_dict['config'] | |
| if model_g is not None: | |
| model_g = load_model_diffsize( | |
| model_g, | |
| checkpoint_dict['model_g']) | |
| if model_d is not None: | |
| model_d = load_model_diffsize( | |
| model_d, | |
| checkpoint_dict['model_d']) | |
| if rank == 0: | |
| logger.info( | |
| "Loaded checkpoint '{}' (iteration {})".format( | |
| checkpoint_path, | |
| iteration | |
| ) | |
| ) | |
| del checkpoint_dict | |
| return model_g, model_d, learning_rate, iteration, config | |
| def load_model_diffsize(model, model_state_dict): | |
| if hasattr(model, 'module'): | |
| state_dict = model.module.state_dict() | |
| else: | |
| state_dict = model.state_dict() | |
| for k, v in model_state_dict.items(): | |
| if k in state_dict and state_dict[k].size() == v.size(): | |
| state_dict[k] = v | |
| if hasattr(model, 'module'): | |
| model.module.load_state_dict(state_dict, strict=False) | |
| else: | |
| model.load_state_dict(state_dict, strict=False) | |
| return model | |
| def load_model(model, model_state_dict, optim, optim_state_dict): | |
| if optim is not None: | |
| optim.load_state_dict(optim_state_dict) | |
| if hasattr(model, 'module'): | |
| state_dict = model.module.state_dict() | |
| else: | |
| state_dict = model.state_dict() | |
| for k, v in model_state_dict.items(): | |
| if k in state_dict and state_dict[k].size() == v.size(): | |
| state_dict[k] = v | |
| if hasattr(model, 'module'): | |
| model.module.load_state_dict(state_dict) | |
| else: | |
| model.load_state_dict(state_dict) | |
| return model, optim | |
| def save_checkpoint(net_g, optim_g, net_d, optim_d, hps, epoch, learning_rate, save_path): | |
| def get_state_dict(model): | |
| if hasattr(model, 'module'): | |
| state_dict = model.module.state_dict() | |
| else: | |
| state_dict = model.state_dict() | |
| return state_dict | |
| torch.save({'model_g': get_state_dict(net_g), | |
| 'model_d': get_state_dict(net_d), | |
| 'optimizer_g': optim_g.state_dict(), | |
| 'optimizer_d': optim_d.state_dict(), | |
| 'config': str(hps), | |
| 'iteration': epoch, | |
| 'learning_rate': learning_rate}, save_path) | |
| def summarize(writer, global_step, scalars={}, histograms={}, images={}, audios={}, audio_sampling_rate=22050): | |
| for k, v in scalars.items(): | |
| writer.add_scalar(k, v, global_step) | |
| for k, v in histograms.items(): | |
| writer.add_histogram(k, v, global_step) | |
| for k, v in images.items(): | |
| writer.add_image(k, v, global_step, dataformats='HWC') | |
| for k, v in audios.items(): | |
| writer.add_audio(k, v, global_step, audio_sampling_rate) | |
| def plot_spectrogram_to_numpy(spectrogram): | |
| global MATPLOTLIB_FLAG | |
| if not MATPLOTLIB_FLAG: | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| MATPLOTLIB_FLAG = True | |
| mpl_logger = logging.getLogger('matplotlib') | |
| mpl_logger.setLevel(logging.WARNING) | |
| import matplotlib.pylab as plt | |
| import numpy as np | |
| fig, ax = plt.subplots(figsize=(10, 2)) | |
| im = ax.imshow(spectrogram, aspect="auto", origin="lower", | |
| interpolation='none') | |
| plt.colorbar(im, ax=ax) | |
| plt.xlabel("Frames") | |
| plt.ylabel("Channels") | |
| plt.tight_layout() | |
| fig.canvas.draw() | |
| data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='') | |
| data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) | |
| plt.close() | |
| return data | |
| def plot_alignment_to_numpy(alignment, info=None): | |
| global MATPLOTLIB_FLAG | |
| if not MATPLOTLIB_FLAG: | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| MATPLOTLIB_FLAG = True | |
| mpl_logger = logging.getLogger('matplotlib') | |
| mpl_logger.setLevel(logging.WARNING) | |
| import matplotlib.pylab as plt | |
| import numpy as np | |
| fig, ax = plt.subplots(figsize=(6, 4)) | |
| im = ax.imshow(alignment.transpose(), aspect='auto', origin='lower', | |
| interpolation='none') | |
| fig.colorbar(im, ax=ax) | |
| xlabel = 'Decoder timestep' | |
| if info is not None: | |
| xlabel += '\n\n' + info | |
| plt.xlabel(xlabel) | |
| plt.ylabel('Encoder timestep') | |
| plt.tight_layout() | |
| fig.canvas.draw() | |
| data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='') | |
| data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) | |
| plt.close() | |
| return data | |
| def load_wav_to_torch(full_path): | |
| sampling_rate, wav = read(full_path) | |
| if len(wav.shape) == 2: | |
| wav = wav[:, 0] | |
| if wav.dtype == np.int16: | |
| wav = wav / 32768.0 | |
| elif wav.dtype == np.int32: | |
| wav = wav / 2147483648.0 | |
| elif wav.dtype == np.uint8: | |
| wav = (wav - 128) / 128.0 | |
| wav = wav.astype(np.float32) | |
| return torch.FloatTensor(wav), sampling_rate | |
| def load_filepaths_and_text(filename, split="|"): | |
| with open(filename, encoding='utf-8') as f: | |
| filepaths_and_text = [line.strip().split(split) for line in f] | |
| return filepaths_and_text | |
| def get_hparams(args, init=True): | |
| config = OmegaConf.load(args.config) | |
| hparams = HParams(**config) | |
| model_dir = os.path.join(hparams.train.log_path, args.model) | |
| if not os.path.exists(model_dir): | |
| os.makedirs(model_dir) | |
| hparams.model_name = args.model | |
| hparams.model_dir = model_dir | |
| config_save_path = os.path.join(model_dir, "config.yaml") | |
| if init: | |
| OmegaConf.save(config, config_save_path) | |
| return hparams | |
| def get_hparams_from_file(config_path): | |
| config = OmegaConf.load(config_path) | |
| hparams = HParams(**config) | |
| return hparams | |
| def check_git_hash(model_dir): | |
| source_dir = os.path.dirname(os.path.realpath(__file__)) | |
| if not os.path.exists(os.path.join(source_dir, ".git")): | |
| logger.warn("{} is not a git repository, therefore hash value comparison will be ignored.".format( | |
| source_dir | |
| )) | |
| return | |
| cur_hash = subprocess.getoutput("git rev-parse HEAD") | |
| path = os.path.join(model_dir, "githash") | |
| if os.path.exists(path): | |
| saved_hash = open(path).read() | |
| if saved_hash != cur_hash: | |
| logger.warn("git hash values are different. {}(saved) != {}(current)".format( | |
| saved_hash[:8], cur_hash[:8])) | |
| else: | |
| open(path, "w").write(cur_hash) | |
| def get_logger(model_dir, filename="train.log"): | |
| global logger | |
| logger = logging.getLogger(os.path.basename(model_dir)) | |
| logger.setLevel(logging.DEBUG) | |
| formatter = logging.Formatter( | |
| "%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s") | |
| if not os.path.exists(model_dir): | |
| os.makedirs(model_dir) | |
| h = logging.FileHandler(os.path.join(model_dir, filename)) | |
| h.setLevel(logging.DEBUG) | |
| h.setFormatter(formatter) | |
| logger.addHandler(h) | |
| return logger | |
| class HParams(): | |
| def __init__(self, **kwargs): | |
| for k, v in kwargs.items(): | |
| if type(v) == dict: | |
| v = HParams(**v) | |
| self[k] = v | |
| def keys(self): | |
| return self.__dict__.keys() | |
| def items(self): | |
| return self.__dict__.items() | |
| def values(self): | |
| return self.__dict__.values() | |
| def __len__(self): | |
| return len(self.__dict__) | |
| def __getitem__(self, key): | |
| return getattr(self, key) | |
| def __setitem__(self, key, value): | |
| return setattr(self, key, value) | |
| def __contains__(self, key): | |
| return key in self.__dict__ | |
| def __repr__(self): | |
| return self.__dict__.__repr__() | |