jax-stylegan3

name	jax-stylegan3
description	Use this skill when constructing, sampling, training, saving, loading, or documenting the JAX plus Equinox StyleGAN3 implementation from the stylegan3_jax repository.

JAX StyleGAN3

This skill documents the JAX StyleGAN3 implementation from src/jax_stylegan3. It vendors a snapshot of that package in jax_stylegan3/ and includes four verified examples.

• construct.py builds a tiny generator and discriminator.
• generate.py samples single, truncated, and batched images.
• train.py runs a minimal synthetic GAN training step.
• save_load.py saves, lists, and reloads a pretrained generator.

Run the examples from an environment with the project dependencies installed. The training example needs the project train extra because it imports optax.

uv sync --extra train --extra dev
uv run python ~/.claude/skills/jax-stylegan3/construct.py
uv run python ~/.claude/skills/jax-stylegan3/generate.py
uv run python ~/.claude/skills/jax-stylegan3/train.py
uv run python ~/.claude/skills/jax-stylegan3/save_load.py

Module Map

The generator API lives in jax_stylegan3.generator. The discriminator API lives in jax_stylegan3.discriminator. The generator config is jax_stylegan3.config.StyleGAN3Config. Training losses live in jax_stylegan3.stylegan2_loss. The moving average helper is jax_stylegan3.mapping.update_w_avg.

The full repository also has reference training loops under models/dsprites and models/mnist. Use those for dataset IO and checkpointing patterns.

Construct

Create a StyleGAN3Config, then pass it to Generator.init. The discriminator is initialized separately. For tiny batches or small smoke tests, disable minibatch standard deviation grouping through epilogue_kwargs.

cfg = StyleGAN3Config(
  z_dim=16,
  w_dim=16,
  c_dim=0,
  img_resolution=32,
  img_channels=1,
  num_layers=4,
  num_critical=2,
  channel_base=1024,
  channel_max=16,
  margin_size=4,
  num_fp16_res=0,
  smooth_generator=True,
)
key_g, key_d = jax.random.split(key)
generator = Generator.init(key_g, cfg)
discriminator = Discriminator.init(
  key_d,
  c_dim=cfg.c_dim,
  img_resolution=cfg.img_resolution,
  img_channels=cfg.img_channels,
  channel_base=1024,
  channel_max=16,
  epilogue_kwargs={
    "mbstd_group_size": None,
    "mbstd_num_channels": 1,
  },
)

See construct.py for a complete runnable version with output shape assertions.

Generate

The generator accepts batched latent vectors with shape (batch, z_dim). For unconditional models, pass None for the label input. sample_generator creates random latent vectors and calls the generator.

single = sample_generator(G=generator, key=key_single, batch_size=1)
truncated = sample_generator(
  G=generator,
  key=key_truncated,
  batch_size=1,
  truncation_psi=0.7,
)

For a per-key batch pattern, vmap a one-sample function.

def sample_one(sample_key: jax.Array) -> jnp.ndarray:
  return sample_generator(G=generator, key=sample_key, batch_size=1)[0]

batched = jax.vmap(sample_one)(keys)

See generate.py for a complete runnable version.

Train

JIT the top-level training step with eqx.filter_jit. Use eqx.filter_value_and_grad on the discriminator and generator modules. Filter parameters and gradients with eqx.is_inexact_array, update with optax, and apply updates with eqx.apply_updates.

The verified smoke test follows the same structure as the dSprites loop. It uses synthetic random images only to exercise the API.

loss_d, grads_d = eqx.filter_value_and_grad(d_loss_fn)(discriminator)
params_d = eqx.filter(discriminator, eqx.is_inexact_array)
grads_d = eqx.filter(grads_d, eqx.is_inexact_array)
updates_d, opt_state_d = optim_d.update(grads_d, opt_state_d, params_d)
discriminator = eqx.apply_updates(discriminator, updates_d)

Use jax.lax.stop_gradient on fake images for the discriminator step. Use r1_penalty_grad on an interval when you want R1 regularization. Update the mapping average with update_w_avg.

fake = generator(z, None)
fake_sg = jax.lax.stop_gradient(fake)
loss_real = loss_d_real(discriminator(real_img, None))
loss_fake = loss_d_gen(discriminator(fake_sg, None))

See train.py for the complete verified training step.

Save And Load

The pretrained model path in this skill means a generator trained and saved from this JAX codebase. There is no PyTorch to JAX weight converter in the repository.

save_load.py saves one generator per directory under MODEL_ROOT. The default MODEL_ROOT is Path("pretrained_models"), relative to the current working directory. Each saved model directory contains these files.

• weights.eqx
• metadata.json

The metadata stores name, description, saved_at, step, loss, num_parameters, and the full StyleGAN3Config as config. The config is the source of truth for reconstructing the skeleton before loading weights.

eqx.tree_serialise_leaves(weights_path, generator)
metadata = {
  "name": name,
  "description": description,
  "saved_at": datetime.now(timezone.utc).isoformat(),
  "step": int(step),
  "loss": float(loss),
  "num_parameters": count_array_parameters(generator),
  "config": asdict(cfg),
}

To load, rebuild the skeleton from the saved config, then deserialize the leaves.

metadata = read_metadata(model_dir)
cfg = StyleGAN3Config(**metadata["config"])
skeleton = Generator.init(key, cfg)
generator = eqx.tree_deserialise_leaves(model_dir / "weights.eqx", skeleton)

list_saved_models(model_root) returns metadata dictionaries for directories that contain metadata.json. save_load.py prints a short summary and checks exact output parity after loading.

Notes For Agents

Keep JAX code in this style. Use Equinox modules. Use eqx.filter_jit at the top-level call that represents the full computation. Use jax.vmap or eqx.filter_vmap when batching a pure per-example function helps the code stay simple and compiled. Run Python through the project uv environment.