1 files changed, 17 insertions, 15 deletions

diff --git a/public/bundle.js b/public/bundle.js
index f2e7ef2..3c6d78d 100644
--- a/public/bundle.js
+++ b/public/bundle.js
@@ -9656,34 +9656,29 @@ var live_task = exports.live_task = function live_task(sequence, checkpoint, opt
   };
 };
 
-var augment_task = exports.augment_task = function augment_task(opt) {
+var augment_task = exports.augment_task = function augment_task(dataset, opt) {
   return function (dispatch) {
     var task = {
       module: _pix2pixhd2.default.name,
       activity: 'augment',
-      dataset: sequence,
-      checkpoint: checkpoint,
-      opt: _extends({}, opt, {
-        poll_delay: 0.01
-      })
+      dataset: dataset,
+      opt: _extends({}, opt)
     };
     console.log(task);
-    console.log('add live task');
+    console.log('add augment task');
     return _actions2.default.queue.add_task(task);
   };
 };
 
-var clear_recursive_task = exports.clear_recursive_task = function clear_recursive_task(opt) {
+var clear_recursive_task = exports.clear_recursive_task = function clear_recursive_task(dataset) {
   return function (dispatch) {
     var task = {
       module: _pix2pixhd2.default.name,
       activity: 'clear_recursive',
-      dataset: sequence,
-      checkpoint: checkpoint,
-      opt: _extends({}, opt)
+      dataset: dataset
     };
     console.log(task);
-    console.log('add live task');
+    console.log('add clear recursive task');
     return _actions2.default.queue.add_task(task);
   };
 };
@@ -10933,6 +10928,7 @@ var Pix2PixHDTrain = function (_Component) {
     _this.state = {
       checkpoint_name: '',
       epoch: 'latest',
+      augment_name: '',
       augment_take: 100,
       augment_make: 20
     };
@@ -11076,11 +11072,17 @@ var Pix2PixHDTrain = function (_Component) {
               min: '1',
               max: '1000'
             }),
+            (0, _preact.h)(_common.TextInput, {
+              name: 'augment_name',
+              title: 'Tag this epoch',
+              value: this.state.augment_name,
+              onChange: this.handleChange
+            }),
             (0, _preact.h)(_common.Button, {
               title: 'Augment dataset',
               value: 'Augment',
               onClick: function onClick() {
-                return remote.augment_task(dataset, pix2pixhd.folder_id, 1);
+                return remote.augment_task(_this2.state.checkpoint_name, _this2.state);
               }
             })
           ),
@@ -11091,7 +11093,7 @@ var Pix2PixHDTrain = function (_Component) {
               title: 'Train one epoch',
               value: 'Train',
               onClick: function onClick() {
-                return remote.train_task(dataset, pix2pixhd.folder_id, 1);
+                return remote.train_task(_this2.state.checkpoint_name, pix2pixhd.folder_id, 1);
               }
             })
           ),
@@ -11102,7 +11104,7 @@ var Pix2PixHDTrain = function (_Component) {
               title: 'Delete recursive frames',
               value: 'Clear',
               onClick: function onClick() {
-                return remote.clear_recursive_task(dataset, pix2pixhd.folder_id, 1);
+                return remote.clear_recursive_task(_this2.state.checkpoint_name);
               }
             })
           )