Improve the Autoencoder KL model using the Stable Diffusion approach for the adversarial component and using a better learning rate scheduling.