Add backwards compatibility w/ older versions of diffusers (<0.25.0) (d…

…eepspeedai#5083)

This PR adds backwards compatibility for older versions of `diffusers`
(`<0.25.0`) by updating the `vae` container import logic to account for
changes between the various versions.

deepspeed/module_inject/containers/vae.py

@@ -13,11 +13,17 @@ def __init__(self):
         super().__init__()
         try:
             import diffusers
-            if hasattr(diffusers.models.autoencoders.vae, "AutoencoderKL"):
-                self._orig_layer_class = diffusers.models.autoencoders.vae.AutoencoderKL
-            else:
-                # Diffusers >= 0.12.0 changes location of AutoencoderKL
+            if hasattr(diffusers.models, "autoencoders"):
+                # Diffusers >= 0.25.0
+                # Changes location to 'autoencoders' directory
                 self._orig_layer_class = diffusers.models.autoencoders.autoencoder_kl.AutoencoderKL
+            elif hasattr(diffusers.models.vae, "AutoencoderKL"):
+                # Diffusers < 0.12.0
+                self._orig_layer_class = diffusers.models.vae.AutoencoderKL
+            else:
+                # Diffusers >= 0.12.0 & < 0.25.0
+                # Changes location of AutoencoderKL
+                self._orig_layer_class = diffusers.models.autoencoder_kl.AutoencoderKL
         except ImportError:
             self._orig_layer_class = None