Skip to main content

Generate Realistic Images using StyleGAN3

Introduction

With StyleGAN3 we can generate realistic looking images or videos. It can generate not only human faces, but also animals, cars and landscapes.

StyleGAN is based on Generative Adversarial Networks (GANs), which includes a generator and discriminator network that has been trained to differentiate images generated by generator from real images. However, during the training the generator tries to fool the discriminator, which results in the generatation of realistic looking images.

Prerequisite

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

Generate an image using pre-trained AFHQv2 model

python gen_images.py --outdir=out --trunc=1 --seeds=2 --network=stylegan3-r-afhqv2-512x512.pkl

Viewing the output image

Building the docker image

Instead of building from scratch we add the following lines to the existing docker file in the repo. Copy and paste the following lines of code to the container and set your variables.

COPY . /scratch

WORKDIR /scratch

ENV HOME /scratch

This step is done in the docker run command but since we’re running the command on a remote machine mounting local directory is not possible, so we add this step to the Dockerfile.

To build the docker container, run the docker build command

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

Replace the following:

  • &lt;hub-user>: with your docker hub username, If you don’t have a docker hub account Follow these instructions to create docker account, and use the username of the account you created

  • &lt;repo-name>: This is the name of the container, you can name it anything you want

  • &lt;tag>: This is not required but you can use the latest tag

After you have build the container, the next step is to test it locally and then push it docker hub.

You'll need to push this repository to the registry designated by its name or tag.

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

After the repo image has been pushed to docker hub, we can now run the container on Bacalhau.

Running the container on Bacalhau

Now we're ready to submit a Bacalhau job. The code below runs a job, downloads the results and prints the stdout.

The command below is similar to what we have run locally but we changed the output directory to the outputs folder so that the results are saved to IPFS.

We use the --gpu flag to denote the no of GPU we are going to use

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

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.

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

Parameters for gen_video.py animation length and seed keyframes:

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.

Get your Bacalhau job information

Command to get your 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
4f758052-0543-40b5-bd86-6ab41e77389a
bacalhau list --id-filter ${JOB_ID} --wide
 CREATED   ID        JOB                      STATE      VERIFIED  PUBLISHED               
 17:33:46  4f758052  Docker jsacex/stable...  Completed   /ipfs/QmcQEQPg934Pow... 

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

To find out more information about your job, run the following command:

bacalhau describe ${JOB_ID}

Since there is no error, we won't see any error instead we see the state of our job to be complete, that means we can download the results! Next is to create a temporary directory to save our results

To download the results of your job, run the following command:

rm -rf results && mkdir -p results
bacalhau get $JOB_ID --output-dir results
17:38:25.343 | INF bacalhau/get.go:67 > Fetching results of job '4f758052-0543-40b5-bd86-6ab41e77389a'...
2022/09/29 17:38:25 failed to sufficiently increase receive buffer size (was: 208 kiB, wanted: 2048 kiB, got: 416 kiB). See https://github.com/lucas-clemente/quic-go/wiki/UDP-Receive-Buffer-Size for details.
17:38:35.851 | INF ipfs/downloader.go:115 > Found 1 result shards, downloading to temporary folder.
17:38:37.1 | INF ipfs/downloader.go:195 > Combining shard from output volume 'outputs' to final location: '/content/results'

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

ls results/
shards  stderr  stdout  volumes