1 files changed, 31 insertions, 2 deletions

diff --git a/app/client/modules/pix2pixhd/pix2pixhd.actions.js b/app/client/modules/pix2pixhd/pix2pixhd.actions.js
index 8e481d3..c1cd2b1 100644
--- a/app/client/modules/pix2pixhd/pix2pixhd.actions.js
+++ b/app/client/modules/pix2pixhd/pix2pixhd.actions.js
@@ -21,7 +21,11 @@ export const load_directories = (id) => (dispatch) => {
     // actions.socket.disk_usage({ module, dir: 'datasets' }),
   ], (percent, i, n) => {
     console.log('pix2pixhd load progress', i, n)
-    dispatch({ type: types.app.load_progress, progress: { i, n }})
+    dispatch({
+      type: types.app.load_progress,
+      progress: { i, n },
+      data: { module: 'pix2pixhd' },
+    })
   }).then(res => {
     const [datasetApiReport, sequences, datasets, checkpoints] = res //, datasets, results, output, datasetUsage, lossReport] = res
     const {
@@ -157,7 +161,11 @@ export const load_results = (id) => (dispatch) => {
     actions.socket.list_directory({ module, dir: 'renders' }),
   ], (percent, i, n) => {
     console.log('pix2pixhd load progress', i, n)
-    dispatch({ type: types.app.load_progress, progress: { i, n }})
+    dispatch({
+      type: types.app.load_progress,
+      progress: { i, n },
+      data: { module: 'pix2pixhd' },
+    })
   }).then(res => {
     const [folders, files, results, renders] = res //, datasets, results, output, datasetUsage, lossReport] = res
     console.log(files, results, renders)
@@ -171,4 +179,25 @@ export const load_results = (id) => (dispatch) => {
       }
     })
   })
+}
+
+const G_NET_REGEXP = new RegExp('_net_G.pth$')
+
+export const list_epochs = (checkpoint_name) => (dispatch) => {
+  const module = pix2pixhdModule.name
+  actions.socket.list_directory({ module, dir: 'checkpoints/' + checkpoint_name }).then(files => {
+    // console.log(files)
+    const epochs = files.map(f => {
+      if (!f.name.match(G_NET_REGEXP)) return null
+      return f.name.replace(G_NET_REGEXP, '')
+    }).filter(f => !!f)
+    // console.log(epochs)
+    dispatch({
+      type: types.dataset.list_epochs,
+      data: {
+        epochs,
+        module
+      },
+    })
+  })
 }
\ No newline at end of file