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
To run StyleGAN3 you need to match these requirements
The Bacalhau client - Installation instructions
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:
<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>: This is the name of the container, you can name it anything you want<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
[92;100m CREATED [0m[92;100m ID [0m[92;100m JOB [0m[92;100m STATE [0m[92;100m VERIFIED [0m[92;100m PUBLISHED [0m
[97;40m 17:33:46 [0m[97;40m 4f758052 [0m[97;40m Docker jsacex/stable... [0m[97;40m Completed [0m[97;40m [0m[97;40m /ipfs/QmcQEQPg934Pow... [0m
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
[90m17:38:25.343 |[0m [32mINF[0m [1mbacalhau/get.go:67[0m[36m >[0m 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.
[90m17:38:35.851 |[0m [32mINF[0m [1mipfs/downloader.go:115[0m[36m >[0m Found 1 result shards, downloading to temporary folder.
[90m17:38:37.1 |[0m [32mINF[0m [1mipfs/downloader.go:195[0m[36m >[0m 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