2 files changed, 16 insertions, 15 deletions
diff --git a/app/client/modules/pix2pixhd/pix2pixhd.tasks.js b/app/client/modules/pix2pixhd/pix2pixhd.tasks.js
index 92c0ff4..bd51f2b 100644
--- a/app/client/modules/pix2pixhd/pix2pixhd.tasks.js
+++ b/app/client/modules/pix2pixhd/pix2pixhd.tasks.js
@@ -53,34 +53,28 @@ export const live_task = (sequence, checkpoint, opt) => dispatch => {
   return actions.queue.add_task(task)
 }
 
-export const augment_task = (opt) => dispatch => {
+export const augment_task = (dataset, opt) => dispatch => {
   const task = {
     module: module.name,
     activity: 'augment',
-    dataset: sequence,
-    checkpoint,
+    dataset,
     opt: {
       ...opt,
-      poll_delay: 0.01,
     }
   }
   console.log(task)
-  console.log('add live task')
+  console.log('add augment task')
   return actions.queue.add_task(task)
 }
 
-export const clear_recursive_task = (opt) => dispatch => {
+export const clear_recursive_task = (dataset) => dispatch => {
   const task = {
     module: module.name,
     activity: 'clear_recursive',
-    dataset: sequence,
-    checkpoint,
-    opt: {
-      ...opt,
-    }
+    dataset,
   }
   console.log(task)
-  console.log('add live task')
+  console.log('add clear recursive task')
   return actions.queue.add_task(task)
 }
 
diff --git a/app/client/modules/pix2pixhd/views/pix2pixhd.train.js b/app/client/modules/pix2pixhd/views/pix2pixhd.train.js
index cd6507f..df3a1f2 100644
--- a/app/client/modules/pix2pixhd/views/pix2pixhd.train.js
+++ b/app/client/modules/pix2pixhd/views/pix2pixhd.train.js
@@ -28,6 +28,7 @@ class Pix2PixHDTrain extends Component {
     this.state = {
       checkpoint_name: '',
       epoch: 'latest',
+      augment_name: '',
       augment_take: 100,
       augment_make: 20,
     }
@@ -131,10 +132,16 @@ class Pix2PixHDTrain extends Component {
               min="1"
               max="1000"
             />
+            <TextInput
+              name="augment_name"
+              title="Tag this epoch"
+              value={this.state.augment_name}
+              onChange={this.handleChange}
+            />
             <Button
               title="Augment dataset"
               value="Augment"
-              onClick={() => remote.augment_task(dataset, pix2pixhd.folder_id, 1)}
+              onClick={() => remote.augment_task(this.state.checkpoint_name, this.state)}
             />
           </Group>
 
@@ -142,7 +149,7 @@ class Pix2PixHDTrain extends Component {
             <Button
               title="Train one epoch"
               value="Train"
-              onClick={() => remote.train_task(dataset, pix2pixhd.folder_id, 1)}
+              onClick={() => remote.train_task(this.state.checkpoint_name, pix2pixhd.folder_id, 1)}
             />
           </Group>
 
@@ -150,7 +157,7 @@ class Pix2PixHDTrain extends Component {
             <Button
               title="Delete recursive frames"
               value="Clear"
-              onClick={() => remote.clear_recursive_task(dataset, pix2pixhd.folder_id, 1)}
+              onClick={() => remote.clear_recursive_task(this.state.checkpoint_name)}
             />
           </Group>
         </div>