Generate Realistic Images using StyleGAN3 and Bacalhau

Introduction

In this example tutorial, we will show you how to generate realistic images with StyleGAN3 and Bacalhau. StyleGAN is based on Generative Adversarial Networks (GANs), which include a generator and discriminator network that has been trained to differentiate images generated by the generator from real images. However, during the training, the generator tries to fool the discriminator, which results in the generation of realistic-looking images. With StyleGAN3 we can generate realistic-looking images or videos. It can generate not only human faces but also animals, cars, and landscapes.

TL;DR​

bacalhau docker run \
    --wait \
    --id-only \
    --gpu 1 \
    --timeout 3600 \
    --wait-timeout-secs 3600 \
    jsacex/stylegan3 \
    -- python gen_images.py --outdir=../outputs --trunc=1 --seeds=2 --network=stylegan3-r-afhqv2-512x512.pkl

Prerequisite​

To get started, you need to install the Bacalhau client, see more information here

Running StyleGAN3 locally​

To run StyleGAN3 locally, you'll need to clone the repo, install dependencies and download the model weights.

git clone https://github.com/NVlabs/stylegan3
cd stylegan3
conda env create -f environment.yml
conda activate stylegan3
wget https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl

Now you can generate an image using a pre-trained AFHQv2 model. Here is an example of the image we generated:

Containerize Script with Docker​

To build your own docker container, create a Dockerfile, which contains instructions to build your image.

FROM nvcr.io/nvidia/pytorch:21.08-py3

COPY . /scratch

WORKDIR /scratch

ENV HOME /scratch

See more information on how to containerize your script/app here

Build the container​

We will run docker build command to build the container:

docker build -t <hub-user>/<repo-name>:<tag> .

Before running the command replace:

1. hub-user with your docker hub username, If you don’t have a docker hub account follow these instructions to create docker account (https://docs.docker.com/docker-id/), and use the username of the account you created

2. repo-name with the name of the container, you can name it anything you want

3. tag this is not required but you can use the latest tag

In our case:

docker build -t jsacex/stylegan3 

Push the container​

Next, upload the image to the registry. This can be done by using the Docker hub username, repo name or tag.

docker push <hub-user>/<repo-name>:<tag>

In our case:

docker push jsacex/stylegan3 

Running a Bacalhau Job​

Structure of the command​

Some of the jobs presented in the Examples section may require more resources than are currently available on the demo network. Consider starting your own network or running less resource-intensive jobs on the demo network

To submit a job run the Bacalhau command with following structure:

1. export JOB_ID=$( ... ) exports the job ID as environment variable

2. bacalhau docker run: call to Bacalhau

3. The --gpu 1 flag is set to specify hardware requirements, a GPU is needed to run such a job

4. The --id-only flag is set to print only job id

5. jsacex/stylegan3: the name and the tag of the docker image we are using

6. python gen_images.py: execute the script with following parameters:

   1. --trunc=1 --seeds=2 --network=stylegan3-r-afhqv2-512x512.pkl: The animation length is either determined based on the --seeds value or explicitly specified using the --num-keyframes option. When num keyframes are specified with --num-keyframes, the output video length will be num_keyframes * w_frames frames.

   2. ../outputs: path to the output

export JOB_ID=$(bacalhau docker run \
    --wait \
    --id-only \
    --gpu 1 \
    --timeout 3600 \
    --wait-timeout-secs 3600 \
    jsacex/stylegan3 \
    -- python gen_images.py --outdir=../outputs --trunc=1 --seeds=2 --network=stylegan3-r-afhqv2-512x512.pkl)

Declarative job description​

The same job can be presented in the declarative format. In this case, the description will look like this:

name: StyleGAN3
type: batch
count: 1
tasks:
  - name: My main task
    Engine:
      type: docker
      params:
        Image: "jsacex/stylegan3" 
        Parameters:
          - python gen_images.py --outdir=../outputs --trunc=1 --seeds=2 --network=stylegan3-r-afhqv2-512x512.pkl
    Resources:
      GPU: "1"

The job description should be saved in .yaml format, e.g. stylegan3.yaml, and then run with the command:

bacalhau job run stylegan3.yaml

Render a latent vector interpolation video​

You can also run variations of this command to generate videos and other things. In the following command below, we will render a latent vector interpolation video. This will render a 4x2 grid of interpolations for seeds 0 through 31.

Structure of the command​

Let's look closely at the command below:

1. export JOB_ID=$( ... ) exports the job ID as environment variable

2. bacalhau docker run: call to bacalhau

3. The --gpu 1 flag is set to specify hardware requirements, a GPU is needed to run such a job

4. The --id-only flag is set to print only job id

5. jsacex/stylegan3 the name and the tag of the docker image we are using

6. python gen_images.py: execute the script with following parameters:

   1. --trunc=1 --seeds=2 --network=stylegan3-r-afhqv2-512x512.pkl: The animation length is either determined based on the --seeds value or explicitly specified using the --num-keyframes option. When num keyframes is specified with --num-keyframes, the output video length will be num_keyframes * w_frames frames. If --num-keyframes is not specified, the number of seeds given with --seeds must be divisible by grid size W*H (--grid). In this case, the output video length will be # seeds/(w*h)*w_frames frames.

   2. ../outputs: path to the output

export JOB_ID=$(bacalhau docker run \
    jsacex/stylegan3 \
    --gpu 1 \
    --timeout 3600 \
    --wait-timeout-secs 3600 \
    -- python gen_video.py --output=../outputs/lerp.mp4 --trunc=1 \
        --seeds=0-31 \
        --grid=4x2 \
        --network=stylegan3-r-afhqv2-512x512.pkl)

When a job is submitted, Bacalhau prints out the related job_id. We store that in an environment variable so that we can reuse it later on.

Checking the State of your Jobs​

Job status​

You can check the status of the job using bacalhau job list.

bacalhau job list --id-filter ${JOB_ID}

When it says Completed, that means the job is done, and we can get the results.

Job information​

You can find out more information about your job by using bacalhau job describe.

bacalhau job describe ${JOB_ID}

Job download​

You can download your job results directly by using bacalhau job get. Alternatively, you can choose to create a directory to store your results. In the command below, we created a directory and downloaded our job output to be stored in that directory.

rm -rf results && mkdir -p results
bacalhau job get $JOB_ID --output-dir results

After the download has finished you should see the following contents in results directory

Viewing your Job Output​

Now you can find the file in the results/outputs folder.

Support

If you have questions or need support or guidance, please reach out to the Bacalhau team via Slack (#general channel).