Real-time Fast Style Transfer model with Tensorflow.js and Next.js in the Browser

I wanted to share a project that I have been working on using Next.js and React. My challenge was to create an interactive component using React/Next.JS that anyone could try from a browser to play with the Fast Style Transfer model in real-time, via a web camera all within the browser.

Let’s just say that this was not an easy task. But I did it for a very specific reason. I needed to learn Next.JS and fast.

What is Next.JS?

Next.JS is a flavor of REACT, and if you haven’t heard of REACT before, REACT was developed by Facebook for Single Page Applications with a high element of interactivity capable of running on any device with access to browser, this includes iPhone, Android, desktop applications for Windows and Mac, and many more.

One of the trademarks of React applications is that all the rendering work is done by the client. This is good, especially for native applications running on a mobile phone or even a browser, but it is not good for websites that rely on SEO to be discovered by Google.

The defining characteristic of a REACT Single Page Application is that a browser receives only a skelleton of the HTML and the remainder is built by the Browser using a Javascript intepreter. This is all done in the client-side. But the search bots typically do not have the ability of running Javascript code before rendering an HTML page.

This is mainly why Next.JS and other REACT flavours like Gatsby came about.

Server Side Rendering

Unlike traditional REACT applications, Next.JS can do all the pre-rendering work of building the HTML page, in the server side, using a Javascript interpreter(NodeJS). The benefits of doing this is that not only we solve the SEO problem, as the search bots are not required to run Javascript and we also help reduce performance requirements on the client, allowing thinner devices to render the site quickly.

The nice thing is that with Next.JS we can strategically decide whether to use server-side rendering and/or client-side rendering. And there is one more option. Let’s go into that next.

Static Site Generation

Instead of rendering the HTML server side, each time we get a website visitor, Next.js(and so can Gatsby), we can also pre-render all pages in the site, at build time, using static site generation. This means that we can just copy the site to a CDN and serve it from anywhere without even needing a backend server running. Isn’t that awesome?

Hybrid Next.JS App

The power of Next.jS is that it allows us to create a hybrid of Server Side Rendering and Static Site Generation. Typically, pages that are very dynamic and can change often probably should be server-side rendered and pages that change very little should be generated at build time.

Hydration

It is not necessary to build 100% of the HTML before serving the page to the user, to get the benefits of server side rendering or static site generation.

By using a process known as hydration, the browser can add more content to the HTML once it has downloaded the pre-rendered HTML. This is typically useful for example to get content specific to the user or for example when it is important that some of the content in the page is more recent.

What is Tensorflow.JS?

Tensorflow.js is a client-side framework designed to enable training and inference with a machine-learning directly from the browser, using Javascript.

With Tensorflow.js we can:

  • Run pre-trained models, originally developed for Tensorflow directly from the browser. Tensorflow.js provides a command-line tool to convert almost any existing pre-trained model.
  • Train new or pre-existing models with Transfer Learning using our own data with a Javascript API.

Tensorflow.js can also run from a Node.js backend, but for me the focus will be on running Tensorflow.js in the client-side.

How to convert a Model from Tensorflow to Tensorflow.jS

Before you can use a model with Tensorflow.js, you need to convert it to another format so it can be served to a web browser and usable with Tensorflow.js. Thankfully I have already done this beforehand, so if you want to find out how I converted the Fast Style Transfer model to Tensorflow.js you just need to read my tutorial:

Building the Tensorflow.JS App

Now let’s go to the fun part. I will try to describe all the steps I went through to create a Next.JS app using Typescript. If you don’t know Typescript, Typescript is a strongly typed version of Javascript. It is a pain to learn but far easier to resolve errors as you get information on what is wrong with your javsacript code, most times without even running it.

How to create a Starter Next.JS app using Typescript.

To create a Next.JS app using typescript, we can simply run:

npx [email protected] --ts

You should get a similar output if all goes well:

Need to install the following packages:
  [email protected]
Ok to proceed? (y) y
✔ What is your project named? … fast-style-transfer-nextjs
Creating a new Next.js app in /home/codemental/projects/fast-style-transfer-nextjs.

Using npm.

Installing dependencies:
- react
- react-dom
- next


added 16 packages, and audited 17 packages in 7s

2 packages are looking for funding
  run `npm fund` for details

found 0 vulnerabilities

Installing devDependencies:
- eslint
- eslint-config-next
- typescript
- @types/react
- @types/node


added 208 packages, and audited 225 packages in 11s

63 packages are looking for funding
  run `npm fund` for details

found 0 vulnerabilities

Initialized a git repository.

Success! Created fast-style-transfer-nextjs at /home/codemental/projects/fast-style-transfer-nextjs
Inside that directory, you can run several commands:

  npm run dev
    Starts the development server.

  npm run build
    Builds the app for production.

  npm start
    Runs the built app in production mode.

We suggest that you begin by typing:

  cd fast-style-transfer-nextjs
  npm run dev

Starting the Next.js App locally for Development

To develop our widget we need to run locally our Next.js app.
In order to run locally our Next.js app all we need to do is to open a terminal window and under the same directory where the app was created execute:

  cd fast-style-transfer-nextjs
  npm run dev

> dev
> next dev

ready - started server on 0.0.0.0:3000, url: http://localhost:3000
event - compiled client and server successfully in 392 ms (158 modules)

We are good to go! Open localhost:3000 on your browser just to confirm that you can open it and that you see a Hello World page.

Understanding the structure of the Fast Style Transfer Next.js App

Even though I have already touched on significant concepts about Next.JS and REACT, I need to remind myself that I am not writing a tutorial on Next.js.

Instead, I will briefly describe the major building blocks of the “Fast Style Transfer” app and then we can discuss each one, and focus on Tensorflow.js.

REACT Component Structure of the Fast Style Transfer Next.js App

In this diagram, you can see the main REACT components that I have created to build the user interface. The core concept of REACT is to build the UI as a set of reusable blocks(REACT Components.

FastStyleTransferModel

The FastStyleTransferModel is a REACT component with the main job of loading the Tensoflow.js model from the network. This model is quite heavy(70mb+), so it takes a significant amount of time to download it, depending on the network bandwidth available.

Therefore this component, while it downloads the model from the network, it displays a “Loading Model” message inside a modal. When it is finished with the download, only then does it download the remainder of the REACT components that display the full user interface.

useEffect() React Hook

We only want to load the model when the page loads for the first time, since it is a very heavy download. For that reason we use the React Hook useEffect(). If we use pass an empty array to the second parameter of useEffect(), the function we pass in the first parameter only gets called once per page load.

    useEffect(() => {
        console.log("Loading model!");
        loadModel().then((loadedModel: tf.GraphModel) => {
            console.log("Finished loading model");
            let model = loadedModel;
            setModel(loadedModel);
        });
    }, [])

useState() React Hook

While the model is loading in the UI we show a modal saying that the model is being loaded. Once the model is loaded we show the UI that the user can use to experiment with the Style Transfer Model. This is all accomplished by using useState().

const [model, setModel] = useState(null as tf.GraphModel | null)


When the model has finished loading we change the REACT state using setModel. Once the model is added to the REACT state, we change the UI as follows:

 return (
        <>
            {!model &&
                <LoadingModel />
            }

            {model &&
                [ children(doStyleTransferCallback) ]
            }
        </>
    )

Demo