1 files changed, 27 insertions, 3 deletions
diff --git a/app/relay/modules/pix2pixhd.js b/app/relay/modules/pix2pixhd.js
index a01e321..43f839e 100644
--- a/app/relay/modules/pix2pixhd.js
+++ b/app/relay/modules/pix2pixhd.js
@@ -4,7 +4,7 @@ import fs from 'fs'
 const name = 'pix2pixhd'
 const cwd = process.env.PIX2PIXHD_CWD || path.join(process.env.HOME, 'code/' + name + '/')
 const env = {
-  LD_LIBRARY_PATH: '/usr/local/cuda/lib64:' + process.env.HOME + '/Downloads/TensorRT-4.0.0.3/lib',
+  LD_LIBRARY_PATH: '/usr/local/cuda/lib64:' + process.env.TENSORRT_LIB_PATH,
 }
 
 const fetch = {
@@ -49,8 +49,9 @@ const train = {
     const datasets_path = path.join(cwd, 'datasets', dataset)
     const checkpoints_path = path.join(cwd, 'checkpoints', dataset)
     const iter_txt = path.join(checkpoints_path, 'iter.txt')
+    const checkpoint_path = path.join(checkpoints_path, 'latest_net_G.pth')
     console.log(dataset, iter_txt)
-    if (fs.existsSync(iter_txt)) {
+    if (fs.existsSync(iter_txt) && fs.existsSync(checkpoint_path)) {
       const iter = fs.readFileSync(iter_txt).toString().split('\n');
       console.log(iter)
       epoch = iter[0] || 0
@@ -234,7 +235,7 @@ const render_recursive = {
   }
 }
 const splice = {
-  type: 'perl_gpu',
+  type: 'perl',
   script: 'splice.pl',
   params: (task) => {
     console.log(task.opt.selection)
@@ -250,6 +251,28 @@ const splice = {
   },
   after: 'train'
 }
+const uprez = {
+  type: 'pytorch',
+  script: 'uprez.py',
+  params: task => {
+    console.log(task)
+    /*
+     set folder based on whether the file is a result or recursive mp4
+    */
+    let folder
+    if (task.opt.is_result) {
+      folder = './results/' + task.dataset
+    } else if (task.opt.is_recursive) {
+      folder = './recursive/' + task.dataset + '/' + task.opt.name + '/'
+    }
+    return [
+      '--folder', folder,
+      '--folder_id', task.opt.folder_id,
+      '--dataset', task.dataset,
+      '--out_fn', task.dataset + "_" + Date.now(),
+    ]
+  }
+}
 
 function pad(num, size) {
   if (isNaN(parseInt(num))) return num;
@@ -270,5 +293,6 @@ export default {
     render,
     render_recursive,
     splice,
+    uprez,
   }
 }