separate onnx & add pipeline for SDXL Img2Img

modules/onnx_ep.py

@@ -0,0 +1,71 @@
+from enum import Enum
+from typing import Tuple, List
+import onnxruntime as ort
+from installer import log
+
+
+class ExecutionProvider(str, Enum):
+    CPU = "CPUExecutionProvider"
+    DirectML = "DmlExecutionProvider"
+    CUDA = "CUDAExecutionProvider"
+    ROCm = "ROCMExecutionProvider"
+    OpenVINO = "OpenVINOExecutionProvider"
+
+
+available_execution_providers: List[ExecutionProvider] = ort.get_available_providers()
+EP_TO_NAME = {
+    ExecutionProvider.CPU: "cpu",
+    ExecutionProvider.DirectML: "gpu-dml",
+    ExecutionProvider.CUDA: "gpu-?", # TODO
+    ExecutionProvider.ROCm: "gpu-rocm",
+    ExecutionProvider.OpenVINO: "gpu", # Other devices can use --use-openvino instead of olive
+}
+
+
+def get_default_execution_provider() -> ExecutionProvider:
+    from modules import devices
+
+    if devices.backend == "cpu":
+        return ExecutionProvider.CPU
+    elif devices.backend == "directml":
+        return ExecutionProvider.DirectML
+    elif devices.backend == "cuda":
+        return ExecutionProvider.CUDA
+    elif devices.backend == "rocm":
+        if ExecutionProvider.ROCm in available_execution_providers:
+            return ExecutionProvider.ROCm
+        else:
+            log.warning("Currently, there's no pypi release for onnxruntime-rocm. Please download and install .whl file from https://download.onnxruntime.ai/")
+    elif devices.backend == "ipex" or devices.backend == "openvino":
+        return ExecutionProvider.OpenVINO
+    return ExecutionProvider.CPU
+
+
+def get_execution_provider_options():
+    from modules.shared import cmd_opts, opts
+
+    execution_provider_options = {
+        "device_id": int(cmd_opts.device_id or 0),
+    }
+
+    if opts.onnx_execution_provider == ExecutionProvider.ROCm:
+        if ExecutionProvider.ROCm in available_execution_providers:
+            execution_provider_options["tunable_op_enable"] = 1
+            execution_provider_options["tunable_op_tuning_enable"] = 1
+        else:
+            log.warning("Currently, there's no pypi release for onnxruntime-rocm. Please download and install .whl file from https://download.onnxruntime.ai/ The inference will be fall back to CPU.")
+    elif opts.onnx_execution_provider == ExecutionProvider.OpenVINO:
+        from modules.intel.openvino import get_device as get_raw_openvino_device
+        raw_openvino_device = get_raw_openvino_device()
+        if opts.onnx_olive_float16 and not opts.openvino_hetero_gpu:
+            raw_openvino_device = f"{raw_openvino_device}_FP16"
+        execution_provider_options["device_type"] = raw_openvino_device
+        del execution_provider_options["device_id"]
+
+    return execution_provider_options
+
+
+def get_provider() -> Tuple:
+    from modules.shared import opts
+
+    return (opts.onnx_execution_provider, get_execution_provider_options(),)

modules/onnx_utils.py

