From 7ef342933e4c720383fdc053d3e4e7c2649d0b0d Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 16 Aug 2025 08:25:35 -0400 Subject: [PATCH] cleanup onnx init Signed-off-by: Vladimir Mandic --- modules/onnx_impl/__init__.py | 72 ++++++++----------- ...nx_stable_diffusion_xl_img2img_pipeline.py | 1 + modules/onnx_impl/utils.py | 3 +- modules/ui_settings.py | 3 +- 4 files changed, 35 insertions(+), 44 deletions(-) diff --git a/modules/onnx_impl/__init__.py b/modules/onnx_impl/__init__.py index 5a009a741..adfc0af5f 100644 --- a/modules/onnx_impl/__init__.py +++ b/modules/onnx_impl/__init__.py @@ -3,6 +3,7 @@ import numpy as np import torch import diffusers import onnxruntime as ort +from installer import log, installed, install initialized = False @@ -194,7 +195,6 @@ def initialize_onnx(): global initialized # pylint: disable=global-statement if initialized: return - from installer import log, installed from modules import devices if not installed('onnx', quiet=True): return @@ -203,53 +203,43 @@ def initialize_onnx(): from .execution_providers import ExecutionProvider, TORCH_DEVICE_TO_EP, available_execution_providers if devices.backend == "rocm": TORCH_DEVICE_TO_EP["cuda"] = ExecutionProvider.ROCm - from .pipelines.onnx_stable_diffusion_pipeline import OnnxStableDiffusionPipeline - 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 - - OnnxRuntimeModel.__module__ = 'diffusers' # OnnxRuntimeModel Hijack. - diffusers.OnnxRuntimeModel = OnnxRuntimeModel - - diffusers.OnnxStableDiffusionPipeline = OnnxStableDiffusionPipeline - diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionPipeline - - diffusers.OnnxStableDiffusionImg2ImgPipeline = OnnxStableDiffusionImg2ImgPipeline - diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionImg2ImgPipeline - - diffusers.OnnxStableDiffusionInpaintPipeline = OnnxStableDiffusionInpaintPipeline - diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionInpaintPipeline - - diffusers.OnnxStableDiffusionUpscalePipeline = OnnxStableDiffusionUpscalePipeline - log.debug(f'ONNX: version={ort.__version__}, 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 - - diffusers.OnnxStableDiffusionXLImg2ImgPipeline = OnnxStableDiffusionXLImg2ImgPipeline - diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion-xl"] = diffusers.OnnxStableDiffusionXLImg2ImgPipeline - - diffusers.ORTStableDiffusionXLPipeline = diffusers.OnnxStableDiffusionXLPipeline # Huggingface model compatibility - diffusers.ORTStableDiffusionXLImg2ImgPipeline = diffusers.OnnxStableDiffusionXLImg2ImgPipeline - - optimum.onnxruntime.modeling_diffusion.ORTPipelinePart.to = ORTPipelinePart_to # pylint: disable=protected-access - except Exception as e: - log.debug(f'ONNX failed to initialize XL pipelines: {e}') + log.error(f'ONNX initialization: {e}') initialized = True +def initialize_onnx_pipelines(): + try: # may fail on onnx import + import onnx # pylint: disable=unused-import + OnnxRuntimeModel.__module__ = 'diffusers' # OnnxRuntimeModel Hijack. + diffusers.OnnxRuntimeModel = OnnxRuntimeModel + from .pipelines.onnx_stable_diffusion_pipeline import OnnxStableDiffusionPipeline + 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 + diffusers.OnnxStableDiffusionPipeline = OnnxStableDiffusionPipeline + diffusers.OnnxStableDiffusionImg2ImgPipeline = OnnxStableDiffusionImg2ImgPipeline + diffusers.OnnxStableDiffusionInpaintPipeline = OnnxStableDiffusionInpaintPipeline + diffusers.OnnxStableDiffusionUpscalePipeline = OnnxStableDiffusionUpscalePipeline + diffusers.OnnxStableDiffusionXLPipeline = OnnxStableDiffusionXLPipeline + diffusers.OnnxStableDiffusionXLImg2ImgPipeline = OnnxStableDiffusionXLImg2ImgPipeline + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionPipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionImg2ImgPipeline + diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["onnx-stable-diffusion"] = diffusers.OnnxStableDiffusionInpaintPipeline + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion-xl"] = diffusers.OnnxStableDiffusionXLPipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["onnx-stable-diffusion-xl"] = diffusers.OnnxStableDiffusionXLImg2ImgPipeline + diffusers.ORTStableDiffusionXLPipeline = diffusers.OnnxStableDiffusionXLPipeline # Huggingface model compatibility + diffusers.ORTStableDiffusionXLImg2ImgPipeline = diffusers.OnnxStableDiffusionXLImg2ImgPipeline + except Exception as e: + log.error(f'ONNX initialization: {e}') + + def install_olive(): - from installer import installed, install, log if installed("olive-ai"): return try: diff --git a/modules/onnx_impl/pipelines/onnx_stable_diffusion_xl_img2img_pipeline.py b/modules/onnx_impl/pipelines/onnx_stable_diffusion_xl_img2img_pipeline.py index 5606da06e..2627ba074 100644 --- a/modules/onnx_impl/pipelines/onnx_stable_diffusion_xl_img2img_pipeline.py +++ b/modules/onnx_impl/pipelines/onnx_stable_diffusion_xl_img2img_pipeline.py @@ -2,6 +2,7 @@ from typing import Optional, Dict, Any import numpy as np import torch import onnxruntime as ort + import optimum.onnxruntime from modules.onnx_impl.pipelines import CallablePipelineBase from modules.onnx_impl.pipelines.utils import randn_tensor diff --git a/modules/onnx_impl/utils.py b/modules/onnx_impl/utils.py index b86e1ed70..80b75cb4a 100644 --- a/modules/onnx_impl/utils.py +++ b/modules/onnx_impl/utils.py @@ -4,7 +4,6 @@ import importlib from typing import Type, Tuple, Union, List, Dict, Any import torch import diffusers -import onnxruntime as ort def extract_device(args: List, kwargs: Dict): @@ -18,7 +17,7 @@ def extract_device(args: List, kwargs: Dict): return device -def move_inference_session(session: ort.InferenceSession, device: torch.device): +def move_inference_session(session, device: torch.device): # session: ort.InferenceSession from modules.devices import device as default_device from modules.devices import backend as default_backend diff --git a/modules/ui_settings.py b/modules/ui_settings.py index f15fd8985..f1def3d89 100644 --- a/modules/ui_settings.py +++ b/modules/ui_settings.py @@ -129,8 +129,9 @@ def run_settings(*args): if shared.opts.set(key, value): changed.append(key) if shared.opts.cuda_compile_backend == "olive-ai": - from modules.onnx_impl import install_olive + from modules.onnx_impl import install_olive, initialize_onnx_pipelines install_olive() + initialize_onnx_pipelines() if shared.cmd_opts.use_directml: from modules.dml import directml_override_opts directml_override_opts()