text: # Copyright 2020 LMNT, Inc. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  
  import numpy as np
  import os
  import torch
  import torch.nn as nn
  
  from torch.utils.tensorboard import SummaryWriter
  from tqdm import tqdm
  
  from wavegrad.model import WaveGrad
  
  def set_init_dict(model_dict, checkpoint_state):
      # Partial initialization: if there is a mismatch with new and old layer, it is skipped.
      for k, v in checkpoint_state.items():
          if k not in model_dict:
              print(" | > Layer missing in the model definition: {}".format(k))
      # 1. filter out unnecessary keys
      pretrained_dict = {
          k: v
          for k, v in checkpoint_state.items() if k in model_dict
      }
      # 2. filter out different size layers
      pretrained_dict = {
          k: v
          for k, v in pretrained_dict.items()
          if v.numel() == model_dict[k].numel()
      }
      # 4. overwrite entries in the existing state dict
      model_dict.update(pretrained_dict)
      print(" | > {} / {} layers are restored.".format(len(pretrained_dict),
                                                       len(model_dict)))
      return model_dict
  
  def _nested_map(struct, map_fn):
    if isinstance(struct, tuple):
      return tuple(_nested_map(x, map_fn) for x in struct)
    if isinstance(struct, list):
      return [_nested_map(x, map_fn) for x in struct]
    if isinstance(struct, dict):
      return { k: _nested_map(v, map_fn) for k, v in struct.items() }
    return map_fn(struct)
  
  
  class WaveGradLearner:
    def __init__(self, model_dir, model, dataset, optimizer, params, *args, **kwargs):
      os.makedirs(model_dir, exist_ok=True)
      self.model_dir = model_dir
      self.model = model
      self.dataset = dataset
      self.optimizer = optimizer
      self.params = params
      self.autocast = torch.cuda.amp.autocast(enabled=kwargs.get('fp16', False))
      self.scaler = torch.cuda.amp.GradScaler(enabled=kwargs.get('fp16', False))
      self.step = 0
  
      beta = np.array(self.params.noise_schedule)
      noise_level = np.cumprod(1 - beta)**0.5
      noise_level = np.concatenate([[1.0], noise_level], axis=0)
      self.noise_level = torch.tensor(noise_level.astype(np.float32))
      self.loss_fn = nn.L1Loss()
      self.summary_writer = None
  
    def state_dict(self):
      return {
          'step': self.step,
          'model': { k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in self.model.state_dict().items() },
          'optimizer': { k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in self.optimizer.state_dict().items() },
          'params': dict(self.params),
          'scaler': self.scaler.state_dict(),
      }
  
    def load_state_dict(self, state_dict):
      try:
        self.model.load_state_dict(state_dict['model'])
        self.optimizer.load_state_dict(state_dict['optimizer'])
        print(" > Total model initialization.")
      except:
        print(" > Partial model initialization, and the optimizer was not loaded !")
        model_dict = self.model.state_dict()
        model_dict = set_init_dict(model_dict, state_dict['model'])
        self.model.load_state_dict(model_dict)
        del model_dict
  
      self.scaler.load_state_dict(state_dict['scaler'])
      self.step = state_dict['step']
  
    def save_to_checkpoint(self, filename='weights'):
      save_basename = f'{filename}-{self.step}.pt'
      save_name = f'{self.model_dir}/{save_basename}'
      link_name = f'{self.model_dir}/{filename}.pt'
      torch.save(self.state_dict(), save_name)
      try:
        if os.name == 'nt':
          torch.save(self.state_dict(), link_name)
        else:
          if os.path.islink(link_name):
            os.unlink(link_name)
          os.symlink(save_basename, link_name)
      except OSError:
          pass
  
    def restore_from_checkpoint(self, filename='weights'):
      try:
        checkpoint = torch.load(f'{self.model_dir}/{filename}.pt')
        self.load_state_dict(checkpoint)
        return True
      except FileNotFoundError:
        return False
  
    def train(self, max_steps=None):
      device = next(self.model.parameters()).device
      while True:
        for features in tqdm(self.dataset, desc=f'Epoch {self.step // len(self.dataset)}'):
          if max_steps is not None and self.step >= max_steps:
            return
          features = _nested_map(features, lambda x: x.to(device) if isinstance(x, torch.Tensor) else x)
          loss = self.train_step(features)
          if torch.isnan(loss).any():
            raise RuntimeError(f'Detected NaN loss at step {self.step}.')
          if self.step % 100 == 0:
            self._write_summary(self.step, features, loss)
          if self.step % (len(self.dataset)*50) == 0:
            self.save_to_checkpoint()
          self.step += 1
  
    def train_step(self, features):
      self.optimizer.zero_grad()
  
      audio = features['audio']
      spectrogram = features['spectrogram']
  
      N, T = audio.shape
      S = 1000
      device = audio.device
      self.noise_level = self.noise_level.to(device)
  
      with self.autocast:
        s = torch.randint(1, S + 1, [N], device=audio.device)
        l_a, l_b = self.noise_level[s-1], self.noise_level[s]
        noise_scale = l_a + torch.rand(N, device=audio.device) * (l_b - l_a)
        noise_scale = noise_scale.unsqueeze(1)
        noise = torch.randn_like(audio)
        noisy_audio = noise_scale * audio + (1.0 - noise_scale**2)**0.5 * noise
  
        predicted = self.model(noisy_audio, spectrogram, noise_scale.squeeze(1))
        loss = self.loss_fn(noise, predicted.squeeze(1))
  
      self.scaler.scale(loss).backward()
      self.scaler.unscale_(self.optimizer)
      self.grad_norm = nn.utils.clip_grad_norm_(self.model.parameters(), self.params.max_grad_norm)
      self.scaler.step(self.optimizer)
      self.scaler.update()
      return loss
  
    def _write_summary(self, step, features, loss):
      writer = self.summary_writer or SummaryWriter(self.model_dir, purge_step=step)
      writer.add_audio('audio/reference', features['audio'][0], step, sample_rate=self.params.sample_rate)
      writer.add_scalar('train/loss', loss, step)
      writer.add_scalar('train/grad_norm', self.grad_norm, step)
      writer.flush()
      self.summary_writer = writer
  
  
  def train(dataset, args, params):
    model = WaveGrad(params).cuda()
    opt = torch.optim.Adam(model.parameters(), lr=params.learning_rate)
  
    learner = WaveGradLearner(args.model_dir, model, dataset, opt, params, fp16=args.fp16)
    learner.restore_from_checkpoint()
    learner.train(max_steps=args.max_steps)
  