Reproduce distill experiment on Wan2.2

[Lubenwei-nb123]

First thanks for this great work on distilling vae with only little psnr loss. I wonder how to design the student vae for Wan2.2 and what parameter should I change.

[ZouYa99]

Hi,
Thank you for your attention and reproduction efforts. In fact, I had also designed a student VAE decoder for Wan 2.2, but it was not completed due to time constraints. You only need to modify the configuration file.

Regarding the "decoder_block_out_channels" and "decoder_layers_per_block" parameters, you can adjust them according to your specific needs. I typically determine these values based on both the parameter count and the VRAM usage during decoding. To ensure that videos of the target resolution do not cause OOM errors on mobile devices during decoding, I set the VRAM ceiling to around 3.8GB.

Note that you need to modify "latent_channels" to match the corresponding latent channel count. For example, Wan 2.2 VAE corresponds to 48 channels.

The product of the number of "true" values in "spatio_temporal_scaling" and the "patch_size" collectively determines the spatial downsampling factor. For Wan 2.2 VAE's 16× spatial downsampling, I recommend using a "patch_size" of 2 and setting "spatio_temporal_scaling": [
true,
true,
true,
false
].

The "decoder_spatio_only" parameter is used to adjust the temporal downsampling factor. Since layers where "spatio_temporal_scaling" is set to true perform both spatial and temporal downsampling simultaneously, setting "decoder_spatio_only" to true indicates that no temporal downsampling is performed in that layer. For Wan 2.2 VAE's 4× temporal downsampling, you can use "decoder_spatio_only": [
true,
false,
false,
false
]

For the feature projection layer settings, you can design them based on the dimensions of the corresponding blocks in the teacher and student models.

Additionally, I would like to remind you that to achieve feature alignment, we need the features returned by the teacher model during decoding. You can refer to the addition of the "feature_enabled" parameter in the diffusers_vae.src.diffusers.models.autoencoders.autoencoder_kl_ltx.py file to implement this.

[eisneim]

@Lubenwei-nb123 have you got any experiment results? i'm planing to do the same