cleanup onnx init

Signed-off-by: Vladimir Mandic <mandic00@live.com>

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:

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

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
 

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()