diff --git a/configs/olive/sd/text_encoder.json b/configs/olive/sd/text_encoder.json index 91a3d57d3..4a27eec4c 100644 --- a/configs/olive/sd/text_encoder.json +++ b/configs/olive/sd/text_encoder.json @@ -115,7 +115,7 @@ } } }, - "pass_flows": [["optimize"]], + "pass_flows": [["optimize_AutoExecutionProvider"]], "engine": { "search_strategy": { "execution_order": "joint", diff --git a/configs/olive/sd/unet.json b/configs/olive/sd/unet.json index 907504b06..ac8430f78 100644 --- a/configs/olive/sd/unet.json +++ b/configs/olive/sd/unet.json @@ -132,7 +132,7 @@ } } }, - "pass_flows": [["optimize"]], + "pass_flows": [["optimize_AutoExecutionProvider"]], "engine": { "search_strategy": { "execution_order": "joint", diff --git a/configs/olive/sd/vae_decoder.json b/configs/olive/sd/vae_decoder.json index 2e21c0f79..daea46948 100644 --- a/configs/olive/sd/vae_decoder.json +++ b/configs/olive/sd/vae_decoder.json @@ -122,7 +122,7 @@ } } }, - "pass_flows": [["optimize"]], + "pass_flows": [["optimize_AutoExecutionProvider"]], "engine": { "search_strategy": { "execution_order": "joint", diff --git a/configs/olive/sd/vae_encoder.json b/configs/olive/sd/vae_encoder.json index 2f012b187..fdfb1d56b 100644 --- a/configs/olive/sd/vae_encoder.json +++ b/configs/olive/sd/vae_encoder.json @@ -122,7 +122,7 @@ } } }, - "pass_flows": [["optimize"]], + "pass_flows": [["optimize_AutoExecutionProvider"]], "engine": { "search_strategy": { "execution_order": "joint", diff --git a/configs/olive/sdxl/text_encoder.json b/configs/olive/sdxl/text_encoder.json index e1fc1d672..5a80f087d 100644 --- a/configs/olive/sdxl/text_encoder.json +++ b/configs/olive/sdxl/text_encoder.json @@ -130,7 +130,7 @@ } } }, - "pass_flows": [["optimize"]], + "pass_flows": [["optimize_AutoExecutionProvider"]], "engine": { "search_strategy": { "execution_order": "joint", diff --git a/configs/olive/sdxl/text_encoder_2.json b/configs/olive/sdxl/text_encoder_2.json index 479c48e4a..287283513 100644 --- a/configs/olive/sdxl/text_encoder_2.json +++ b/configs/olive/sdxl/text_encoder_2.json @@ -170,7 +170,7 @@ } } }, - "pass_flows": [["optimize"]], + "pass_flows": [["optimize_AutoExecutionProvider"]], "engine": { "search_strategy": { "execution_order": "joint", diff --git a/configs/olive/sdxl/unet.json b/configs/olive/sdxl/unet.json index f5baf75f3..440c7e1fe 100644 --- a/configs/olive/sdxl/unet.json +++ b/configs/olive/sdxl/unet.json @@ -120,7 +120,7 @@ } } }, - "pass_flows": [["optimize"]], + "pass_flows": [["optimize_AutoExecutionProvider"]], "engine": { "search_strategy": { "execution_order": "joint", diff --git a/configs/olive/sdxl/vae_decoder.json b/configs/olive/sdxl/vae_decoder.json index d5409a511..40a3eedd9 100644 --- a/configs/olive/sdxl/vae_decoder.json +++ b/configs/olive/sdxl/vae_decoder.json @@ -130,7 +130,7 @@ } } }, - "pass_flows": [["optimize"]], + "pass_flows": [["optimize_AutoExecutionProvider"]], "engine": { "search_strategy": { "execution_order": "joint", diff --git a/configs/olive/sdxl/vae_encoder.json b/configs/olive/sdxl/vae_encoder.json index 3e871d524..12404142c 100644 --- a/configs/olive/sdxl/vae_encoder.json +++ b/configs/olive/sdxl/vae_encoder.json @@ -110,7 +110,7 @@ } } }, - "pass_flows": [["optimize"]], + "pass_flows": [["optimize_AutoExecutionProvider"]], "engine": { "search_strategy": { "execution_order": "joint", diff --git a/modules/onnx_pipelines.py b/modules/onnx_pipelines.py index 2a0a25722..b2d40e46f 100644 --- a/modules/onnx_pipelines.py +++ b/modules/onnx_pipelines.py @@ -274,22 +274,23 @@ class OnnxRawPipeline(OnnxPipelineBase): log.info(f"\nProcessing {submodel}") with open(os.path.join(sd_configs_path, "olive", 'sdxl' if self._is_sdxl else 'sd', f"{submodel}.json"), "r") as config_file: - olive_config = json.load(config_file) - pass_key = f"optimize_{shared.opts.onnx_execution_provider}" + olive_config: Dict[str, Dict[str, Dict]] = json.load(config_file) + for flow in olive_config["pass_flows"]: for i in range(len(flow)): - if flow[i] == "optimize": - flow[i] = pass_key + flow[i] = flow[i].replace("AutoExecutionProvider", shared.opts.onnx_execution_provider) olive_config["input_model"]["config"]["model_path"] = os.path.abspath(os.path.join(in_dir, submodel, "model.onnx")) olive_config["engine"]["execution_providers"] = [shared.opts.onnx_execution_provider] - if pass_key in olive_config["passes"]: - float16 = shared.opts.olive_float16 and not (submodel == "vae_encoder" and shared.opts.olive_vae_encoder_float32) - olive_config["passes"][pass_key]["config"]["float16"] = float16 - if shared.opts.onnx_execution_provider == ExecutionProvider.CUDA or shared.opts.onnx_execution_provider == ExecutionProvider.ROCm: - if version.parse(ort.__version__) < version.parse("1.17.0"): - olive_config["passes"][pass_key]["config"]["optimization_options"] = {"enable_skip_group_norm": False} - if float16: - olive_config["passes"][pass_key]["config"]["keep_io_types"] = False + + for pass_key in olive_config["passes"]: + if olive_config["passes"][pass_key]["type"] == "OrtTransformersOptimization": + float16 = shared.opts.olive_float16 and not (submodel == "vae_encoder" and shared.opts.olive_vae_encoder_float32) + olive_config["passes"][pass_key]["config"]["float16"] = float16 + if shared.opts.onnx_execution_provider == ExecutionProvider.CUDA or shared.opts.onnx_execution_provider == ExecutionProvider.ROCm: + if version.parse(ort.__version__) < version.parse("1.17.0"): + olive_config["passes"][pass_key]["config"]["optimization_options"] = {"enable_skip_group_norm": False} + if float16: + olive_config["passes"][pass_key]["config"]["keep_io_types"] = False run(olive_config)