Combined loss term for VQ-VAE (`diffusers.VQModel`)

For training the VQ-VAE component of a latent diffusion model a la `CompVis/ldm-celebahq-256` (which uses `diffusers.VQModel`), is there a combined loss term for each of the losses as described by the authors: reconstruction loss, vq loss, and commitment loss?

I see the vq loss term is collected in `VectorQuantizer`, but it does not seem to be used anywhere else.
https://github.com/huggingface/diffusers/blob/ebc99a77aad647c5d33eb36a33c23f7b3949cb40/src/diffusers/models/autoencoders/vae.py#L726-L730

I'm also open to alternatives to `VQModel` like `AutoEncoderKL`, if they can collect the loss terms more easily.

Thank you!