@@ -0,0 +1,93 @@
+import os
+import importlib
+from typing import Type, Tuple, Union, List, Dict, Any
+import diffusers
+import onnxruntime as ort
+from installer import log
+
+
+def get_sess_options(batch_size: int, height: int, width: int, is_sdxl: bool) -> ort.SessionOptions:
+    sess_options = ort.SessionOptions()
+    sess_options.enable_mem_pattern = False
+    sess_options.add_free_dimension_override_by_name("unet_sample_batch", batch_size * 2)
+    sess_options.add_free_dimension_override_by_name("unet_sample_channels", 4)
+    sess_options.add_free_dimension_override_by_name("unet_sample_height", height // 8)
+    sess_options.add_free_dimension_override_by_name("unet_sample_width", width // 8)
+    sess_options.add_free_dimension_override_by_name("unet_time_batch", 1)
+    sess_options.add_free_dimension_override_by_name("unet_hidden_batch", batch_size * 2)
+    sess_options.add_free_dimension_override_by_name("unet_hidden_sequence", 77)
+    if is_sdxl:
+        sess_options.add_free_dimension_override_by_name("unet_text_embeds_batch", batch_size * 2)
+        sess_options.add_free_dimension_override_by_name("unet_text_embeds_size", 1280)
+        sess_options.add_free_dimension_override_by_name("unet_time_ids_batch", batch_size * 2)
+        sess_options.add_free_dimension_override_by_name("unet_time_ids_size", 6)
+    return sess_options
+
+
+def load_init_dict(cls: Type[diffusers.DiffusionPipeline], path: os.PathLike):
+    merged: Dict[str, Any] = {}
+    extracted = cls.extract_init_dict(diffusers.DiffusionPipeline.load_config(path))
+    for dict in extracted:
+        merged.update(dict)
+    merged = merged.items()
+    R: Dict[str, Tuple[str]] = {}
+    for k, v in merged:
+        if isinstance(v, list):
+            if v[0] is None or v[1] is None:
+                log.debug(f"Skipping {k} while loading init dict of '{path}': {v}")
+                continue
+            R[k] = v
+    return R
+
+
+def check_pipeline_sdxl(cls: Type[diffusers.DiffusionPipeline]) -> bool:
+    return 'XL' in cls.__name__
+
+
+def load_submodel(path: os.PathLike, is_sdxl: bool, submodel_name: str, item: List[Union[str, None]], **kwargs_ort):
+    lib, atr = item
+    if lib is None or atr is None:
+        return None
+    library = importlib.import_module(lib)
+    attribute = getattr(library, atr)
+    path = os.path.join(path, submodel_name)
+    if issubclass(attribute, diffusers.OnnxRuntimeModel):
+        return diffusers.OnnxRuntimeModel.load_model(
+            os.path.join(path, "model.onnx"),
+            **kwargs_ort,
+        ) if is_sdxl else diffusers.OnnxRuntimeModel.from_pretrained(
+            path,
+            **kwargs_ort,
+        )
+    return attribute.from_pretrained(path)
+
+
+def load_submodels(path: os.PathLike, is_sdxl: bool, init_dict: Dict[str, Type], **kwargs_ort):
+    loaded = {}
+    for k, v in init_dict.items():
+        if not isinstance(v, list):
+            loaded[k] = v
+            continue
+        try:
+            loaded[k] = load_submodel(path, is_sdxl, k, v, **kwargs_ort)
+        except Exception:
+            pass
+    return loaded
+
+
+def patch_kwargs(cls: Type[diffusers.DiffusionPipeline], kwargs: Dict) -> Dict:
+    from modules import onnx_pipelines as pipelines
+    if cls == pipelines.OnnxStableDiffusionPipeline or cls == pipelines.OnnxStableDiffusionImg2ImgPipeline or cls == pipelines.OnnxStableDiffusionInpaintPipeline:
+        kwargs["safety_checker"] = None
+        kwargs["requires_safety_checker"] = False
+    if cls == pipelines.OnnxStableDiffusionXLPipeline or cls == pipelines.OnnxStableDiffusionXLImg2ImgPipeline:
+        kwargs["config"] = {}
+
+    return kwargs
+
+
+def load_pipeline(cls: Type[diffusers.DiffusionPipeline], path: os.PathLike, **kwargs_ort):
+    if os.path.isdir(path):
+        return cls(**patch_kwargs(cls, load_submodels(path, check_pipeline_sdxl(cls), load_init_dict(cls, path), **kwargs_ort)))
+    else:
+        return cls.from_single_file(path)

modules/processing_diffusers.py

@@ -7,6 +7,7 @@ import torch
 import torchvision.transforms.functional as TF
 import diffusers
 from modules import shared, devices, processing, sd_samplers, sd_models, images, errors, masking, prompt_parser_diffusers, sd_hijack_hypertile, processing_correction, processing_vae
+from modules.onnx_pipelines import OnnxStableDiffusionPipeline
 
 
 debug = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None

modules/shared.py

@@ -16,7 +16,7 @@ from rich.console import Console
 from modules import errors, shared_items, shared_state, cmd_args, theme
 from modules.paths import models_path, script_path, data_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, extensions_dir, extensions_builtin_dir # pylint: disable=W0611
 from modules.dml import memory_providers, default_memory_provider, directml_do_hijack
-from modules.onnx import available_execution_providers, get_default_execution_provider
+from modules.onnx_ep import available_execution_providers, get_default_execution_provider
 import modules.interrogate
 import modules.memmon
 import modules.styles

modules/shared_items.py

@@ -26,8 +26,11 @@ def list_crossattention():
 
 def get_pipelines():
     import diffusers
-    import modules.onnx # pylint: disable=unused-import
     from installer import log
+    from modules.onnx import initialize as initialize_onnx_pipelines
+
+    initialize_onnx_pipelines()
+
     pipelines = { # note: not all pipelines can be used manually as they require prior pipeline next to decoder pipeline
         'Autodetect': None,
         'Stable Diffusion': getattr(diffusers, 'StableDiffusionPipeline', None),