StyleGAN2(-ada) is known for generating photorealistic images of e.g. portraits [Figure 1] and is now widely used in research, education and entertainment. Deployment of these models can be either by providing the files directly, dedicated model hosting services or in-browser. The latter has so far been the most challenging, yet a good trade-off between usability and cost. Since my last try a couple of months ago, I’ve finally seen new technologies to properly achieve this. So in this post I will explain the steps required to run StyleGAN2-ada models in your browser using onnxruntime with the current knowledge and technologies (October 2021). Tested on Windows 10. This approach may work for StyleGAN3 in the future as NVLabs stated on their StyleGAN3 git: “This repository is an updated version of stylegan2-ada-pytorch”. Unfortunately some of the ops used in StyleGAN3 are not supported in ONNX yet (affine_grid_generator). In short, this post is on how you can run StyleGAN2 in a browser.
You can go to https://github.com/Goblaski/stylegan2-web-convert to download the relevant files.
- StyleGAN2(-ada) PyTorch edition
- PyTorch 1.9.1 (this differs from the StyleGAN2-ada recommendations)
- A StyleGAN2 model (either torch or tensorflow based)
- For tensorflow based models you need to convert them with the legacy converter
It is important that you have your models in SyleGAN2-ada PyTorch format. Either you have trained your model previously in StyleGAN2(-ada) PyTorch or you can convert the tensorflow versions using the legacy converter. This should allow you to use any official StyleGAN2 version models. Example code can be found below.
python legacy.py \ --source=https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-cat-config-f.pkl \ --dest=stylegan2-cat-config-f.pkl
PyTorch model to ONNX
For the next step we need to convert your model to the ONNX format. I’ve written a snippet here to load the model, assign z and c as input variables and have Y as the output variable. It uses opset 10 of ONNX. The code will probably give you some warnings, but in the end you will end up with a file with the onnx extension. This configuration was the only configuration that worked properly for me. Alternatively, on github you can find a converter script.
import dnnlib import numpy as np import torch import legacy import functools # load model on cpu device = torch.device('cpu') with dnnlib.util.open_url(r'path\to\your\network.pkl') as f: GG = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore # enforce 32 bit floats (might be redundant in later version) GG.forward = functools.partial(GG.forward, force_fp32=True) # set input and output names in_names = [ "z" ] + [ "c"] out_names = [ "Y" ] #use a dummy input and label to determine graph dummy_input = torch.from_numpy(np.random.RandomState(0).randn(1, GG.z_dim)).to(device) label = torch.zeros([1, GG.c_dim], device=device) # export as onnx file torch.onnx.export(model=GG, args=(dummy_input,label), f="your_onnx_model.onnx", input_names=in_names, output_names=out_names, verbose=True, opset_version=10, export_params=True, do_constant_folding=False, use_external_data_format=False, operator_export_type=torch.onnx.OperatorExportTypes.ONNX)
Minimalist html + js
Opening the index.html file on a remote server should result in the images. An example [Figure 2] is given below as used in the Abstract Art Generator V2.
Pros & Cons
There are some discussion points to mention using this approach.
- Cheap solution requiring no additional hardware on your side
- Very broad support of devices and platforms
- The ONNX format allows you to use this in other deployments as well
- Fairly lightweight (no GPU required)
- Allows for style-transfer (See the https://github.com/Goblaski/stylegan2-web-convert page for that).
- Slow on CPU. I can’t export the model with the GPU ops at the moment and no GPU activity is used during generation
- Very old devices with low specs might not work
- High bandwidth usage. Each model is tens to hundreds of megabytes depending on the resolution used
- This might be an issue for some people: your models can easily be downloaded and used by others