In VAE context, KL is computed for latent z, acting as a regularizer. In that scenario, the simplest distribution can be used as a prior - standard independent gaussian.
Your snippet already computes std. gaussian KL, and you only have to switch to some loss based on pixel color distances (seems that simple MSE can be used for this).