neuralstyle.js

'use strict';

let firebase = require('firebase');
let fs = require('fs');
let spawn = require('child_process').spawn;
let throttle = require('lodash.throttle');
let request = require('superagent');
let AWS = require('aws-sdk');

let ref = new Firebase('https://neural.firebaseio.com/');
let lock = false;
let queue = [];

fs.watch('out', throttle((evt, file) => {
  let name = Math.random().toString(36).substring(7) + '.png';

  let s3obj = new AWS.S3({params: {Bucket: 'neuralstyle', Key: name, ContentType: 'image/png'}});
  s3obj.upload({Body: fs.createReadStream('out/' + file)})
    .send(err => {
      if (err)
        return;

      console.log('http://neuralstyle.s3-website-eu-west-1.amazonaws.com/' + name);
      ref.child('output').push({
        url: 'http://neuralstyle.s3-website-eu-west-1.amazonaws.com/' + name,
        iter: file.split('.')[0].substr(4)
      })
    });

  console.log(evt, file);
}, 2000, {leading: false}));

ref.child('input').on('child_added', snap => {
  let data = snap.val();

  let styles = Promise.all(data.styleUrls.map((image, i) => new Promise(accept => {
    let path = 'in/style' + i;
    request.get(image.url)
      .pipe(fs.createWriteStream(path))
      .on('close', () => accept(path));
  })));

  let content = new Promise(accept => request.get(data.contentUrl)
    .pipe(fs.createWriteStream('in/content'))
    .on('close', () => accept('in/content')));

  content.then(console.log.bind(console))
  styles.then(console.log.bind(console))

  Promise.all([styles, content]).then(images => {
    console.log(images);
    data.style = images[0].join(',');
    data.content = images[1];
    data.weights = data.styleUrls.map(image => image.weight).join(',');
    job(data);
    ref.child('input').child(snap.key()).remove();
  });
});

let job = data => {
  if (lock) {
    console.log('Got job while working. Enqueue!');
    queue.unshift(data);
    return;
  }
  lock = true;

  console.log('Starting job!');

  data.size = data.size || 512;
  data.smooth = data.smooth || 0.001;
  data.iterations = data.iterations || 100;
  data.numPics = data.numPics || 5;
  data.random = data.random ? 'random' : 'image';
  data.scale = data.scale || 1;
  data.contentWeight = data.contentWeight || 5;
  data.styleWeight = data.styleWeight || 100;

  let neural = spawn('th', [
    'neural_style.lua',
    '-gpu', '0',
    '-backend', 'cudnn', '-cudnn_autotune',
    '-output_image', 'out/out.png',
    '-style_image', data.style,
    '-content_image', data.content,
    '-style_blend_weights', data.weights,
    '-image_size', data.size,
    '-num_iterations', data.iterations * data.numPics,
    '-save_iter', data.iterations,
    '-init', data.random,
    '-style_scale', data.scale,
    '-content_weight', data.contentWeight,
    '-style_weight', data.styleWeight,
  ], {cwd: process.cwd()})

  neural.stderr.on('data', data => console.log(data.toString()));

  neural.on('close', code => {
    console.log(code);
    lock = false;
    let nextData = queue.pop();
    if (nextData) {
      console.log('Picking up job from queue!');
      return job(nextData);
    }
  });
}