Have model ddsp_export do renaming instead of colab notebook.

jesseengel · Magenta Team · commit 0f6132f00b28 · 2022-05-09T23:16:56.000-07:00
PiperOrigin-RevId: 447651496
diff --git a/ddsp/colab/demos/Train_VST.ipynb b/ddsp/colab/demos/Train_VST.ipynb
@@ -273,17 +273,13 @@
         "  export_path = os.path.join(model_dir, model_name)\n",
         "\n",
         "  !ddsp_export \\\n",
+        "  --name=$model_name \\\n",
         "  --model_path=$model_dir \\\n",
         "  --save_dir=$export_path \\\n",
         "  --inference_model=vst_stateless_predict_controls \\\n",
         "  --tflite \\\n",
         "  --notfjs\n",
         "\n",
-        "  # Rename tflite model.\n",
-        "  tflite_old_fp = os.path.join(export_path, 'tflite', 'model.tflite')\n",
-        "  tflite_new_fp = os.path.join(export_path, 'tflite', f'{model_name}.tflite')\n",
-        "  !mv $tflite_old_fp $tflite_new_fp\n",
-        "\n",
         "  # Zip the whole directory.\n",
         "  zip_fname = f'{model_name}.zip'\n",
         "  zip_fp = os.path.join(model_dir, zip_fname)\n",
@@ -361,7 +357,7 @@
         "  # ------------------------------------------------------------------------------\n",
         "  print('Installing DDSP...')\n",
         "  print('This should take about 2 minutes...')\n",
-        "  !pip install -U ddsp[data_preparation]==3.4.1 \u0026\u003e /dev/null\n",
+        "  !pip install -U ddsp[data_preparation]==3.4.3 \u0026\u003e /dev/null\n",
         "\n",
         "  # ------------------------------------------------------------------------------\n",
         "  # Import DDSP\n",
diff --git a/ddsp/training/ddsp_export.py b/ddsp/training/ddsp_export.py
@@ -48,10 +48,14 @@
 import tensorflow as tf
 from tensorflowjs.converters import converter
 
+# pylint: disable=pointless-string-statement
 
 from tflite_support import metadata as _metadata
+# pylint: enable=pointless-string-statement
 
-
+flags.DEFINE_string(
+    'name', '', 'Name of your model to use as folder and filename on export. '
+    'Defaults to "export/" and "model.tflite" if none is provided.')
 flags.DEFINE_string(
     'model_path', '', 'Path to checkpoint or SavedModel directory. If no '
     'SavedModel is found, will search for latest checkpoint '
@@ -236,7 +240,10 @@ def saved_model_to_tfjs(input_dir, save_dir):
   print('TFJS Conversion Success!')
 
 
-def saved_model_to_tflite(input_dir, save_dir, metadata_file=None):
+def saved_model_to_tflite(input_dir,
+                          save_dir,
+                          metadata_file=None,
+                          name=''):
   """Convert SavedModel to TFLite model."""
   print(f'\nConverting to TFLite:\nInput:{input_dir}\nOutput:{save_dir}\n')
   # Convert the model.
@@ -247,7 +254,8 @@ def saved_model_to_tflite(input_dir, save_dir, metadata_file=None):
   ]
   tflite_model = tflite_converter.convert()  # Byte string.
   # Save the model.
-  save_path = os.path.join(save_dir, 'model.tflite')
+  name = name if name else 'model'
+  save_path = os.path.join(save_dir, f'{name}.tflite')
   with tf.io.gfile.GFile(save_path, 'wb') as f:
     f.write(tflite_model)
 
@@ -302,6 +310,7 @@ def main(unused_argv):
   is_ckpt = not tf.io.gfile.isdir(model_path)
 
   # Infer save directory path.
+  export_name = FLAGS.name if FLAGS.name else 'export'
   if FLAGS.save_dir:
     save_dir = FLAGS.save_dir
   else:
@@ -310,10 +319,10 @@ def main(unused_argv):
       save_dir = model_path
     elif is_ckpt:
       # If model_path is a checkpoint file, use the directory of the file.
-      save_dir = os.path.join(os.path.dirname(model_path), 'export')
+      save_dir = os.path.join(os.path.dirname(model_path), export_name)
     else:
       # If model_path is a checkpoint directory, use child export directory.
-      save_dir = os.path.join(model_path, 'export')
+      save_dir = os.path.join(model_path, export_name)
 
   # Make a new save directory.
   save_dir = train_util.expand_path(save_dir)
@@ -344,8 +353,10 @@ def main(unused_argv):
   if FLAGS.tflite:
     tflite_dir = os.path.join(save_dir, 'tflite')
     ensure_exits(tflite_dir)
-    saved_model_to_tflite(save_dir, tflite_dir,
-                          metadata_path if FLAGS.metadata else '')
+    saved_model_to_tflite(save_dir,
+                          tflite_dir,
+                          metadata_path if FLAGS.metadata else '',
+                          name=FLAGS.name)
 
 
 def console_entry_point():
diff --git a/ddsp/training/postprocessing.py b/ddsp/training/postprocessing.py
@@ -254,24 +254,22 @@ def fit_transform(self, x):
 def compute_dataset_statistics(data_provider,
                                batch_size=1,
                                power_frame_size=1024,
-                               power_frame_rate=50,
-                               legacy=False):
+                               power_frame_rate=50):
   """Calculate dataset stats.
 
   Args:
     data_provider: A DataProvider from ddsp.training.data.
     batch_size: Iterate over dataset with this batch size.
     power_frame_size: Calculate power features on the fly with this frame size.
     power_frame_rate: Calculate power features on the fly with this frame rate.
-    legacy: Use the 'audio' key instead of 'audio_16k'.
 
   Returns:
     Dictionary of dataset statistics. This is an overcomplete set of statistics,
     as there are now several different tone transfer implementations (js, colab,
     vst) that need different statistics for normalization.
   """
   print('Calculating dataset statistics for', data_provider)
-  data_iter = iter(data_provider.get_batch(batch_size, repeats=1))
+  ds = data_provider.get_batch(batch_size, repeats=1)
 
   # Unpack dataset.
   i = 0
@@ -281,9 +279,10 @@ def compute_dataset_statistics(data_provider,
   f0_conf = []
   audio = []
 
-  audio_key = 'audio' if legacy else 'audio_16k'
+  batch = next(iter(ds))
+  audio_key = 'audio_16k' if 'audio_16k' in batch.keys() else 'audio'
 
-  for batch in data_iter:
+  for batch in iter(ds):
     loudness.append(batch['loudness_db'])
     power.append(
         spectral_ops.compute_power(batch[audio_key],
diff --git a/ddsp/version.py b/ddsp/version.py
@@ -18,4 +18,4 @@
 pulling in all the dependencies in __init__.py.
 """
 
-__version__ = '3.4.2'
+__version__ = '3.4.3'
