Skip to content

Pass explicit latent layout to the GenCeption VAE decoder - #13

Open
sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix/explicit-vae-latent-layout
Open

sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix/explicit-vae-latent-layout

Conversation

@sylvesterkaczmarek

Copy link
Copy Markdown
Contributor

Pass channels-last latents explicitly from GenCeptionPipeline.decode_video() to the VAE, whose preferred input layout is [B, T, H, W, C].

The pipeline's input contract is channels-first. The VAE currently guesses the layout from whether the last dimension equals z_dim. When latent width equals the 16 channels, it treats width as channels and decodes the wrong temporal/spatial axes. With a real reduced-size VAE, input (1, 16, 2, 1, 16) produces (1, 61, 16, 8, 3) instead of (1, 5, 8, 128, 3).

The fix adds an explicit transpose after latent denormalization. No change to the VAE's public layout heuristic, normalization, or checkpoint loading.

Validation

CPU, macOS arm64, Python 3.12, JAX 0.7.2, Flax 0.11.2:

  • The real-VAE width-16 regression fails on unchanged 9912d29; width 3 is an unchanged control.
  • Both new tests pass. Coverage includes exact decoder-input values and real VAE outputs for batches 1/2 and latent widths 3/16.
  • Built the wheel, installed it into the test environment, and reran both tests outside the source checkout: passed.
  • Ruff E9/F, Pyink on the new test, compilation, and git diff --check passed.

Tests use small random VAE weights, not pretrained weights or a transformer forward pass. GPU inference was not run. The unpinned JAX 0.11.2/Flax 0.12.9 environment fails in Flax NNX import (HiPrimitive) even before the patch; dependency compatibility is unchanged and was not bundled into this fix.

@sylvesterkaczmarek
sylvesterkaczmarek marked this pull request as ready for review September 19, 2026 18:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant