Apply LeanVAE to other Diffuser?

[Aoko955]

If we apply the LeanVAE to other Diffusers, such as Wan or Hunyuan Diffusers, should we finetune the Dit to make it adapt to the new VAE?

[hyc9]

Yes, their latent spaces are not aligned.

[hyc9]

There are basically two approaches to make LeanVAE work with other diffusion backbones like Wan or Hunyuan, and they can also be combined:

1. Align LeanVAE's latent space to the target VAE

Instead of finetuning the DiT, you can train LeanVAE so that its latent space is aligned with the target VAE's latent space. This lets LeanVAE act as a lightweight, drop-in replacement for the (much heavier) original VAE, without touching the DiT.

We've just released two experimental branches that do exactly this for Wan:

Branch	Compression / Channels	Alignment target
align-wan2.1	4×8×8, latent_dim=16	Wan2.1 VAE (z_dim=16)
align-wan2.2	4×16×16, latent_dim=48	Wan2.2 VAE

For the align-wan2.1 case, you can start training from our released LeanVAE weights, which significantly speeds up convergence. You can follow these branches as a reference for aligning LeanVAE to other VAEs as well.

2. Finetune / adapt the DiT to the new VAE

Alternatively, you can keep LeanVAE fixed and train the DiT to adapt to LeanVAE's latent space. A good reference for this direction is DC-Gen.

These two approaches are not mutually exclusive — you can also combine latent-space alignment with light DiT finetuning to get the best of both worlds.

Hope this helps!

[bee992]

Hi, I noticed that the align-wan-2-2 branch README says 4×16×16, but the actual spatial downsampling seems to be 4×8×8 — DWT does spatial ÷2, patch_size=(2,4,4) does spatial ÷4, total 2×4=8× per axis. For a 480×832 input, LeanVAE latent is 60×104 while Wan2.2 is 30×52. Is this a bug or an intentional design choice? What would be the correct configuration to train a LeanVAE that can truly drop-in replace the Wan2.2 VAE? Thanks!

[hyc9]

Hi, I noticed that the align-wan-2-2 branch README says 4×16×16, but the actual spatial downsampling seems to be 4×8×8 — DWT does spatial ÷2, patch_size=(2,4,4) does spatial ÷4, total 2×4=8× per axis. For a 480×832 input, LeanVAE latent is 60×104 while Wan2.2 is 30×52. Is this a bug or an intentional design choice? What would be the correct configuration to train a LeanVAE that can truly drop-in replace the Wan2.2 VAE? Thanks!

Thanks! They are bugs and have been fixed.