NUM_INST = 20 class SynthCoder(tf.keras.Model):
  """The DDSP Inference module."""

  def __init__(self, encoder, decoder):
    super().__init__()
    self.encoder = encoder
    self.decoder = decoder
    self.model_type = 'SynthCoder'

  def encode(self, inputs, training=False):
    z = self.encoder(inputs, training=training)
    return z

  def decode(self, z, inputs):
    synth_params = self.decoder([z, inputs])
    return synth_params

  def call(self, inputs, training=None):
    z = self.encode(inputs, training=training)
    synth_params = self.decode(z, inputs)
    return synth_params

  def _build(self, inputs):
    inputs, kwargs = inputs
    self(inputs, **kwargs) class MIDIExpressionAE(tf.keras.Model):
  """
  The DDSP Inference + Synthesis Generator module.
  The naming of the module is not exactly the same as in the MIDI-DDSP paper,
  but one can find correspondence in following crude diagram:

                                       Note Sequence
                                             │
          Note Expression Controls           │       │
            (conditioning_dict)    ──────────┤       │
                    ▲                        │       │
      Feature       │                        ├────┐  │
     Extraction     │                        │    │  │
  (gen_cond_dict)   │                        ▼    │  │ Synthesis Generator
           Synthesis Parameters             f0    │  │   (midi_decoder)
            (f0,amps,hd,noise)               │    │  │
                    ▲                        ├────┘  │
      DDSP          │                        │       │
   Inference        │                        │       │
  (SynthCoder)      │                        │       │
                    │                        ▼
              Audio Feature         Synthesis Parameters
               (f0,ld,mel)           (f0,amps,hd,noise)
                    ▲                        │
      CREPE         │                        │
   A-Weighting      │                        │ DDSP
     Mel-STFT       │                        │
                    │                        ▼
                  Audio                    Audio
                  Input               Reconstruction
  """

  def __init__(self, synth_coder, midi_decoder, n_frames=1000, frame_size=64,
               sample_rate=16000, reverb_module=None,
               use_f0_ld=False):
    super().__init__()
    self.synth_coder = synth_coder
    self.midi_decoder = midi_decoder
    self.reverb_module = reverb_module
    self.model_type = 'MIDIExpressionAE'

    self.processor_group = get_process_group(n_frames, frame_size, sample_rate,
                                             use_angular_cumsum=False)
    self.n_frames = n_frames
    self.frame_size = frame_size
    self.sample_rate = sample_rate

    self.use_f0_ld = use_f0_ld

  def train_synth_coder_only(self):
    self.midi_decoder.trainable = False
    self.synth_coder.trainable = True

  def freeze_synth_coder(self):
    self.midi_decoder.trainable = True
    self.synth_coder.trainable = False
    if self.reverb_module is not None:
      self.reverb_module.trainable = False

  def run_synth_coder(self, features, training=True):
    synth_params = self.synth_coder(features, training=training)
    control_params = self.processor_group.get_controls(synth_params,
                                                       verbose=False)
    synth_audio = self.processor_group.get_signal(control_params)
    if self.reverb_module is not None:
      synth_audio = self.reverb_module(synth_audio,
                                       reverb_number=features['instrument_id'],
                                       training=training)
    return synth_params, control_params, synth_audio

  @staticmethod
  def get_gt_midi(features):
    """Prepare the ground truth MIDI."""
    f0_midi = ddsp.core.hz_to_midi(features['f0_hz'])
    q_pitch = tf.cast(features['midi'][..., tf.newaxis], tf.float32)

    q_vel = q_pitch * 0.0  # (Jesse) Don't use velocities for now.

    f0_loss_weights = tf.cast(tf.abs(f0_midi - q_pitch) < 2.0, tf.float32)

    return q_pitch, q_vel, f0_loss_weights, features['onsets'], features[
      'offsets']

  def gen_cond_dict_from_feature(self, features, training=False):
    synth_params, control_params, synth_audio = self.run_synth_coder(features,
                                                                     training=
                                                                     training)

    # synth_params_normalized: scaled and normalized synth params
    synth_params_normalized = extract_harm_controls(control_params)

    midi_features = self.get_gt_midi(features)

    conditioning_dict = self.midi_decoder.gen_cond_dict(synth_params_normalized,
                                                        midi_features)

    return synth_params_normalized, midi_features, conditioning_dict

  def gen_audio_from_cond_dict(self, conditioning_dict, midi_features,
                               instrument_id=None,
                               use_angular_cumsum=True,
                               display_progressbar=False):
    z_midi_decoder, params_pred = self.midi_decoder.gen_params_from_cond(
      conditioning_dict, midi_features,
      instrument_id=instrument_id,
      display_progressbar=display_progressbar)

    self.processor_group = get_process_group(z_midi_decoder.shape[1],
                                             self.frame_size, self.sample_rate,
                                             use_angular_cumsum=
                                             use_angular_cumsum)

    if self.use_f0_ld:
      midi_synth_params, midi_control_params, midi_audio = self.run_synth_coder(
        {'f0_hz': params_pred['f0_hz'],
         'loudness_db': params_pred['ld'],
         'instrument_id': instrument_id},
        training=False)
      f0_pred, amps_pred, hd_pred, noise_pred = extract_harm_controls(
        midi_control_params)

    else:
      midi_synth_params = self.processor_group.get_controls(
        {'amplitudes': params_pred['amplitudes'],
         'harmonic_distribution': params_pred[
           'harmonic_distribution'],
         'noise_magnitudes': params_pred['noise_magnitudes'],
         'f0_hz': params_pred['f0_hz'], },
        verbose=False)

      midi_audio = self.processor_group.get_signal(midi_synth_params)
      if self.reverb_module is not None:
        midi_audio = self.reverb_module(midi_audio, reverb_number=instrument_id,
                                        training=False)
      f0_pred, amps_pred, hd_pred, noise_pred = extract_harm_controls(
        midi_synth_params)

    midi_control_params = (f0_pred, amps_pred, hd_pred, noise_pred)
    return midi_audio, midi_control_params, midi_synth_params

  def call(self, features, training=False, run_synth_coder_only=False):
    """Run the network to get a prediction and optionally get losses."""

    synth_params, control_params, synth_audio = self.run_synth_coder(features,
                                                                     training=
                                                                     training)
    if run_synth_coder_only:  # only run synth coder branch
      outputs = {
        'synth_params': synth_params,
        'synth_audio': synth_audio,
      }
      return outputs

    # synth_params_normalized: scaled and normalized synth params
    synth_params_normalized = extract_harm_controls(control_params,
                                                    stop_gradient=True)

    midi_features = self.get_gt_midi(features)

    # --- MIDI Decoding
    conditioning_dict, params_pred = self.midi_decoder(features,
                                                       synth_params_normalized,
                                                       midi_features,
                                                       training=training,
                                                       synth_params=
                                                       control_params)

    if self.use_f0_ld:
      midi_synth_params, midi_control_params, midi_audio = self.run_synth_coder(
        {'f0_hz': params_pred['f0_hz'],
         'loudness_db': params_pred['ld'],
         'instrument_id': features[
           'instrument_id']
         },
        training=training)

      f0_pred, amps_pred, hd_pred, noise_pred = extract_harm_controls(
        midi_control_params)

    else:
      midi_synth_params = self.processor_group.get_controls(
        {'amplitudes': params_pred['amplitudes'],
         'harmonic_distribution': params_pred[
           'harmonic_distribution'],
         'noise_magnitudes': params_pred['noise_magnitudes'],
         'f0_hz': params_pred['f0_hz'], },
        verbose=False)
      # --- MIDI Audio Losses
      midi_audio = self.processor_group.get_signal(midi_synth_params)

      if self.reverb_module is not None:
        midi_audio = self.reverb_module(midi_audio,
                                        reverb_number=features['instrument_id'],
                                        training=training)

      f0_pred, amps_pred, hd_pred, noise_pred = extract_harm_controls(
        midi_synth_params)

    # unpack things
    f0, amps, hd, noise = synth_params_normalized
    q_pitch, q_vel, f0_loss_weights, onsets, offsets = midi_features

    # --- Finalize and return
    outputs = {
      'synth_params': synth_params,
      'control_params': control_params,
      'synth_audio': synth_audio,
      'midi_synth_params': midi_synth_params,
      'midi_audio': midi_audio,
      'q_pitch': q_pitch,
      'q_vel': q_vel,
      'conditioning_dict': conditioning_dict,
      'amps': amps,
      'hd': hd,
      'noise': noise,
      'amps_pred': amps_pred,
      'hd_pred': hd_pred,
      'noise_pred': noise_pred,
      'f0_loss_weights': f0_loss_weights,
      'params_pred': params_pred,
    }

    return outputs

  def _build(self, inputs):
    inputs, kwargs = inputs
    self(inputs, **kwargs) class MelF0LDEncoder(tfkl.Layer):
  """The encoder in DDSP Inference.
  The MelF0LDEncoder takes input of audio, loudness and f0.
  The MelF0LDEncoder extract features from audio using an 8-layer CNN,
  and extract features from loudness and f0 using fully-connected layers.
  Then, a bi-lstm is used to extract contextual features from the extracted
  features.
  """

  def __init__(self, cnn, nhid, sample_rate, win_length, hop_length, n_fft,
               num_mels, fmin):
    super().__init__()
    self.nhid = nhid
    self.cnn = cnn
    self.z_fc = tfkl.Dense(nhid)
    self.f0_ld_fc = tfkl.Dense(nhid)
    self.rnn = tfkl.Bidirectional(
      tfkl.LSTM(units=nhid, return_sequences=True), name='bilstm'
    )
    # TODO(yusongwu): change emb dim to 64
    self.instrument_emb = tfkl.Embedding(NUM_INST, 128)

    # mel-spec parameters
    self.sample_rate = sample_rate
    self.win_length = win_length
    self.hop_length = hop_length
    self.n_fft = n_fft
    self.num_mels = num_mels
    self.fmin = fmin

  def call(self, inputs, training=False):
    mel = tf_log_mel(inputs['audio'],
                     self.sample_rate,
                     self.win_length,
                     self.hop_length,
                     self.n_fft,
                     self.num_mels,
                     self.fmin)
    z_cnn = self.cnn(mel, training=training)
    z_reduce = self.z_fc(z_cnn)
    instrument_z = tf.tile(
      self.instrument_emb(inputs['instrument_id'])[:, tf.newaxis, :],
      [1, z_cnn.shape[1], 1])
    x = tf.concat([ddsp.core.hz_to_midi(inputs['f0_hz']) / F0_RANGE,
                   inputs['loudness_db'] / DB_RANGE], -1)
    x_z = self.f0_ld_fc(x)
    z_out = self.rnn(tf.concat([x_z, z_reduce, instrument_z], -1))
    return z_out class F0LDEncoder(tfkl.Layer):
  """The encoder of original DDSP autoencoder."""

  # TODO: (yusongwu) To be removed and use the decoders.RnnFcDecoder
  def __init__(self):
    super().__init__()
    self.nhid = 512
    self.f0_fc = nn.FcStack(self.nhid, layers=3)
    self.ld_fc = nn.FcStack(self.nhid, layers=3)
    self.instrument_emb = tfkl.Embedding(NUM_INST, 128)
    self.rnn = tfkl.GRU(
      units=self.nhid, return_sequences=True,  # dropout=0.2,
    )

  def call(self, inputs, training=False):
    z_f0 = self.f0_fc(ddsp.core.hz_to_midi(inputs['f0_hz']) / F0_RANGE)
    z_ld = self.ld_fc(inputs['loudness_db'] / DB_RANGE)
    instrument_z = tf.tile(
      self.instrument_emb(inputs['instrument_id'])[:, tf.newaxis, :],
      [1, z_ld.shape[1], 1])
    x_z = tf.concat([z_f0, z_ld, instrument_z], -1)
    z_out = self.rnn(x_z)
    z_out = tf.concat([x_z, z_out], -1)
    return z_out class FCHarmonicDecoder(tfkl.Layer):
  """The decoder in DDSP Inference.
  The FCHarmonicDecoder takes input of a feature sequence,
  and output the synthesis parameters for DDSP through fully-connected layers.
  """

  def __init__(self, nhramonic=100, nnoise=65):
    super().__init__()

    self.harmonic_amp_fc = tfkl.Dense(1, bias_initializer='ones')
    self.harmonic_distribution_fc = tfkl.Dense(nhramonic)

    self.noise_mag_fc = tfkl.Dense(nnoise)

  def get_synth_params(self, inputs):
    z, data = inputs

    harmonic_amp = self.harmonic_amp_fc(z)
    harmonic_distribution = self.harmonic_distribution_fc(z)
    noise_mag = self.noise_mag_fc(z)

    synth_params = {
      'f0_hz': data['f0_hz'],
      'amplitudes': harmonic_amp,
      'harmonic_distribution': harmonic_distribution,
      'noise_magnitudes': noise_mag,
    }

    return synth_params

  def call(self, inputs):
    synth_params = self.get_synth_params(inputs)

    return synth_params class FCStackHarmonicDecoder(tfkl.Layer):
  """The decoder original DDSP autoencoder.
  The FCStackHarmonicDecoder takes input of a feature sequence,
  and output the synthesis parameters for DDSP through stacked MLP.
  """

  def __init__(self, nharmonic=100, nnoise=65):
    super().__init__()

    self.output_splits = (
      ('amplitudes', 1), ('harmonic_distribution', nharmonic),
      ('noise_magnitudes', nnoise))
    self.n_out = sum([v[1] for v in self.output_splits])
    self.out_stack = nn.FcStack(512, layers=3)
    self.dense_out = tfkl.Dense(self.n_out)

  def get_synth_params(self, inputs):
    z, data = inputs

    z_output = self.out_stack(z)
    synth_params = nn.split_to_dict(self.dense_out(z_output),
                                    self.output_splits)

    synth_params['f0_hz'] = data['f0_hz']

    return synth_params

  def call(self, inputs):
    synth_params = self.get_synth_params(inputs)
    return synth_params class Cnn8(tfkl.Layer):
  """
  A tensorflow implementation of CNN8 used in audioset classification.
  This CNN has better performance when used in spectrogram feature
  extraction for audio tagging. Adapted from pytorch implementation:
  https://github.com/qiuqiangkong/audioset_tagging_cnn.
  paper: https://arxiv.org/abs/1912.10211.
  """

  def __init__(self, pool_size=(2, 2), dropout=0.2):
    super().__init__()

    self.conv_block1 = ConvBlock(out_channels=64, pool_size=pool_size)
    self.conv_block2 = ConvBlock(out_channels=128, pool_size=pool_size)
    self.conv_block3 = ConvBlock(out_channels=256, pool_size=pool_size)
    self.conv_block4 = ConvBlock(out_channels=512, pool_size=pool_size)

    self.dropout = tfkl.Dropout(rate=dropout)

  def call(self, x, training=None):
    x = x[..., tf.newaxis]

    x = self.conv_block1(x, pool_type='avg', training=training)
    x = self.dropout(x, training=training)
    x = self.conv_block2(x, pool_type='avg', training=training)
    x = self.dropout(x, training=training)
    x = self.conv_block3(x, pool_type='avg', training=training)
    x = self.dropout(x, training=training)
    x = self.conv_block4(x, pool_type='avg', training=training)
    x = self.dropout(x, training=training)
    x = tf.reshape(x, [x.shape[0], x.shape[1], -1])

    return x class ReverbModules(tfkl.Layer):
  """A reverb module for multi-instrument training.
  An unique reverb parameter is used for each instrument id.
  """

  def __init__(self, num_reverb=1, reverb_length=48000):
    super().__init__()
    self.num_reverb = num_reverb
    self.reverb_length = reverb_length
    self.reverb = ddsp.effects.Reverb(trainable=False,
                                      reverb_length=reverb_length)
    initializer = tf.random_normal_initializer(mean=0, stddev=1e-6)
    self.magnitudes_embedding = tfkl.Embedding(num_reverb, reverb_length,
                                               embeddings_initializer=
                                               initializer)

  def call(self, audio, reverb_number=0, training=False):
    batch_size = audio.shape[0]
    if isinstance(reverb_number, int):
      reverb_number = tf.repeat(tf.constant([reverb_number], dtype=tf.int64),
                                batch_size)
    if self.num_reverb == 1 or reverb_number is None:  # unified reverb
      reverb_number = tf.repeat(tf.constant([0], dtype=tf.int64), batch_size)
    ir_magnitudes = self.magnitudes_embedding(reverb_number)
    if not training:
      ir_magnitudes = ir_magnitudes * get_exp_decay(
        self.reverb_length, 16000, 4)[tf.newaxis, ...]
    output = self.reverb(audio, ir_magnitudes)
    return output class MidiToSynthAutoregDecoder(tfkl.Layer):
  """Predicting synthesis parameters from MIDI and note expression controls
   using autoregressive RNN for f0 and dilated convolution networks for
   amplitudes, harmonic distribution and noise magnitudes."""
  def __init__(self, nhid=256,
               hd_noise_output_splits=(('amplitudes', 1),
                                       ('harmonic_distribution', 60),
                                       ('noise_magnitudes', 65))):
    super().__init__()
    self.q_pitch_emb = tfkl.Dense(64)
    self.midi_to_f0 = MidiExpreToF0AutoregDecoder(nhid=nhid)
    self.midi_f0_to_harmonic = MidiExpreToSynthDecoder(
      net=nn.DilatedConvStack(
        ch=128,
        layers_per_stack=5,
        stacks=4,
        norm_type='layer',
        conditional=False,
      ),
      output_splits=hd_noise_output_splits,
    )

  def call(self, q_pitch, z_midi_decoder, conditioning_dict, out=None,
           training=False, display_progressbar=False):
    f0_output = self.midi_to_f0(q_pitch, z_midi_decoder, conditioning_dict,
                                out=out, training=training,
                                display_progressbar=display_progressbar)

    if training:
      f0_midi = ddsp.core.hz_to_midi(out['f0_hz'])
    else:
      f0_midi = f0_output['f0_midi']

    z = tf.concat([z_midi_decoder, self.q_pitch_emb(f0_midi / 127)], -1)
    outputs = self.midi_f0_to_harmonic(z, training=training)

    outputs['f0_output'] = f0_output
    if training:
      outputs['f0_hz'] = out['f0_hz']
    else:
      outputs['f0_hz'] = ddsp.core.midi_to_hz(f0_output['f0_midi'],
                                             midi_zero_silence=True)

    return outputs class MidiToF0LDAutoregDecoder(TwoLayerCondAutoregRNN):
  """Autoregressive prediction of f0 and loudnes from z_midi_decoder that
   only contains MIDI. (The implementation of MIDI2Params.)"""
  def __init__(self, nhid=256, norm=True, dropout=0.5, sampling_method='top_p'):
    super().__init__(
      nhid=nhid,
      n_out=RELATIVE_F0_BINS + ABSOLUTE_LD_BINS,
      input_dropout=True,
      input_dropout_p=0.5,
      dropout=dropout,
    )
    self.birnn = tfkl.Bidirectional(tfkl.GRU(
      units=nhid, return_sequences=True, dropout=dropout
    ), )
    self.dense_out = nn.FcStackOut(ch=nhid, layers=2, n_out=self.n_out)
    self.dense_in = nn.FcStack(ch=nhid, layers=2)
    self.norm = nn.Normalize('layer') if norm else None
    self.sampling_method = sampling_method

  def encode_z(self, z, training=False):
    return self.dense_in(z)

  def split_teacher_force_output(self, output):
    f0_midi_dv_logits = output[..., :RELATIVE_F0_BINS]
    f0_midi_dv_onehot = get_onehot_f0_dv(get_float_f0_dv(f0_midi_dv_logits))
    ld_logits = output[...,
                RELATIVE_F0_BINS:RELATIVE_F0_BINS + ABSOLUTE_LD_BINS]
    ld_onehot = get_onehot_ld(get_float_ld(ld_logits))
    output = {
      'f0_midi_dv_logits': f0_midi_dv_logits,
      'f0_midi_dv_onehot': f0_midi_dv_onehot,
      'ld_logits': ld_logits,
      'ld_onehot': ld_onehot,
    }
    return output

  def encode_cond(self, cond, training=False):
    cond = self.birnn(cond, training=training)
    return cond

  def decode_out(self, z_out, training=False):
    if self.norm is not None:
      z_out = self.norm(z_out)
    output = self.dense_out(z_out)
    return output

  def sample_out(self, out, training=False):
    f0_midi_dv_logits = out[..., :RELATIVE_F0_BINS]
    f0_midi_dv_onehot = sample_from(f0_midi_dv_logits, return_onehot=True,
                                    sample_method=self.sampling_method)
    ld_logits = out[..., RELATIVE_F0_BINS:RELATIVE_F0_BINS + ABSOLUTE_LD_BINS]
    ld_onehot = sample_from(ld_logits, return_onehot=True,
                            sample_method=self.sampling_method)
    output = {
      'f0_midi_dv_logits': f0_midi_dv_logits,
      'f0_midi_dv_onehot': f0_midi_dv_onehot,
      'ld_logits': ld_logits,
      'ld_onehot': ld_onehot,
      'curr_out': tf.concat([f0_midi_dv_onehot, ld_onehot], -1)
    }
    return output

  def preprocess(self, q_pitch, out):
    f0_midi = ddsp.core.hz_to_midi(out['f0_hz'])
    f0_midi_dv = f0_midi - q_pitch  # f0 residual
    f0_midi_dv_onehot = get_onehot_f0_dv(f0_midi_dv)
    ld = out['loudness_db']
    ld_onehot = get_onehot_ld(ld)
    return tf.concat([f0_midi_dv_onehot, ld_onehot], -1)

  def postprocess(self, outputs, q_pitch, training=False):
    outputs['f0_midi_dv'] = get_float_f0_dv(outputs['f0_midi_dv_onehot'])
    outputs['f0_midi'] = outputs['f0_midi_dv'] + q_pitch
    outputs['f0_hz'] = ddsp.core.midi_to_hz(outputs['f0_midi'])
    outputs['ld'] = get_float_ld(outputs['ld_onehot'])
    return outputs

  def call(self, q_pitch, z_midi_decoder, out=None, training=False):

    if training:
      out = self.preprocess(q_pitch, out)
      outputs = self.teacher_force(z_midi_decoder, out, training=training)
    else:
      outputs = self.autoregressive(z_midi_decoder, training=training)

    outputs = self.postprocess(outputs, q_pitch)

    return outputs

  def partial_teacher_force(self, q_pitch, z_midi_decoder, out,
                            teacher_force_mask, training=False,
                            display_progressbar=False):
    """For the expression generator reconstruction comparison.
    (Table 1 (c) in paper)"""
    # right shift and encode out
    out = self.preprocess(q_pitch, out)
    go_frame = tf.tile([[0.0]], [out.shape[0], out.shape[-1]])[:, tf.newaxis, :]
    out = tf.concat([go_frame, out[:, :-1, :]], 1)
    out = self.encode_out(out)

    cond = self.encode_cond(z_midi_decoder, training=training)
    batch_size = cond.shape[0]
    length = cond.shape[1]
    prev_out = tf.tile([[0.0]], [batch_size, self.n_out])[:, tf.newaxis,
               :]  # go_frame
    prev_states = (None, None)
    overall_outputs = []

    teacher_force_batch_size = teacher_force_mask.shape[0]

    for i in tqdm(range(length), position=0, leave=True, desc='Generating: ',
                  disable=not display_progressbar):
      curr_cond = tf.tile(cond[:, i:i + 1, :], [teacher_force_batch_size, 1, 1])
      prev_out = self.encode_out(prev_out) * (
            1 - teacher_force_mask[:, i:i + 1, :]) + \
                 out[:, i:i + 1, :] * teacher_force_mask[:, i:i + 1, :]
      curr_out, curr_states = self._one_step(curr_cond, prev_out, prev_states,
                                             training=training)
      curr_out = self.sample_out(curr_out)
      overall_outputs.append(curr_out)
      prev_out, prev_states = curr_out['curr_out'], curr_states

    outputs = {}
    for k in curr_out:
      outputs[k] = tf.concat([x[k] for x in overall_outputs], 1)

    return self.postprocess(outputs, q_pitch) class F0LDAutoregDecoder(TwoLayerCondAutoregRNN):
  """Autoregressive prediction of f0 and loudness.
  (A fully unconditional model without MIDI input.)"""
  def __init__(self, nhid=256, norm=True, dropout=0.5, sampling_method='top_p'):
    self.n_f0 = ABSOLUTE_F0_BINS
    self.n_ld = ABSOLUTE_LD_BINS
    super().__init__(
      nhid=nhid,
      n_out=self.n_f0 + self.n_ld,
      input_dropout=True,
      input_dropout_p=0.5,
      dropout=dropout,
    )
    self.birnn = tfkl.Bidirectional(tfkl.GRU(
      units=nhid, return_sequences=True, dropout=dropout
    ), )
    self.dense_out = nn.FcStackOut(ch=nhid, layers=2, n_out=self.n_out)
    self.dense_in = nn.FcStack(ch=nhid, layers=2)
    self.norm = nn.Normalize('layer') if norm else None
    self.sampling_method = sampling_method

  def split_teacher_force_output(self, output):
    # split outputs
    f0_midi_logits = output[..., :self.n_f0]
    f0_midi_onehot = get_onehot_absolute_f0(
      get_float_absolute_f0(f0_midi_logits))
    ld_logits = output[..., self.n_f0:self.n_f0 + self.n_ld]
    ld_onehot = get_onehot_ld(get_float_ld(ld_logits))
    output = {
      'f0_midi_logits': f0_midi_logits,
      'f0_midi_onehot': f0_midi_onehot,
      'ld_logits': ld_logits,
      'ld_onehot': ld_onehot,
    }
    return output

  def encode_z(self, z, training=False):
    return self.dense_in(z)

  def decode_out(self, z_out, training=False):
    if self.norm is not None:
      z_out = self.norm(z_out)
    output = self.dense_out(z_out)
    return output

  def sample_out(self, out, training=False):
    f0_midi_logits = out[..., :self.n_f0]
    f0_midi_onehot = sample_from(f0_midi_logits, return_onehot=True,
                                 sample_method=self.sampling_method)
    ld_logits = out[..., self.n_f0:self.n_f0 + self.n_ld]
    ld_onehot = sample_from(ld_logits, return_onehot=True,
                            sample_method=self.sampling_method)
    output = {
      'f0_midi_logits': f0_midi_logits,
      'f0_midi_onehot': f0_midi_onehot,
      'ld_logits': ld_logits,
      'ld_onehot': ld_onehot,
      'curr_out': tf.concat([f0_midi_onehot, ld_onehot], -1)
    }
    return output

  def preprocess(self, out, training=False):
    f0_midi = ddsp.core.hz_to_midi(out['f0_hz'])
    f0_midi_onehot = get_onehot_absolute_f0(f0_midi)
    ld = out['loudness_db']
    ld_onehot = get_onehot_ld(ld)
    return tf.concat([f0_midi_onehot, ld_onehot], -1)

  def postprocess(self, outputs, training=False):
    outputs['f0_midi'] = get_float_absolute_f0(outputs['f0_midi_onehot'])
    outputs['f0_hz'] = ddsp.core.midi_to_hz(outputs['f0_midi'])
    outputs['ld'] = get_float_ld(outputs['ld_onehot'])
    return outputs

  def call(self, q_pitch, z_midi_decoder, out=None, training=False):
    z_midi_decoder = z_midi_decoder[..., 66:]

    if training:
      out = self.preprocess(out)
      outputs = self.teacher_force(z_midi_decoder, out, training=training)
    else:
      outputs = self.autoregressive(z_midi_decoder, training=training)

    outputs = self.postprocess(outputs)

    return outputs class MidiNoiseToHarmonicDecoder(nn.DictLayer):
  """Decodes MIDI notes (& velocities) to f0, amps, hd, noise."""

  def __init__(self,
               net=None,
               f0_residual=True,
               norm=True,
               output_splits=(('f0_midi', 1),
                              ('amplitudes', 1),
                              ('harmonic_distribution', 60),
                              ('magnitudes', 65)),
               **kwargs):
    """Constructor."""
    self.output_splits = output_splits
    self.n_out = sum([v[1] for v in output_splits])
    output_keys = [v[0] for v in output_splits] + ['f0_hz']
    super().__init__(output_keys=output_keys, **kwargs)

    # Layers.
    self.net = net
    self.f0_residual = f0_residual
    self.dense_out = tfkl.Dense(self.n_out)
    self.norm = nn.Normalize('layer') if norm else None

  def call(self, noise, z_pitch, z=None):
    """Forward pass for the MIDI decoder.

    Args:
      z_pitch: Tensor containing encoded pitch in MIDI scale. [batch, time, 1].
      z_vel: Tensor containing encoded velocity in MIDI scale. [batch, time, 1].
      z: Additional non-MIDI latent tensor. [batch, time, n_z]

    Returns:
      A dictionary to feed into a processor group.
    """
    x = noise
    x = self.net(x) if z is None else self.net([x, z])

    if self.norm is not None:
      x = self.norm(x)

    x = self.dense_out(x)

    outputs = nn.split_to_dict(x, self.output_splits)

    if self.f0_residual:
      outputs['f0_midi'] += z_pitch

    outputs['f0_hz'] = ddsp.core.midi_to_hz(outputs['f0_midi'])
    return outputs class ExpressionMidiDecoder(tfkl.Layer):
  """Decoder wrapper for predicting synthesis parameters from MIDI and
  note expression controls."""

  def __init__(self, decoder, z_preconditioning_stack, multi_instrument=False,
               position_code='none', decoder_type='dilated_conv',
               without_note_expression=False):
    self.decoder_type = decoder_type
    self.z_preconditioning_stack = z_preconditioning_stack
    self.decoder = decoder
    self.position_code = position_code
    self.multi_instrument = multi_instrument
    if multi_instrument:
      self.instrument_emb = tfkl.Embedding(NUM_INST, 64)
    self.without_note_expression = without_note_expression
    super().__init__()

  def gen_params_from_cond(self, conditioning_dict, midi_features,
                           instrument_id=None, synth_params=None,
                           training=False, display_progressbar=False):
    q_pitch, q_vel, f0_loss_weights, onsets, offsets = midi_features
    # note-wise conditioning

    if self.without_note_expression:
      z_conditioning = tf.concat([q_pitch / 127,
                                  tf.cast(onsets, tf.float32)[..., tf.newaxis],
                                  tf.cast(offsets, tf.float32)[..., tf.newaxis]
                                  ], -1)
    else:
      z_conditioning = tf.stop_gradient(
        tf.concat(list(conditioning_dict.values()), -1))
      z_conditioning = tf.concat([z_conditioning,
                                  q_pitch / 127,
                                  tf.cast(onsets, tf.float32)[..., tf.newaxis],
                                  tf.cast(offsets, tf.float32)[..., tf.newaxis]
                                  ], -1)

    if self.position_code == 'index_length':
      note_mask = ddsp.training.nn.get_note_mask_from_onset(q_pitch, onsets)
      each_note_idx = tf.cumsum(note_mask, axis=1) * tf.cast(~(note_mask == 0),
                                                             tf.float32)
      each_note_len = tf.reduce_max(each_note_idx, axis=1,
                                    keepdims=True) * tf.cast(each_note_idx > 0,
                                                             tf.float32)
      each_note_idx = tf.reduce_sum(each_note_idx, axis=-1)[..., tf.newaxis]
      each_note_len = tf.reduce_sum(each_note_len, axis=-1)[..., tf.newaxis]
      relative_position = tf.math.divide_no_nan(each_note_idx, each_note_len)
      z_conditioning = tf.concat([z_conditioning, relative_position], -1)
    elif self.position_code == 'sinusoidal':
      note_mask = ddsp.training.nn.get_note_mask_from_onset(q_pitch, onsets)
      each_note_idx = tf.cumsum(note_mask, axis=1) * tf.cast(~(note_mask == 0),
                                                             tf.float32)
      each_note_idx = tf.reduce_sum(each_note_idx, axis=-1)[..., tf.newaxis]
      pos_code = positional_encoding(each_note_idx.numpy().astype(np.int64), 64)
      z_conditioning = tf.concat([z_conditioning, pos_code], -1)

    # --- Precondition
    z_midi_decoder = self.z_preconditioning_stack(z_conditioning)
    if self.multi_instrument:
      instrument_z = tf.tile(
        self.instrument_emb(instrument_id)[:, tf.newaxis, :],
        [1, z_midi_decoder.shape[1], 1])
      z_midi_decoder = tf.concat([z_midi_decoder, instrument_z], -1)

    # --- MIDI Decoding
    if self.decoder_type == 'dilated_conv':
      params_pred = self.decoder(q_pitch, q_vel, z_midi_decoder)
    elif self.decoder_type == 'noise_dilated_conv':
      noise = tf.random.normal(
        [z_midi_decoder.shape[0], z_midi_decoder.shape[1], 100])
      params_pred = self.decoder(noise, q_pitch, z_midi_decoder)
    elif 'rnn' in self.decoder_type:
      params_pred = self.decoder(q_pitch, z_midi_decoder, conditioning_dict,
                                 out=synth_params, training=training,
                                 display_progressbar=display_progressbar)
    # midi_decoder: [q_pitch, z_midi_decoder] -> synth params
    return z_midi_decoder, params_pred

  def gen_cond_dict(self, synth_params_normalized, midi_features):
    f0, amps, hd, noise = synth_params_normalized
    q_pitch, q_vel, f0_loss_weights, onsets, offsets = midi_features

    f0_midi_gt = ddsp.core.midi_to_hz(q_pitch, midi_zero_silence=True)
    conditioning = get_interpretable_conditioning(f0_midi_gt, f0, amps, hd,
                                                  noise)

    # --- Z Note Encoding
    conditioning_dict = get_conditioning_dict(conditioning, q_pitch, onsets,
                                              pool_type='note_pooling')
    return conditioning_dict

  def call(self, features, synth_params_normalized, midi_features,
           training=False, synth_params=None):
    conditioning_dict = self.gen_cond_dict(synth_params_normalized,
                                           midi_features)

    instrument_id = features['instrument_id'] if self.multi_instrument else None
    if self.decoder_type == 'rnn_f0_ld':
      synth_params = features
    z_midi_decoder, params_pred = self.gen_params_from_cond(conditioning_dict,
                                                            midi_features,
                                                            instrument_id=
                                                            instrument_id,
                                                            synth_params=
                                                            synth_params,
                                                            training=training)

    params_pred['z_midi_decoder'] = z_midi_decoder

    return conditioning_dict, params_pred class MidiDecoder(tfkl.Layer):
  """Decoder wrapper for predicting synthesis parameters from only MIDI."""

  def __init__(self, decoder, multi_instrument=False):
    super().__init__()
    self.decoder = decoder
    self.multi_instrument = multi_instrument
    if multi_instrument:
      self.instrument_emb = tfkl.Embedding(NUM_INST, 64)
    self.pitch_emb = tfkl.Embedding(128, 64)

  def call(self, features, synth_params_normalized, midi_features,
           training=False, synth_params=None):
    instrument_id = features['instrument_id'] if self.multi_instrument else None
    q_pitch, q_vel, f0_loss_weights, onsets, offsets = midi_features
    instrument_z = tf.tile(self.instrument_emb(instrument_id)[:, tf.newaxis, :],
                           [1, q_pitch.shape[1], 1])
    z_midi_decoder = tf.concat(
      [self.pitch_emb(tf.cast(q_pitch, tf.int64)[..., 0]),
       tf.cast(onsets, tf.float32)[..., tf.newaxis],
       tf.cast(offsets, tf.float32)[..., tf.newaxis],
       instrument_z],
      -1)  # HACK
    params_pred = self.decoder(q_pitch, z_midi_decoder, out=features,
                               training=training)
    return {}, params_pred
