Skip to content

vvmnnnkv/nodejs-style-transfer

Repository files navigation

Image Style Transfer with Node.JS

Demo of image style transfer app based on Node.js and Express that uses fast neural style transfer model created and traced with Pytorch.

Usage

Open the site on your mobile or desktop, shoot or upload something, stylize :-)

Image Style Transfer with Node.js

Live demo: https://nodejs-style-transfer.herokuapp.com

Why?

This app is inspired by Lessons 6 & 9 of Facebook Pytorch Scholarship. It shows how easy it is possible to deploy high-performance ML models these days combining such popular platforms as Node.js and NPM with newly released Torch Script feature of Pytorch.

Apart from that, Node.js non-blocking I/O seems to be good fit for running NN inference in Worker Pool.

The heavy lifting of image processing is done by libtorch C++ library, which is wrapped in Node.js module libtorchjs created specially for this app to expose some bits of libtorch in JS.

Running

Currently supported platforms are Linux and Windows. libtorch does have Mac variant, but it's not included in libtorchjs yet. GPU is also not supported yet.

Linux or Windows

Checkout repo, install npm libraries, start app:

$ npm i
$ npm run start

Open http://localhost:3000/

Docker

Checkout repo, build image and run:

$ docker build --tag=nodejs-style-transfer .
$ docker run -P nodejs-style-transfer

Open http://localhost:3000/

How to Add Style

This app has just one style included.
The easiest way to create more models is to install fast_neural_style from Pytorch examples and add following code in the end of stylyze function in fast_neural_style/neural_style/neural_style.py:

traced_script_module = torch.jit.trace(style_model, content_image)
traced_script_module.save("style_model.pt")

And after that, execute stylization script on any image and this will write out style_model.pt. Note that traced model will keep the size of your image and will be able to work only with same dimensions. For this demo app, image size is fairly small (224x224) so it's fast even on CPU.

See libtorch tutorial for more info about model tracing.

Acknowledgements

Disclamer

Note that it's just a toy project and libtorchjs is too immature to use in production.