Generate Realistic Images using StyleGAN3 and Bacalhau
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 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. With StyleGAN3 we can generate realistic looking images or videos. It can generate not only human faces, but also animals, cars and landscapes.
TD;LR
Generative images with Bacalhau
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.
%%bash
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
Viewing the output image
Containerize Script with Docker
To build your own docker container, create a Dockerfile
, which contains instructions to build your image.
COPY . /scratch
WORKDIR /scratch
ENV HOME /scratch
See more information on how to containerize your script/apphere
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;
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
repo-name with the name of the container, you can name it anything you want
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
To submit a job, run the following Bacalhau command:
%%bash --out 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
Structure of the command
Let's look closely at the command above:
bacalhau docker run
: call to bacalhau--gpu 1
: No of GPUsjsacex/stylegan3
: the name and the tag of the docker image we are using../outputs
: path to outputpython gen_images.py
: execute script--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.
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.
%%bash --out 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
Structure of the command
Let's look closely at the command above:
bacalhau docker run
: call to bacalhau--gpu 1
: No of GPUsjsacex/stylegan3
: the name and the tag of the docker image we are using../outputs
: path to outputpython gen_images.py
: execute script--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 'numkeyframes*w_frames' frames. If --num-keyframes is not specified, the number of seeds given with --seeds_ must be divisible by grid size WH (--grid). In this case the output video length will be '# seeds/(wh)*w_frames' frames.
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.
%env JOB_ID={job_id}
Checking the State of your Jobs
- Job status: You can check the status of the job using
bacalhau list
.
%%bash
bacalhau list --id-filter ${JOB_ID} --wide
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 describe
.
%%bash
bacalhau describe ${JOB_ID}
- Job download: You can download your job results directly by using
bacalhau 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.
%%bash
rm -rf results && mkdir -p results
bacalhau get $JOB_ID --output-dir results
After the download has finished you should see the following contents in results directory
Viewing your Job Output
To view the file, run the following command:
%%bash
ls results/outputs
import IPython.display as display
display.Image("results/outputs/seed0002.png")