[Vae] Make sure all vae's work with latent diffusion models (#5880)
* add comments to explain the code better
* add comments to explain the code better
* add comments to explain the code better
* add comments to explain the code better
* add comments to explain the code better
* fix more
* fix more
* fix more
* fix more
* fix more
* fix more