import ddsp.training
import tensorflow as tf
from ddsp.training import nn
from midi_ddsp.data_handling.instrument_name_utils import NUM_INST
from midi_ddsp.modules.model import SynthCoder, MIDIExpressionAE
from midi_ddsp.modules.ddsp_inference import MelF0LDEncoder, F0LDEncoder, \
  FCHarmonicDecoder, FCStackHarmonicDecoder, Cnn8
from midi_ddsp.modules.reverb_modules import ReverbModules
from midi_ddsp.modules.synth_params_decoder import MidiToSynthAutoregDecoder, \
  MidiToF0LDAutoregDecoder, F0LDAutoregDecoder, \
  MidiNoiseToHarmonicDecoder
from midi_ddsp.modules.midi_decoder import ExpressionMidiDecoder, \
  MidiDecoder
#  Copyright 2022 The MIDI-DDSP Authors.
#  #
#  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.

"""Get Synthesis Generator module from hyperparameters."""




def get_midi_decoder(hp):
  """Get midi_decoder from hyperparameters."""
  z_preconditioning_stack = nn.FcStackOut(ch=512, layers=5, n_out=256)
  if hp.midi_decoder_decoder_net == 'dilated_conv':
    net = nn.DilatedConvStack(
      ch=256,
      layers_per_stack=5,
      stacks=4,
      norm_type='layer',
      conditional=True,
    )
    midi_decoder_decoder = ddsp.training.decoders.MidiToHarmonicDecoder(
      net=net,
      f0_residual=True,
      norm=True,
      output_splits=(('f0_midi', 1),
                     ('amplitudes', 1),
                     ('harmonic_distribution', 60),
                     ('noise_magnitudes', 65)),
    )
  elif hp.midi_decoder_decoder_net == 'noise_dilated_conv':
    net = nn.DilatedConvStack(
      ch=128,
      layers_per_stack=5,
      stacks=4,
      norm_type='layer',
      conditional=True,
    )
    midi_decoder_decoder = MidiNoiseToHarmonicDecoder(
      net=net,
      f0_residual=True,
      norm=True,
      output_splits=(('f0_midi', 1),
                     ('amplitudes', 1),
                     ('harmonic_distribution', 60),
                     ('noise_magnitudes', 65)),
    )
  elif hp.midi_decoder_decoder_net == 'rnn_synth_params':
    midi_decoder_decoder = MidiToSynthAutoregDecoder()

  elif hp.midi_decoder_decoder_net == 'rnn_f0_ld':
    midi_decoder_decoder = MidiToF0LDAutoregDecoder()

  if hp.midi_decoder_type == 'interpretable_conditioning':
    midi_decoder = ExpressionMidiDecoder(
      decoder=midi_decoder_decoder,
      decoder_type=hp.midi_decoder_decoder_net,
      z_preconditioning_stack=z_preconditioning_stack,
      multi_instrument=hp.multi_instrument,
      position_code=hp.position_code,
      without_note_expression=hp.without_note_expression,
    )

  elif hp.midi_decoder_type == 'midi_decoder':
    midi_decoder_decoder = MidiToF0LDAutoregDecoder(
      sampling_method='random')
    midi_decoder = MidiDecoder(
      decoder=midi_decoder_decoder,
      multi_instrument=hp.multi_instrument,
    )
  elif hp.midi_decoder_type == 'unconditioned':
    midi_decoder_decoder = F0LDAutoregDecoder(sampling_method='random')
    midi_decoder = MidiDecoder(
      decoder=midi_decoder_decoder,
      multi_instrument=hp.multi_instrument,
    )

  return midi_decoder


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)