Repository: magenta/magenta

Issue:
MusicVAE - Add generate/export script



Discussion:

Resolving patch:
diff --git a/magenta/models/music_vae/music_vae_generate.py b/magenta/models/music_vae/music_vae_generate.py
new file mode 100644
--- /dev/null
+++ b/magenta/models/music_vae/music_vae_generate.py
@@ -0,0 +1,194 @@
+# Copyright 2018 Google 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.
+"""MusicVAE generation script."""
+
+# TODO(adarob): Add support for models with conditioning.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import sys
+import time
+
+# internal imports
+import numpy as np
+import tensorflow as tf
+
+from magenta import music as mm
+from magenta.models.music_vae import configs
+from magenta.models.music_vae import TrainedModel
+
+flags = tf.app.flags
+logging = tf.logging
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string(
+    'run_dir', None,
+    'Path to the directory where the latest checkpoint will be loaded from.')
+flags.DEFINE_string(
+    'checkpoint_file', None,
+    'Path to the checkpoint file. run_dir will take priority over this flag.')
+flags.DEFINE_string(
+    'output_dir', '/tmp/music_vae/generated',
+    'The directory where MIDI files will be saved to.')
+flags.DEFINE_string(
+    'config', None,
+    'The name of the config to use.')
+flags.DEFINE_string(
+    'mode', 'sample',
+    'Generate mode (either `sample` or `interpolate`).')
+flags.DEFINE_string(
+    'input_midi_1', None,
+    'Path of start MIDI file for interpolation.')
+flags.DEFINE_string(
+    'input_midi_2', None,
+    'Path of end MIDI file for interpolation.')
+flags.DEFINE_integer(
+    'num_outputs', 5,
+    'In `sample` mode, the number of samples to produce. In `interpolate` '
+    'mode, the number of steps (including the endpoints).')
+flags.DEFINE_integer(
+    'max_batch_size', 8,
+    'The maximum batch size to use. Decrease if you are seeing an OOM.')
+flags.DEFINE_float(
+    'temperature', 0.5,
+    'The randomness of the decoding process.')
+flags.DEFINE_string(
+    'log', 'INFO',
+    'The threshold for what messages will be logged: '
+    'DEBUG, INFO, WARN, ERROR, or FATAL.')
+
+
+def _slerp(p0, p1, t):
+  """Spherical linear interpolation."""
+  omega = np.arccos(
+      np.dot(np.squeeze(p0/np.linalg.norm(p0)),
+             np.squeeze(p1/np.linalg.norm(p1))))
+  so = np.sin(omega)
+  return np.sin((1.0-t)*omega) / so * p0 + np.sin(t*omega)/so * p1
+
+
+def run(config_map):
+  """Load model params, save config file and start trainer.
+
+  Args:
+    config_map: Dictionary mapping configuration name to Config object.
+
+  Raises:
+    ValueError: if required flags are missing or invalid.
+  """
+  date_and_time = time.strftime('%Y-%m-%d_%H%M%S')
+
+  if FLAGS.run_dir is None == FLAGS.checkpoint_file is None:
+    raise ValueError(
+        'Exactly one of `--run_dir` or `--checkpoint_file` must be specified.')
+  if FLAGS.output_dir is None:
+    raise ValueError('`--output_dir` is required.')
+  tf.gfile.MakeDirs(FLAGS.output_dir)
+  if FLAGS.mode != 'sample' and FLAGS.mode != 'interpolate':
+    raise ValueError('Invalid value for `--mode`: %s' % FLAGS.mode)
+
+  if FLAGS.config not in config_map:
+    raise ValueError('Invalid config name: %s' % FLAGS.config)
+  config = config_map[FLAGS.config]
+  config.data_converter.max_tensors_per_item = None
+
+  if FLAGS.mode == 'interpolate':
+    if FLAGS.input_midi_1 is None or FLAGS.input_midi_2 is None:
+      raise ValueError(
+          '`--input_midi_1` and `--input_midi_2` must be specified in '
+          '`interpolate` mode.')
+    input_midi_1 = os.path.expanduser(FLAGS.input_midi_1)
+    input_midi_2 = os.path.expanduser(FLAGS.input_midi_2)
+    if not os.path.exists(input_midi_1):
+      raise ValueError('Input MIDI 1 not found: %s' % FLAGS.input_midi_1)
+    if not os.path.exists(input_midi_2):
+      raise ValueError('Input MIDI 2 not found: %s' % FLAGS.input_midi_2)
+    input_1 = mm.midi_to_sequence_proto(tf.gfile.GFile(input_midi_1).read())
+    input_2 = mm.midi_to_sequence_proto(tf.gfile.GFile(input_midi_2).read())
+
+    def _check_extract_examples(input_ns, path, input_number):
+      """Make sure each input returns exactly one example from the converter."""
+      tensors = config.data_converter.to_tensors(input_ns).outputs
+      if not tensors:
+        print(
+            'MusicVAE configs have very specific input requirements. Could not '
+            'extract any valid inputs from `%s`. Try another MIDI file.' % path)
+        sys.exit()
+      elif len(tensors) > 1:
+        basename = os.path.join(
+            FLAGS.output_dir,
+            '%s_input%d-extractions_%s-*-of-%03d.mid' %
+            (FLAGS.config, input_number, date_and_time, len(tensors)))
+        for i, ns in enumerate(config.data_converter.to_notesequences(tensors)):
+          mm.sequence_proto_to_midi_file(ns, basename.replace('*', '%03d' % i))
+        print(
+            '%d valid inputs extracted from `%s`. Outputting these potential '
+            'inputs as `%s`. Call script again with one of these instead.' %
+            (len(tensors), path, basename))
+        sys.exit()
+    logging.info(
+        'Attempting to extract examples from input MIDIs using config `%s`...',
+        FLAGS.config)
+    _check_extract_examples(input_1, FLAGS.input_midi_1, 1)
+    _check_extract_examples(input_2, FLAGS.input_midi_2, 2)
+
+  logging.info('Loading model...')
+  checkpoint_dir_or_path = os.path.expanduser(
+      os.path.join(FLAGS.run_dir, 'train')
+      if FLAGS.run_dir else FLAGS.checkpoint_file)
+  model = TrainedModel(
+      config, batch_size=min(FLAGS.max_batch_size, FLAGS.num_outputs),
+      checkpoint_dir_or_path=checkpoint_dir_or_path)
+
+  if FLAGS.mode == 'interpolate':
+    logging.info('Interpolating...')
+    _, mu, _ = model.encode([input_1, input_2])
+    z = np.array([
+        _slerp(mu[0], mu[1], t) for t in np.linspace(0, 1, FLAGS.num_outputs)])
+    results = model.decode(
+        length=config.hparams.max_seq_len,
+        z=z,
+        temperature=FLAGS.temperature)
+  elif FLAGS.mode == 'sample':
+    logging.info('Sampling...')
+    results = model.sample(
+        n=FLAGS.num_outputs,
+        length=config.hparams.max_seq_len,
+        temperature=FLAGS.temperature)
+
+  basename = os.path.join(
+      FLAGS.output_dir,
+      '%s_%s_%s-*-of-%03d.mid' %
+      (FLAGS.config, FLAGS.mode, date_and_time, FLAGS.num_outputs))
+  logging.info('Outputting %d files as `%s`...', FLAGS.num_outputs, basename)
+  for i, ns in enumerate(results):
+    mm.sequence_proto_to_midi_file(ns, basename.replace('*', '%03d' % i))
+
+  logging.info('Done.')
+
+
+def main(unused_argv):
+  logging.set_verbosity(FLAGS.log)
+  run(configs.CONFIG_MAP)
+
+
+def console_entry_point():
+  tf.app.run(main)
+
+
+if __name__ == '__main__':
+  console_entry_point()
diff --git a/magenta/models/music_vae/trained_model.py b/magenta/models/music_vae/trained_model.py
--- a/magenta/models/music_vae/trained_model.py
+++ b/magenta/models/music_vae/trained_model.py
@@ -17,9 +17,12 @@
 from __future__ import print_function
 
 import copy
