Fixes to Train_VST.ipynb

jesseengel · Magenta Team · commit 9742103dc11c · 2022-08-26T12:10:19.000-07:00
* No longer crashes, can now handle model directory paths that have spaces.

Seems like something changed about the way !bash commands were parsing strings/variables that caused things to crashed. Fixed by just doing all string manipulation in python explicitly.

PiperOrigin-RevId: 470303927
diff --git a/ddsp/colab/demos/Train_VST.ipynb b/ddsp/colab/demos/Train_VST.ipynb
@@ -128,7 +128,7 @@
         "from google.colab import drive\n",
         "import tensorflow as tf\n",
         "\n",
-        "!pip install ipyfilechooser \u0026\u003e /dev/null\n",
+        "!pip install ipyfilechooser==0.6.0 \u0026\u003e /dev/null\n",
         "from ipyfilechooser import FileChooser\n",
         "\n",
         "# ------------------------------------------------------------------------------\n",
@@ -227,9 +227,12 @@
         "    print('(10 minutes of training audio -\u003e 20-30 minutes)')\n",
         "\n",
         "    audio_filepattern = os.path.join(audio_dir, '*')\n",
+        "    audio_fp_str = f'\"{audio_filepattern}\"'    \n",
+        "    tfrecord_path_str = f'\"{data_dir}/train.tfrecord\"'\n",
+        "\n",
         "    !ddsp_prepare_tfrecord \\\n",
-        "    --input_audio_filepatterns=$audio_filepattern \\\n",
-        "    --output_tfrecord_path=$data_dir/train.tfrecord \\\n",
+        "    --input_audio_filepatterns=$audio_fp_str \\\n",
+        "    --output_tfrecord_path=$tfrecord_path_str \\\n",
         "    --num_shards=10 \\\n",
         "    --sample_rate=$sample_rate \\\n",
         "    --frame_rate=$frame_rate \\\n",
@@ -241,12 +244,13 @@
         "\n",
         "def train(model_dir, data_dir, steps=30000):\n",
         "  file_pattern = os.path.join(data_dir, 'train.tfrecord*')\n",
+        "  fp_str = f\"TFRecordProvider.file_pattern='{file_pattern}'\"\n",
         "  !ddsp_run \\\n",
         "  --mode=train \\\n",
         "  --save_dir=\"$model_dir\" \\\n",
         "  --gin_file=models/vst/vst.gin \\\n",
         "  --gin_file=datasets/tfrecord.gin \\\n",
-        "  --gin_param=\"TFRecordProvider.file_pattern='$file_pattern'\" \\\n",
+        "  --gin_param=\"$fp_str\" \\\n",
         "  --gin_param=\"TFRecordProvider.centered=True\" \\\n",
         "  --gin_param=\"TFRecordProvider.frame_rate=50\" \\\n",
         "  --gin_param=\"batch_size=16\" \\\n",
@@ -261,21 +265,25 @@
         "\n",
         "\n",
         "def reset_state(data_dir, audio_dir, model_dir):\n",
+        "  model_dir_str = f'\"{model_dir}\"'\n",
         "  if tf.io.gfile.exists(data_dir):\n",
         "    !rm -r $data_dir\n",
         "    !rm -r $audio_dir\n",
         "  !mkdir -p $data_dir\n",
         "  !mkdir -p $audio_dir\n",
-        "  !mkdir -p $model_dir\n",
+        "  !mkdir -p $model_dir_str\n",
         "\n",
         "\n",
         "def export_and_download(model_dir, model_name=Name):\n",
         "  export_path = os.path.join(model_dir, model_name)\n",
         "\n",
+        "  model_dir_str=f'\"{model_dir}\"'\n",
+        "  export_path_str=f'\"{export_path}\"'\n",
+        "  \n",
         "  !ddsp_export \\\n",
         "  --name=$model_name \\\n",
-        "  --model_path=$model_dir \\\n",
-        "  --save_dir=$export_path \\\n",
+        "  --model_path=$model_dir_str \\\n",
+        "  --save_dir=$export_path_str \\\n",
         "  --inference_model=vst_stateless_predict_controls \\\n",
         "  --tflite \\\n",
         "  --notfjs\n",
@@ -284,8 +292,7 @@
         "  zip_fname = f'{model_name}.zip'\n",
         "  zip_fp = os.path.join(model_dir, zip_fname)\n",
         "  print(f'Export complete! Zipping {export_path} to {zip_fp}')\n",
-        "  !cd $model_dir \u0026\u0026 zip -r $zip_fname ./$model_name\n",
-        "  # !zip -jr $zip_fp $export_path\n",
+        "  !cd $model_dir_str \u0026\u0026 zip -r $zip_fname ./$model_name\n",
         "\n",
         "  # Download.\n",
         "  print(f'Zipping Complete! Downloading... {zip_fname}')\n",
