From b7cfd09d50c35e490e7baab5d4aa2512b447171b Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Sun, 29 Oct 2023 20:51:02 +0900 Subject: [PATCH] add optimization on gpu --- configs/olive/config_text_encoder.json | 4 ++-- configs/olive/config_unet.json | 4 ++-- configs/olive/config_vae_decoder.json | 4 ++-- configs/olive/config_vae_encoder.json | 4 ++-- modules/{olive.py => onnx.py} | 4 ++++ modules/sd_models.py | 2 +- modules/shared_items.py | 2 +- 7 files changed, 14 insertions(+), 10 deletions(-) rename modules/{olive.py => onnx.py} (98%) diff --git a/configs/olive/config_text_encoder.json b/configs/olive/config_text_encoder.json index e4ae6fcb1..749c29031 100644 --- a/configs/olive/config_text_encoder.json +++ b/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 } diff --git a/configs/olive/config_unet.json b/configs/olive/config_unet.json index d370ccf56..f8d52d752 100644 --- a/configs/olive/config_unet.json +++ b/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 } diff --git a/configs/olive/config_vae_decoder.json b/configs/olive/config_vae_decoder.json index 86095098f..0e726cbbe 100644 --- a/configs/olive/config_vae_decoder.json +++ b/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 } diff --git a/configs/olive/config_vae_encoder.json b/configs/olive/config_vae_encoder.json index e00dc01f6..65e0e1187 100644 --- a/configs/olive/config_vae_encoder.json +++ b/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 } diff --git a/modules/olive.py b/modules/onnx.py similarity index 98% rename from modules/olive.py rename to modules/onnx.py index 4856e6d0b..1f83ab9fd 100644 --- a/modules/olive.py +++ b/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) diff --git a/modules/sd_models.py b/modules/sd_models.py index 71952d552..a0ed5b6b8 100644 --- a/modules/sd_models.py +++ b/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: diff --git a/modules/shared_items.py b/modules/shared_items.py index 5e7a75427..0a9b4ec1b 100644 --- a/modules/shared_items.py +++ b/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,