def get_dataset(hp):
  """Return dataloader based on hyperparameters."""
  data_dir = hp.data_dir
  train_data_loader = UrmpMidi(data_dir, instrument_key=hp.instrument,
                               split='train', suffix='batched')
  training_data = train_data_loader.get_batch(batch_size=hp.batch_size,
                                              shuffle=True,
                                              repeats=-1,
                                              drop_remainder=True)
  test_data_loader = UrmpMidi(data_dir, instrument_key=hp.instrument,
                              split='test', suffix='batched')
  evaluation_data = test_data_loader.get_batch(batch_size=hp.batch_size * 3,
                                               shuffle=True,
                                               repeats=1,
                                               drop_remainder=False)

  length_training_data = get_tfrecord_length(train_data_loader)
  length_evaluation_data = get_tfrecord_length(test_data_loader)
  length_training_data = math.ceil(length_training_data / hp.batch_size)
  length_evaluation_data = math.ceil(
    length_evaluation_data / (hp.batch_size * 3))

  training_data = iter(training_data)

  return training_data, length_training_data, evaluation_data, \
         length_evaluation_data def print_hparams(hp):
  """Print the hyperparameters"""
  attributes = inspect.getmembers(hp, lambda a: not inspect.isroutine(a))
  return {a[0]: a[1] for a in attributes if
          not (a[0].startswith('__') and a[0].endswith('__'))} def set_seed(seed):
  """Set the random seed."""
  os.environ['PYTHONHASHSEED'] = str(seed)
  random.seed(seed)
  np.random.seed(seed)
  tf.random.set_seed(seed)
  os.environ['TF_DETERMINISTIC_OPS'] = '1' def save_results(pred_batch, ref_wav_batch, log_dir, save_folder_name,
                 sample_rate, acoustic_params=None,
                 separate_signal=None, file_name=None, bottleneck_feature=None,
                 note=None, no_pred_suffix=False):
  """Save the evaluation results."""
  if file_name and pred_batch.shape[0] > 1:
    raise RuntimeError('Got a file name to save, but has batch_size > 1')
  pred = pred_batch.numpy()
  if ref_wav_batch is not None:
    ref_wav = ref_wav_batch.numpy()
  if log_dir is None:
    save_dir = save_folder_name
  else:
    save_dir = os.path.join(log_dir,
                            f'results_{os.path.basename(save_folder_name)}')
  os.makedirs(save_dir, exist_ok=True)
  output_features_save_dir = os.path.join(save_dir, 'output_features')
  separate_signal_save_dir = os.path.join(save_dir, 'separate_signal')
  spectrogram_dir = os.path.join(save_dir, 'spectrogram')
  os.makedirs(output_features_save_dir, exist_ok=True)
  os.makedirs(separate_signal_save_dir, exist_ok=True)
  os.makedirs(spectrogram_dir, exist_ok=True)
  for i in range(pred.shape[0]):

    save_file_name = file_name if file_name else i

    pred_file_suffix = '' if no_pred_suffix else '_pred'
    save_wav(pred[i],
             os.path.join(save_dir, f'{save_file_name}{pred_file_suffix}.wav'),
             sample_rate=sample_rate)
    plot_spec(pred[i], sr=sample_rate,
              title=f'{save_file_name}{pred_file_suffix}.wav',
              save_path=os.path.join(spectrogram_dir,
                                     f'{save_file_name}{pred_file_suffix}.png'))
    if ref_wav_batch is not None:
      save_wav(ref_wav[i], os.path.join(save_dir, f'{save_file_name}_ref.wav'),
               sample_rate=sample_rate)
      plot_spec(ref_wav[i], sr=sample_rate, title=f'{save_file_name}_ref.wav',
                save_path=os.path.join(spectrogram_dir,
                                       f'{save_file_name}_ref.png'))
    if acoustic_params:
      for k in acoustic_params.keys():
        np.save(
          os.path.join(output_features_save_dir, f'{save_file_name}_{k}.npy'),
          acoustic_params[k][i].numpy())
    if separate_signal:
      for k in separate_signal.keys():
        save_wav(separate_signal[k][i].numpy(),
                 os.path.join(separate_signal_save_dir,
                              f'{save_file_name}_{k}.wav'),
                 sample_rate=sample_rate)
    if bottleneck_feature:
      for k in bottleneck_feature.keys():
        if 'loss' not in k and 'perplexity' not in k:
          np.save(
            os.path.join(output_features_save_dir, f'{save_file_name}_{k}.npy'),
            bottleneck_feature[k][i].numpy())
    if note:
      for k in note.keys():
        np.save(
          os.path.join(output_features_save_dir, f'{save_file_name}_{k}.npy'),
          note[k][i].numpy()) def str2bool(v):
  """Enable boolean in argparse by passing string."""
  # https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse
  if isinstance(v, bool):
    return v
  if v.lower() in ('yes', 'true', 't', 'y', '1'):
    return True
  elif v.lower() in ('no', 'false', 'f', 'n', '0'):
    return False
  else:
    raise argparse.ArgumentTypeError('Boolean value expected.') def write_tensorboard_audio(writer, data, outputs, step, tag):
  """
  Write audio and spectrogram to tensorboard summary file.
  Args:
    writer: writer of the tensorboard summary file.
    data: the data batch from the dataset.
    outputs: outputs from the model output.
    step: current training step.
    tag: tag for the tensorboard summary. Usually is "train" or "test"

  """
  with writer.as_default():
    spectrogram_summary(data['audio'], outputs['synth_audio'], step, name='',
                        tag=f'{tag}/spectrogram_synth_coder')
    audio_summary(data['audio'], step, sample_rate=16000,
                  name=f'{tag}_audio_original')
    audio_summary(outputs['synth_audio'], step, sample_rate=16000,
                  name=f'{tag}_audio_synth_coder')
    if 'midi_audio' in outputs.keys():
      spectrogram_summary(data['audio'], outputs['midi_audio'], step, name='',
                          tag=f'{tag}/spectrogram_midi_decoder')
      audio_summary(outputs['midi_audio'], step, sample_rate=16000,
                    name=f'{tag}_audio_midi_decoder') class hparams:
  """Hyperparameters for training Synthesis Generator."""
  # Learning parameters
  batch_size = 4
  clip_grad = 2.5
  lr = 3e-4
  seed = 1111

  # Training parameters
  training_steps = 100000
  log_interval = 100
  checkpoint_save_interval = 5000
  eval_interval = 5000
  mode = 'train'  # train, eval
  data_dir = None
  restore_path = None

  # Model Parameters
  nhid = 256
  sequence_length = 1000
  train_synth_coder_first = True
  midi_audio_loss = True
  add_synth_loss = False
  synth_params_loss = False
  midi_decoder_type = 'interpretable_conditioning'
  position_code = 'index_length'
  midi_decoder_decoder_net = 'rnn_synth_params'
  multi_instrument = True
  instrument = 'vn'
  synth_coder_training_steps = 300
  lambda_recon = 1.0
  use_gan = True
  lambda_G = 1
  sg_z = True
  lr_disc = 1e-4
  write_tfrecord_audio = False
  without_note_expression = False
  discriminator_dim = 256

  # Synthesis & DSP parameters
  nhramonic = 60
  nnoise = 65
  reverb = True
  reverb_length = 48000
  num_mels = 64
  n_fft = 1024
  sample_rate = 16000
  frame_size = 64
  hop_length = frame_size
  win_length = hop_length * 2
  frame_shift_ms = 1000 / sample_rate * frame_size
  fmin = 40 class ReconLossHelper(LossHelper):
  """Overall wrapper for getting the reconstruction loss helper
  based on hyperparameters."""

  def __init__(self, hp, eval_recon_loss=False):
    midi_audio_loss = hp.midi_audio_loss
    synth_params_loss = hp.synth_params_loss
    midi_decoder_type = hp.midi_decoder_type
    midi_decoder_decoder_net = hp.midi_decoder_decoder_net
    super().__init__()

    self.loss_list = [
      'loss_amp',
      'loss_f0',
      'loss_hd',
      'loss_noise',
      'loss_spectral_synth',
      'loss_spectral_midi',
      'total_loss'
    ]

    if eval_recon_loss:
      self.loss_group = EvalReconLossGroup()

    elif midi_decoder_type == 'midi_decoder' or \
          midi_decoder_type == 'unconditioned' or \
          midi_decoder_decoder_net == 'rnn_f0_ld':

      self.loss_group = F0LdReconLossGroup(midi_audio_loss=midi_audio_loss,
                                           synth_params_loss=synth_params_loss,
                                           full_f0_loss=
                                           midi_decoder_type == 'unconditioned')
    elif midi_decoder_type == 'interpretable_conditioning':
      if 'dilated_conv' in midi_decoder_decoder_net:
        self.loss_group = EvalReconLossGroup()
      else:
        self.loss_group = ReconLossGroup(midi_audio_loss=midi_audio_loss,
                                         synth_params_loss=synth_params_loss)

    self.init_metrics() class GANLossHelper(LossHelper):
  """Loss helper for handling and logging GAN loss.

  Args:
    lambda_recon: the reconstruction coefficient.
    lambda_feat: the feature matching loss coefficient.
    lambda_G: the generator loss coefficient.
    sg_z: whether to stop gradient on the feature sequence input to the
      Synthesis Generator.
  """

  def __init__(self, lambda_recon=1, lambda_feat=10, lambda_G=1, sg_z=True):
    super().__init__()
    self.loss_list = [
      'disc_loss',
      'gen_loss',
      'fm_loss'
    ]
    self.n_layers_D = 4
    self.num_D = 3
    self.lambda_feat = lambda_feat
    self.lambda_recon = lambda_recon
    self.lambda_G = lambda_G
    self.sg_z = sg_z
    self.init_metrics()

  def compute_disc_loss(self, D_fake_det, D_real):
    loss_D = 0
    for scale in D_fake_det:
      loss_D += tf.reduce_mean(scale[-1] ** 2)

    for scale in D_real:
      loss_D += tf.reduce_mean((scale[-1] - 1) ** 2)

    return {'disc_loss': loss_D}

  def compute_gen_loss(self, D_fake, D_real, total_loss_recon):
    loss_G = 0
    for scale in D_fake:
      loss_G += tf.reduce_mean((scale[-1] - 1) ** 2)

    loss_feat = 0
    feat_weights = 4.0 / (self.n_layers_D + 1)
    D_weights = 1.0 / self.num_D
    wt = D_weights * feat_weights
    for i in range(self.num_D):
      for j in range(len(D_fake[i]) - 1):
        loss_feat += wt * tf.reduce_mean(
          tf.abs(D_fake[i][j] - tf.stop_gradient(D_real[i][j])))  # MAE

    fm_loss = self.lambda_feat * loss_feat

    total_loss = self.lambda_recon * total_loss_recon + \
                 self.lambda_G * loss_G + fm_loss

    return total_loss, {'gen_loss': loss_G, 'fm_loss': fm_loss}

  def get_disc_input(self, outputs):
    if self.sg_z:
      cond = tf.stop_gradient(outputs['params_pred']['z_midi_decoder'])
    else:
      cond = outputs['params_pred']['z_midi_decoder']

    real_outputs = tf.concat([
      ddsp.core.hz_to_midi(outputs['synth_params']['f0_hz'] / F0_RANGE),
      outputs['synth_params']['amplitudes'],
      outputs['synth_params']['harmonic_distribution'],
      outputs['synth_params']['noise_magnitudes']
    ], -1)
    real_outputs = tf.stop_gradient(real_outputs)

    fake_outputs = tf.concat([
      ddsp.core.hz_to_midi(outputs['params_pred']['f0_hz'] / F0_RANGE),
      outputs['params_pred']['amplitudes'],
      outputs['params_pred']['harmonic_distribution'],
      outputs['params_pred']['noise_magnitudes']
    ], -1)

    return cond, real_outputs, fake_outputs def get_synthesis_generator(hp):
  if hp.midi_decoder_decoder_net == 'rnn_f0_ld' or \
        hp.midi_decoder_type == 'midi_decoder' or \
        hp.midi_decoder_type == 'unconditioned':
    use_f0_ld = True
    encoder = F0LDEncoder()
    decoder = FCStackHarmonicDecoder(hp.nhramonic, hp.nnoise)
  else:
    cnn = Cnn8(pool_size=(1, 2))
    encoder = MelF0LDEncoder(cnn, hp.nhid, hp.sample_rate, hp.win_length,
                             hp.hop_length, hp.n_fft, hp.num_mels,
                             hp.fmin)
    decoder = FCHarmonicDecoder(hp.nhramonic, hp.nnoise)
    use_f0_ld = False
  synth_coder = SynthCoder(encoder, decoder)

  midi_decoder = get_midi_decoder(hp)

  if hp.reverb:
    num_reverb = NUM_INST if hp.multi_instrument else 1
    reverb_module = ReverbModules(num_reverb=num_reverb,
                                  reverb_length=hp.reverb_length)
  else:
    reverb_module = None

  model = MIDIExpressionAE(
    synth_coder=synth_coder,
    midi_decoder=midi_decoder,
    n_frames=hp.sequence_length,
    frame_size=hp.frame_size,
    sample_rate=hp.sample_rate,
    reverb_module=reverb_module,
    use_f0_ld=use_f0_ld
  )

  return model def get_fake_data_synthesis_generator(hp):
  """Get the fake data for building the synthesis generator."""
  fake_data = {
    'audio': tf.random.normal([1, hp.sequence_length * hp.frame_size]),
    'mel': tf.random.normal([1, hp.sequence_length, hp.num_mels]),
    'f0_hz': tf.random.normal([1, hp.sequence_length, 1]),
    'loudness_db': tf.random.normal([1, hp.sequence_length, 1]),
    'midi': tf.ones([1, hp.sequence_length], dtype=tf.int64),
    'instrument_id': tf.ones([1], dtype=tf.int64),
    'onsets': tf.zeros([1, hp.sequence_length], dtype=tf.int64),
    'offsets': tf.zeros([1, hp.sequence_length], dtype=tf.int64),
  }
  return fake_data, {'training': True} class Discriminator(tfkl.Layer):
  """The multi-scale discriminator, has num_D discriminators at num_D scale,
   each discriminator consists of num_dblock blocks."""

  def __init__(self, num_D=3, num_dblock=4, nhid=128):
    super().__init__()
    self.model = {}
    for i in range(num_D):
      self.model[f'disc_{i}'] = NLayerDiscriminator(
        num_dblock, nhid
      )

    self.downsample = tfkl.AveragePooling1D(4, 2, padding='same')

  def call(self, inputs):
    results = []
    for _, disc in self.model.items():
      results.append(disc(inputs))
      inputs = (self.downsample(inputs[0]), self.downsample(inputs[1]))
    return results
