support more backends

installer.py

@@ -583,8 +583,7 @@ def install_packages():
     install(clip_package, 'clip')
     invisiblewatermark_package = os.environ.get('INVISIBLEWATERMARK_PACKAGE', "git+https://github.com/patrickvonplaten/invisible-watermark.git@remove_onnxruntime_depedency")
     install(invisiblewatermark_package, 'invisible-watermark')
-    install('olive-ai[directml]', 'olive-ai', ignore=True)
-    install('onnxruntime-directml==1.16.1', 'onnxruntime-directml', ignore=True)
+    install('olive-ai', 'olive-ai', ignore=True)
     install('pi-heif', 'pi_heif', ignore=True)
     tensorflow_package = os.environ.get('TENSORFLOW_PACKAGE', 'tensorflow==2.13.0')
     install(tensorflow_package, 'tensorflow-rocm' if 'rocm' in tensorflow_package else 'tensorflow', ignore=True)

launch.py

@@ -23,8 +23,14 @@ python = sys.executable # used by some extensions to run python
 skip_install = False # parsed by some extensions
 
 
-def init_args():
-    global parser, args # pylint: disable=global-statement
+try:
+    import torch._dynamo
+except ModuleNotFoundError:
+    sys.modules["torch._dynamo"] = {} # HACK torch 1.13.1 does not have _dynamo. will be removed.
+
+
+def init_modules():
+    global parser, args, script_path, extensions_dir # pylint: disable=global-statement
     import modules.cmd_args
     parser = modules.cmd_args.parser
     installer.add_args(parser)

modules/olive.py

@@ -4,6 +4,7 @@ import torch
 import shutil
 import diffusers
 import numpy as np
+import onnxruntime as ort
 from typing import Union, Optional, Callable, List
 from transformers.models.clip.modeling_clip import CLIPTextModel, CLIPTextModelWithProjection
 from installer import log, args
@@ -13,14 +14,23 @@ from modules.sd_models import CheckpointInfo
 
 submodels = ("text_encoder", "unet", "vae_encoder", "vae_decoder",)
 
-execution_provider = "CUDAExecutionProvider"
+available_execution_providers = ort.get_available_providers()
+execution_provider = "CUDAExecutionProvider" if "CUDAExecutionProvider" in available_execution_providers else "CPUExecutionProvider"
+execution_provider_options = {
+    "device_id": int(cmd_opts.device_id or 0),
+}
 if args.use_directml:
     execution_provider = "DmlExecutionProvider"
 elif args.use_rocm:
-    execution_provider = "ROCmExecutionProvider"
-provider = (execution_provider, {
-    "device_id": int(cmd_opts.device_id or 0),
-})
+    if "ROCMExecutionProvider" in available_execution_providers:
+        execution_provider = "ROCMExecutionProvider"
+        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 args.use_ipex or args.use_openvino:
+    execution_provider = "OpenVINOExecutionProvider"
+provider = (execution_provider, execution_provider_options,)
 
 class OnnxRuntimeModel(diffusers.OnnxRuntimeModel):
     config = {}
@@ -38,6 +48,12 @@ class OnnxStableDiffusionPipeline(diffusers.OnnxStableDiffusionPipeline):
     sd_checkpoint_info: CheckpointInfo
     sd_model_checkpoint: str
 
+    @staticmethod
+    def from_pretrained(*args, **kwargs):
+        if "provider" not in kwargs:
+            kwargs["provider"] = provider
+        return diffusers.OnnxStableDiffusionPipeline.from_pretrained(*args, **kwargs)
+
     def apply(self, dummy_pipeline):
         self.sd_model_hash = dummy_pipeline.sd_model_hash
         self.sd_checkpoint_info = dummy_pipeline.sd_checkpoint_info
@@ -229,7 +245,9 @@ class OlivePipeline(diffusers.DiffusionPipeline):
         out_dir = os.path.join(opts.olive_cached_models_path, f"{self.original_filename}-{width}w-{height}h")
         if os.path.isdir(out_dir):
             del self.unoptimized
-            return OnnxStableDiffusionPipeline.from_pretrained(out_dir, provider=provider).apply(self)
+            return OnnxStableDiffusionPipeline.from_pretrained(
+                out_dir,
+            ).apply(self)
 
         try:
             if opts.olive_cache_optimized:
@@ -279,7 +297,7 @@ class OlivePipeline(diffusers.DiffusionPipeline):
             del self.unoptimized
             for submodel in submodels:
                 kwargs[submodel] = diffusers.OnnxRuntimeModel.from_pretrained(
-                    os.path.dirname(optimized_model_paths[submodel]), provider=provider,
+                    os.path.dirname(optimized_model_paths[submodel]),
                 )
 
             pipeline = OnnxStableDiffusionPipeline(

modules/sd_models.py

@@ -793,8 +793,8 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
         if os.path.isdir(checkpoint_info.path):
             if shared.opts.olive_sideloaded_models_path in checkpoint_info.path:
                 try:
-                    from modules.olive import OnnxStableDiffusionPipeline, provider
-                    sd_model = OnnxStableDiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.olive_sideloaded_models_path, provider=provider)
+                    from modules.olive import OnnxStableDiffusionPipeline
+                    sd_model = OnnxStableDiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.olive_sideloaded_models_path)
                     sd_model.model_type = sd_model.__class__.__name__
                 except Exception as e:
                     shared.log.error(f'Failed loading {op}: {checkpoint_info.path} olive={e}')

modules/shared_items.py

@@ -26,7 +26,7 @@ def list_crossattention():
 
 def get_pipelines():
     import diffusers
-    from modules.olive import OlivePipeline
+    from modules.olive import OnnxStableDiffusionPipeline, OlivePipeline
     from installer import log
     pipelines = { # note: not all pipelines can be used manually as they require prior pipeline next to decoder pipeline
         'Autodetect': None,
@@ -39,7 +39,8 @@ def get_pipelines():
         'Stable Diffusion XL Img2Img': getattr(diffusers, 'StableDiffusionXLImg2ImgPipeline', None),
         'Stable Diffusion XL Inpaint': getattr(diffusers, 'StableDiffusionXLInpaintPipeline', None),
         'Stable Diffusion XL Instruct': getattr(diffusers, 'StableDiffusionXLInstructPix2PixPipeline', None),
-        'ONNX Stable Diffusion (Olive)': OlivePipeline,
+        'ONNX Stable Diffusion': OnnxStableDiffusionPipeline,
+        'ONNX Stable Diffusion with Olive': OlivePipeline,
         'Latent Consistency Model': getattr(diffusers, 'LatentConsistencyModelPipeline', None),
         'PixArt Alpha': getattr(diffusers, 'PixArtAlphaPipeline', None),
         'UniDiffuser': getattr(diffusers, 'UniDiffuserPipeline', None),

requirements.txt

@@ -65,7 +65,3 @@ Pillow==9.5.0
 timm==0.9.7
 pydantic==1.10.13
 typing-extensions==4.8.0
-
-torch==1.13.1
-torchvision==0.14.1
-torch-directml==0.1.13.1.dev230413