Fix compatibility with diffusers < 0.25.0 (#952)

huggingface · Oct 15, 2024 · f7d379d · f7d379d
1 parent 93ca007
commit f7d379d
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/optimum/intel/openvino/modeling_diffusion.py b/optimum/intel/openvino/modeling_diffusion.py
@@ -28,7 +28,6 @@
 import openvino
 import torch
 from diffusers.configuration_utils import ConfigMixin
-from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
 from diffusers.pipelines import (
     AutoPipelineForImage2Image,
     AutoPipelineForInpainting,
@@ -63,6 +62,7 @@
 )
 
 from ...exporters.openvino import main_export
+from ..utils.import_utils import is_diffusers_version
 from .configuration import OVConfig, OVQuantizationMethod, OVWeightQuantizationConfig
 from .loaders import OVTextualInversionLoaderMixin
 from .modeling_base import OVBaseModel
@@ -76,6 +76,12 @@
 )
 
 
+if is_diffusers_version(">=", "0.25.0"):
+    from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution
+else:
+    from diffusers.models.vae import DiagonalGaussianDistribution
+
+
 core = Core()
 
 logger = logging.getLogger(__name__)