1 files changed, 14 insertions, 5 deletions
diff --git a/app/client/modules/samplernn/samplernn.actions.js b/app/client/modules/samplernn/samplernn.actions.js
index e3e735f..fe73bf2 100644
--- a/app/client/modules/samplernn/samplernn.actions.js
+++ b/app/client/modules/samplernn/samplernn.actions.js
@@ -24,6 +24,7 @@ export const load_directories = (id) => (dispatch) => {
     const unsortedFolder = {
       id: 0,
       name: 'unsorted',
+      files: [],
       datasets: [],
     }
 
@@ -51,8 +52,13 @@ export const load_directories = (id) => (dispatch) => {
 
     // take all of the folders and put them in a lookup
     const folderLookup = folders.reduce((folderLookup, folder) => {
-      folderLookup[folder.id] = { id: folder.id, name: folder.name, folder, datasets: [] }
-      folder.datasets = []
+      folderLookup[folder.id] = {
+        id: folder.id,
+        name: folder.name,
+        folder,
+        files: [],
+        datasets: [],
+      }
       return folderLookup
     }, {
       unsorted: unsortedFolder
@@ -68,17 +74,21 @@ export const load_directories = (id) => (dispatch) => {
         file.name = (file.opt || {}).token || file.url
       }
       const name = (file.name || 'unsorted').split('.')[0]
-      const dataset = get_dataset(name, folderLookup[file.folder_id] || unsortedFolder, file.date || file.created_at)
+      const folder = folderLookup[file.folder_id] || unsortedFolder
+      const dataset = get_dataset(name, folder, file.date || file.created_at)
       if (file.url.match(file.name)) file.persisted = true
       dataset.input.push(file)
+      folder.files.push(file)
       return datasetLookup
     }, datasetLookup)
 
     // go over the generated files and add addl datasets (if the files were deleted)
     generatedFiles.map(file => {
       const pair = file.name.split('.')[0].split('-')
-      const dataset = get_dataset(pair[0], folderLookup[file.folder_id] || unsortedFolder, file.date || file.created_at)
+      const folder = folderLookup[file.folder_id] || unsortedFolder
+      const dataset = get_dataset(pair[0], folder, file.date || file.created_at)
       dataset.output.push(file)
+      folder.files.push(file)
       file.epoch = file.epoch || pair[1]
     })
 
@@ -145,7 +155,6 @@ export const load_directories = (id) => (dispatch) => {
         datasetLookup,
         folders, files,
         checkpoints,
-        builtDatasets,
         output,
       },
     })