import tensorflow as tf
import time
import os
import sys
import logging
import argparse
from midi_ddsp.data_handling.get_dataset import get_dataset
from midi_ddsp.utils.training_utils import print_hparams, set_seed, \
  save_results, str2bool
from midi_ddsp.utils.summary_utils import write_tensorboard_audio
from midi_ddsp.hparams_synthesis_generator import hparams as hp
from midi_ddsp.modules.recon_loss import ReconLossHelper
from midi_ddsp.modules.gan_loss import GANLossHelper
from midi_ddsp.modules.get_synthesis_generator import get_synthesis_generator, \
  get_fake_data_synthesis_generator
from midi_ddsp.modules.discriminator import Discriminator
                      help='The number of training steps to take '
                           'evaluation on whole evaluation set.')
  parser.add_argument('--checkpoint_save_interval', type=int,
                      default=hp.checkpoint_save_interval,
                      help='The number of training steps to take before save '
                           'the model weights once.')
  parser.add_argument('--data_dir', type=str, default=hp.data_dir,
                      help='The directory to tfrecord data files.')
  parser.add_argument('--restore_path', type=str, default=hp.restore_path,
                      help='The path to the model weights file for restore '
                           'training.')
  parser.add_argument('--midi_audio_loss', type=str2bool, nargs='?', const=True,
                      default=hp.midi_audio_loss,
                      help='Whether to use '
                           'multi-scale spectral loss on audio predicted')
  parser.add_argument('--synth_params_loss', type=str2bool, nargs='?',
                      const=True, default=hp.synth_params_loss,
                      help='Whether to use L1 loss on synthesis parameters '
                           'predicted.')
  parser.add_argument('--train_synth_coder_first', type=str2bool, nargs='?',
                      const=True, default=hp.train_synth_coder_first,
                      help='Whether to train DDSP Inference first or joint '
                           'training DDSP Inference and Synthesis Generator.')
  parser.add_argument('--add_synth_loss', type=str2bool, nargs='?', const=True,
                      default=hp.add_synth_loss,
                      help='Whether to add DDSP Inference loss to the '
                           'total loss for back-prop.')
  parser.add_argument('--multi_instrument', type=str2bool, nargs='?',
                      const=True, default=hp.multi_instrument,
                      help='Whether to train multi-instrument model or '
                           'single-instrument model')
  parser.add_argument('--position_code', type=str, default=hp.position_code,
                      help='Whether to use positional encoding in '
                           'Synthesis Generator, and what type of '
                           'positional encoding.')
  parser.add_argument('--midi_decoder_type', type=str,
                      default=hp.midi_decoder_type,
                      help='The type of Synthesis Generator.')
  parser.add_argument('--midi_decoder_decoder_net', type=str,
                      default=hp.midi_decoder_decoder_net,
                      help='The network type for Synthesis Generator.')
  parser.add_argument('--reverb', type=str2bool, nargs='?', const=True,
                      default=hp.reverb, help='Whether to use reverb.')
  parser.add_argument('--reverb_length', type=int, default=hp.reverb_length,
                      help='The length for the reverb impulse response.')
  parser.add_argument('--instrument', type=str, default=hp.instrument,
                      help='The instrument to train on for single instrument '
                           'setting.')
  parser.add_argument('--synth_coder_training_steps', type=int,
                      default=hp.synth_coder_training_steps,
                      help='The number of steps for training DDSP Inference.')
  parser.add_argument('--use_gan', type=str2bool, nargs='?', const=True,
                      default=hp.use_gan, help='Whether to use GAN training '
                                               'for Synthesis Generator.')
  parser.add_argument('--lambda_recon', type=float, default=hp.lambda_recon,
                      help='The weight for the reconstruction loss.')
  parser.add_argument('--lambda_G', type=float, default=hp.lambda_G,
                      help='The weight for the discriminator loss '
                           'for generator.')
  parser.add_argument('--sg_z', type=str2bool, nargs='?', const=True,
                      default=hp.sg_z, help='Whether to stop the gradient for'
                                            'conditioning sequence (z) in '
                                            'discriminator.')
  parser.add_argument('--lr_disc', type=float, default=hp.lr_disc,
                      help='The learning rate for discriminator.')
  parser.add_argument('--without_note_expression', type=str2bool, nargs='?',
                      const=True, default=hp.without_note_expression,
                      help='Whether to train Synthesis Generator using no note '
                           'expression controls. Used in paper rebuttal.')
  parser.add_argument('--mode', type=str, default=hp.mode,
                      help='The mode to run, train or eval.')
  parser.add_argument('--name', type=str, default='logs_synthesis_generator',
                      help='The name of the log folder.')

  # Change hp according to argparse.
  args = parser.parse_args()
  for k, v in vars(args).items():
    setattr(hp, k, v)

  # Create Synthesis Generator
  model = get_synthesis_generator(hp)
  model._build(get_fake_data_synthesis_generator(hp))

  # Create optimizer, loss helper and discriminator.
  scheduler = tf.keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate=hp.lr, decay_steps=1000,
    decay_rate=0.99)
  optimizer = tf.keras.optimizers.Adam(learning_rate=scheduler)

  loss_helper = ReconLossHelper(hp)
  gan_loss_helper = GANLossHelper(lambda_recon=hp.lambda_recon,
                                  lambda_G=hp.lambda_G, sg_z=hp.sg_z)
  optimizer_disc = tf.keras.optimizers.Adam(learning_rate=hp.lr_disc)

  net_D = Discriminator(nhid=hp.discriminator_dim)
  # 64=instrument_emb_dim
  z_dim = hp.discriminator_dim + int(hp.multi_instrument) * 64
  # synth_params_dim = dim(nharmonic + nnoise + amplitude + f0)
  synth_params_dim = hp.nhramonic + hp.nnoise + 2
  _ = net_D((tf.random.normal([4, 1000, z_dim]),
             tf.random.normal([4, 1000, synth_params_dim])))

  # Load model, create log directory and log file.
  log_dir = f'logs/{args.name}'

  if hp.restore_path:
    model.load_weights(hp.restore_path)
    log_dir = os.path.dirname(hp.restore_path)

  writer = tf.summary.create_file_writer(log_dir)
  log_path = os.path.join(log_dir, 'train.log')
  logging.basicConfig(level=logging.INFO,
                      format='%(asctime)s - %(levelname)s: %(message)s',
                      handlers=[
                        logging.FileHandler(log_path),
                        logging.StreamHandler(sys.stdout)]
                      )

  # Load dataset.
  training_data, length_training_data, \

  evaluation_data, length_evaluation_data = get_dataset(hp)