split onnx init into two passes

This commit is contained in:
Seunghoon Lee
2024-07-02 22:46:13 +09:00
parent 803558991d
commit 55ea007369
+14 -7
View File
@@ -200,8 +200,6 @@ def initialize_onnx():
return
try: # may fail on onnx import
import onnx # pylint: disable=unused-import
import optimum.onnxruntime
optimum.onnxruntime.modeling_diffusion._ORTDiffusionModelPart.to = ORTDiffusionModelPart_to # pylint: disable=protected-access
from .execution_providers import ExecutionProvider, TORCH_DEVICE_TO_EP, available_execution_providers
if devices.backend == "rocm":
TORCH_DEVICE_TO_EP["cuda"] = ExecutionProvider.ROCm
@@ -209,8 +207,6 @@ def initialize_onnx():
from .pipelines.onnx_stable_diffusion_img2img_pipeline import OnnxStableDiffusionImg2ImgPipeline
from .pipelines.onnx_stable_diffusion_inpaint_pipeline import OnnxStableDiffusionInpaintPipeline
from .pipelines.onnx_stable_diffusion_upscale_pipeline import OnnxStableDiffusionUpscalePipeline
from .pipelines.onnx_stable_diffusion_xl_pipeline import OnnxStableDiffusionXLPipeline
from .pipelines.onnx_stable_diffusion_xl_img2img_pipeline import OnnxStableDiffusionXLImg2ImgPipeline
OnnxRuntimeModel.__module__ = 'diffusers' # OnnxRuntimeModel Hijack.
diffusers.OnnxRuntimeModel = OnnxRuntimeModel
@@ -226,6 +222,16 @@ def initialize_onnx():
diffusers.OnnxStableDiffusionUpscalePipeline = OnnxStableDiffusionUpscalePipeline
log.debug(f'ONNX: version={ort.__version__} provider={opts.onnx_execution_provider}, available={available_execution_providers}')
except Exception as e:
log.error(f'ONNX failed to initialize: {e}')
try:
# load xl pipelines. may fail if the user has the latest diffusers (0.30.x)
import optimum.onnxruntime
from .pipelines.onnx_stable_diffusion_xl_pipeline import OnnxStableDiffusionXLPipeline
from .pipelines.onnx_stable_diffusion_xl_img2img_pipeline import OnnxStableDiffusionXLImg2ImgPipeline
diffusers.OnnxStableDiffusionXLPipeline = OnnxStableDiffusionXLPipeline
diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion-xl"] = diffusers.OnnxStableDiffusionXLPipeline
@@ -235,9 +241,10 @@ def initialize_onnx():
diffusers.ORTStableDiffusionXLPipeline = diffusers.OnnxStableDiffusionXLPipeline # Huggingface model compatibility
diffusers.ORTStableDiffusionXLImg2ImgPipeline = diffusers.OnnxStableDiffusionXLImg2ImgPipeline
log.debug(f'ONNX: version={ort.__version__} provider={opts.onnx_execution_provider}, available={available_execution_providers}')
except Exception as e:
log.error(f'ONNX failed to initialize: {e}')
optimum.onnxruntime.modeling_diffusion._ORTDiffusionModelPart.to = ORTDiffusionModelPart_to # pylint: disable=protected-access
except Exception:
pass
initialized = True