+import os
 import re
+import tarfile
 
 # internal imports
+from backports import tempfile
 import numpy as np
 import tensorflow as tf
 
@@ -112,7 +115,19 @@ def __init__(self, config, batch_size, checkpoint_dir_or_path=None,
       # Restore graph
       self._sess = tf.Session(target=session_target)
       saver = tf.train.Saver(var_map)
-      saver.restore(self._sess, checkpoint_path)
+      if tarfile.is_tarfile(checkpoint_path):
+        tf.logging.info('Unbundling checkpoint.')
+        with tempfile.TemporaryDirectory() as temp_dir:
+          tar = tarfile.open(checkpoint_path)
+          tar.extractall(temp_dir)
+          # Assume only a single checkpoint is in the directory.
+          for name in tar.getnames():
+            if name.endswith('.ckpt.index'):
+              checkpoint_path = os.path.join(temp_dir, name[0:-6])
+              break
+          saver.restore(self._sess, checkpoint_path)
+      else:
+        saver.restore(self._sess, checkpoint_path)
 
   def sample(self, n=None, length=None, temperature=1.0, same_z=False,
              c_input=None):
diff --git a/magenta/tools/pip/setup.py b/magenta/tools/pip/setup.py
--- a/magenta/tools/pip/setup.py
+++ b/magenta/tools/pip/setup.py
@@ -34,6 +34,7 @@
 REQUIRED_PACKAGES = [
     'IPython',
     'Pillow >= 3.4.2',
+    'backports.tempfile',
     'bokeh >= 0.12.0',
     'intervaltree >= 2.1.0',
     'joblib >= 0.12',
@@ -75,6 +76,7 @@
     'magenta.models.melody_rnn.melody_rnn_create_dataset',
     'magenta.models.melody_rnn.melody_rnn_generate',
     'magenta.models.melody_rnn.melody_rnn_train',
+    'magenta.models.music_vae.music_vae_generate',
     'magenta.models.music_vae.music_vae_train',
     'magenta.models.nsynth.wavenet.nsynth_generate',
     'magenta.models.nsynth.wavenet.nsynth_save_embeddings',


Test patch: