From e7f7bbc9eaa4ef4e62816d38e68c1ddf99217e08 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Thu, 28 Dec 2023 22:49:17 -0600 Subject: [PATCH] wrap some diffusers imports for version safety --- api/onnx_web/diffusers/patches/vae.py | 3 +-- api/onnx_web/diffusers/version_safe_diffusers.py | 12 +++++++++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/api/onnx_web/diffusers/patches/vae.py b/api/onnx_web/diffusers/patches/vae.py index d7e3e1f7..bff1f5ef 100644 --- a/api/onnx_web/diffusers/patches/vae.py +++ b/api/onnx_web/diffusers/patches/vae.py @@ -4,9 +4,8 @@ from typing import Union import numpy as np import torch from diffusers import OnnxRuntimeModel -from diffusers.models.autoencoder_kl import AutoencoderKLOutput -from diffusers.models.vae import DecoderOutput from diffusers.pipelines.onnx_utils import ORT_TO_NP_TYPE +from ..version_safe_diffusers import AutoencoderKLOutput, DecoderOutput from ...server import ServerContext diff --git a/api/onnx_web/diffusers/version_safe_diffusers.py b/api/onnx_web/diffusers/version_safe_diffusers.py index 8c4da406..d76a11af 100644 --- a/api/onnx_web/diffusers/version_safe_diffusers.py +++ b/api/onnx_web/diffusers/version_safe_diffusers.py @@ -5,7 +5,9 @@ from packaging import version is_diffusers_0_15 = version.parse( version.parse(diffusers.__version__).base_version ) >= version.parse("0.15") - +is_diffusers_0_24 = version.parse( + version.parse(diffusers.__version__).base_version +) >= version.parse("0.24") try: from diffusers import DEISMultistepScheduler @@ -32,3 +34,11 @@ if is_diffusers_0_15: from diffusers.models.attention_processor import AttnProcessor else: from diffusers.models.cross_attention import CrossAttnProcessor as AttnProcessor + + +if is_diffusers_0_24: + from diffusers.models.modeling_outputs import AutoencoderKLOutput + from diffusers.models.autoencoders.vae import DecoderOutput +else: + from diffusers.models.autoencoder_kl import AutoencoderKLOutput + from diffusers.models.vae import DecoderOutput \ No newline at end of file