add optimization on gpu

configs/olive/config_text_encoder.json

@@ -4,7 +4,7 @@
     "config": {
       "model_path": "",
       "model_loader": "text_encoder_load",
-      "model_script": "modules/olive.py",
+      "model_script": "modules/onnx.py",
       "io_config": {
         "input_names": ["input_ids"],
         "output_names": ["last_hidden_state", "pooler_output"],
@@ -29,7 +29,7 @@
           "type": "latency",
           "sub_types": [{ "name": "avg" }],
           "user_config": {
-            "user_script": "modules/olive.py",
+            "user_script": "modules/onnx.py",
             "dataloader_func": "text_encoder_data_loader",
             "batch_size": 1
           }

configs/olive/config_unet.json

@@ -4,7 +4,7 @@
     "config": {
       "model_path": "",
       "model_loader": "unet_load",
-      "model_script": "modules/olive.py",
+      "model_script": "modules/onnx.py",
       "io_config": {
         "input_names": [
           "sample",
@@ -46,7 +46,7 @@
           "type": "latency",
           "sub_types": [{ "name": "avg" }],
           "user_config": {
-            "user_script": "modules/olive.py",
+            "user_script": "modules/onnx.py",
             "dataloader_func": "unet_data_loader",
             "batch_size": 2
           }

configs/olive/config_vae_decoder.json

@@ -4,7 +4,7 @@
     "config": {
       "model_path": "",
       "model_loader": "vae_decoder_load",
-      "model_script": "modules/olive.py",
+      "model_script": "modules/onnx.py",
       "io_config": {
         "input_names": ["latent_sample", "return_dict"],
         "output_names": ["sample"],
@@ -36,7 +36,7 @@
           "type": "latency",
           "sub_types": [{ "name": "avg" }],
           "user_config": {
-            "user_script": "modules/olive.py",
+            "user_script": "modules/onnx.py",
             "dataloader_func": "vae_decoder_data_loader",
             "batch_size": 1
           }

configs/olive/config_vae_encoder.json

@@ -4,7 +4,7 @@
     "config": {
       "model_path": "",
       "model_loader": "vae_encoder_load",
-      "model_script": "modules/olive.py",
+      "model_script": "modules/onnx.py",
       "io_config": {
         "input_names": ["sample", "return_dict"],
         "output_names": ["latent_sample"],
@@ -36,7 +36,7 @@
           "type": "latency",
           "sub_types": [{ "name": "avg" }],
           "user_config": {
-            "user_script": "modules/olive.py",
+            "user_script": "modules/onnx.py",
             "dataloader_func": "vae_encoder_data_loader",
             "batch_size": 1
           }

modules/onnx.py

@@ -22,10 +22,12 @@ if args.use_directml:
     execution_provider_options["device_id"] = int(cmd_opts.device_id or 0)
 elif args.use_rocm:
     if "ROCMExecutionProvider" in available_execution_providers:
+        from olive.hardware.accelerator import AcceleratorLookup
         execution_provider = "ROCMExecutionProvider"
         execution_provider_options["device_id"] = int(cmd_opts.device_id or 0)
         execution_provider_options["tunable_op_enable"] = 1
         execution_provider_options["tunable_op_tuning_enable"] = 1
+        AcceleratorLookup.EXECUTION_PROVIDERS["gpu"].append("ROCMExecutionProvider")
     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:
@@ -213,6 +215,7 @@ class OlivePipeline(diffusers.DiffusionPipeline):
     sd_model_hash: str
     sd_checkpoint_info: CheckpointInfo
     sd_model_checkpoint: str
+    config = {}
 
     unoptimized: diffusers.DiffusionPipeline
     original_filename: str
@@ -270,6 +273,7 @@ class OlivePipeline(diffusers.DiffusionPipeline):
 
                 with open(os.path.join(sd_configs_path, "olive", f"config_{submodel}.json"), "r") as config_file:
                     olive_config = json.load(config_file)
+                olive_config["engine"]["execution_providers"] = [execution_provider]
                 olive_config["passes"]["optimize"]["config"]["float16"] = opts.olive_float16
 
                 run(olive_config)

modules/sd_models.py

@@ -793,7 +793,7 @@ 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
+                    from modules.onnx 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:

modules/shared_items.py

@@ -26,7 +26,7 @@ def list_crossattention():
 
 def get_pipelines():
     import diffusers
-    from modules.olive import OnnxStableDiffusionPipeline, OlivePipeline
+    from modules.onnx 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,