StyleGAN2 in PyTorch, with wavelets and a disclaimer

What it does

This is a ground-up PyTorch rebuild of NVIDIA’s StyleGAN2, meant for training high-resolution image generators without the TensorFlow stack. It ships with tools to convert official checkpoints, project photos back into latent space, and tinker with semantic directions via closed-form factorization. The repo also experimentally supports SWAGAN, a wavelet-based variant the author says trains about twice as fast at 256px.

The interesting bit

The author is refreshingly blunt: they warn they may have missed official details and confess they don’t know why their FID scores diverge. That honesty makes the repo feel more like a working lab notebook than a polished product drop. The closed-form factorization scripts are a nice bonus—they turn weight-matrix eigenvectors into unsupervised knobs for tweaking facial attributes.

Key highlights

• Converts official StyleGAN2 TensorFlow checkpoints into PyTorch format.
• Includes closed_form_factorization.py and apply_factor.py for probing latent semantics without labels.
• Experimental SWAGAN architecture for faster wavelet-driven training.
• Supports distributed training and Weights & Biases logging.
• Pretrained 256px FFHQ checkpoint available (author-reported FID ~4.5, with noted discrepancies from official scores).

Caveats

• The author explicitly warns that some implementation details may differ from the official version; they advise using it “with care.”
• FID scores differ from the official TensorFlow release, and the author is unsure whether preprocessing, resolution, or the training loop is to blame.

Verdict

Worth cloning if you want a PyTorch-native StyleGAN2 for research or latent-space hacking. If you need guaranteed bit-for-bit fidelity or a battle-tested production pipeline, stick with NVIDIA’s official release.