From bdd6b2b8f75520e32a49d45ae3ab364178685484 Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Fri, 29 Dec 2023 02:23:45 +0900 Subject: [PATCH] vae override for onnx --- configs/olive/sd/text_encoder.json | 112 ++++++++++++++++++++- configs/olive/sd/unet.json | 129 ++++++++++++++++++++++++- configs/olive/sd/vae_decoder.json | 119 ++++++++++++++++++++++- configs/olive/sd/vae_encoder.json | 119 ++++++++++++++++++++++- configs/olive/sdxl/text_encoder.json | 2 +- configs/olive/sdxl/text_encoder_2.json | 2 +- configs/olive/sdxl/unet.json | 2 +- configs/olive/sdxl/vae_decoder.json | 2 +- configs/olive/sdxl/vae_encoder.json | 2 +- modules/olive.py | 22 +++-- modules/onnx_ep.py | 1 - modules/onnx_pipelines.py | 11 ++- modules/ui_onnx.py | 27 +++++- 13 files changed, 528 insertions(+), 22 deletions(-) diff --git a/configs/olive/sd/text_encoder.json b/configs/olive/sd/text_encoder.json index a6aa757b0..3e4aeb32b 100644 --- a/configs/olive/sd/text_encoder.json +++ b/configs/olive/sd/text_encoder.json @@ -1 +1,111 @@ -{"input_model": {"type": "PyTorchModel", "config": {"model_path": "", "model_loader": "text_encoder_load", "model_script": "modules/olive.py", "io_config": {"input_names": ["input_ids"], "output_names": ["last_hidden_state", "pooler_output"], "dynamic_axes": {"input_ids": {"0": "batch", "1": "sequence"}}}, "dummy_inputs_func": "text_encoder_conversion_inputs"}}, "systems": {"local_system": {"type": "LocalSystem", "config": {"accelerators": ["gpu"]}}}, "evaluators": {"common_evaluator": {"metrics": [{"name": "latency", "type": "latency", "sub_types": [{"name": "avg"}], "user_config": {"user_script": "modules/olive.py", "dataloader_func": "text_encoder_data_loader", "batch_size": 1}}]}}, "passes": {"optimize_DmlExecutionProvider": {"type": "OrtTransformersOptimization", "disable_search": true, "config": {"model_type": "clip", "opt_level": 0, "float16": true, "use_gpu": true, "keep_io_types": false, "optimization_options": {"enable_gelu": true, "enable_layer_norm": true, "enable_attention": true, "use_multi_head_attention": true, "enable_skip_layer_norm": false, "enable_embed_layer_norm": true, "enable_bias_skip_layer_norm": false, "enable_bias_gelu": true, "enable_gelu_approximation": false, "enable_qordered_matmul": false, "enable_shape_inference": true, "enable_gemm_fast_gelu": false, "enable_nhwc_conv": false, "enable_group_norm": true, "enable_bias_splitgelu": false, "enable_packed_qkv": true, "enable_packed_kv": true, "enable_bias_add": false, "group_norm_channels_last": false}, "force_fp32_ops": ["RandomNormalLike"], "force_fp16_inputs": {"GroupNorm": [0, 1, 2]}}}, "optimize_CUDAExecutionProvider": {"type": "OrtTransformersOptimization", "disable_search": true, "config": {"model_type": "clip", "opt_level": 0, "float16": true, "use_gpu": true, "keep_io_types": false}}, "optimize_ROCMExecutionProvider": {"type": "OrtTransformersOptimization", "disable_search": true, "config": {"model_type": "clip", "opt_level": 0, "float16": true, "use_gpu": true, "keep_io_types": false}}}, "pass_flows": [[]], "engine": {"search_strategy": {"execution_order": "joint", "search_algorithm": "exhaustive"}, "evaluator": "common_evaluator", "evaluate_input_model": false, "host": "local_system", "target": "local_system", "cache_dir": "cache", "output_name": "text_encoder", "output_dir": "footprints", "execution_providers": ["DmlExecutionProvider"]}} \ No newline at end of file +{ + "input_model": { + "type": "PyTorchModel", + "config": { + "model_path": "", + "model_loader": "text_encoder_load", + "model_script": "modules/olive.py", + "io_config": { + "input_names": ["input_ids"], + "output_names": ["last_hidden_state", "pooler_output"], + "dynamic_axes": { "input_ids": { "0": "batch", "1": "sequence" } } + }, + "dummy_inputs_func": "text_encoder_conversion_inputs" + } + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "config": { "accelerators": ["gpu"] } + } + }, + "evaluators": { + "common_evaluator": { + "metrics": [ + { + "name": "latency", + "type": "latency", + "sub_types": [{ "name": "avg" }], + "user_config": { + "user_script": "modules/olive.py", + "dataloader_func": "text_encoder_data_loader", + "batch_size": 1 + } + } + ] + } + }, + "passes": { + "optimize_DmlExecutionProvider": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "clip", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false, + "optimization_options": { + "enable_gelu": true, + "enable_layer_norm": true, + "enable_attention": true, + "use_multi_head_attention": true, + "enable_skip_layer_norm": false, + "enable_embed_layer_norm": true, + "enable_bias_skip_layer_norm": false, + "enable_bias_gelu": true, + "enable_gelu_approximation": false, + "enable_qordered_matmul": false, + "enable_shape_inference": true, + "enable_gemm_fast_gelu": false, + "enable_nhwc_conv": false, + "enable_group_norm": true, + "enable_bias_splitgelu": false, + "enable_packed_qkv": true, + "enable_packed_kv": true, + "enable_bias_add": false, + "group_norm_channels_last": false + }, + "force_fp32_ops": ["RandomNormalLike"], + "force_fp16_inputs": { "GroupNorm": [0, 1, 2] } + } + }, + "optimize_CUDAExecutionProvider": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "clip", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false + } + }, + "optimize_ROCMExecutionProvider": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "clip", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false + } + } + }, + "pass_flows": [["optimize"]], + "engine": { + "search_strategy": { + "execution_order": "joint", + "search_algorithm": "exhaustive" + }, + "evaluator": "common_evaluator", + "evaluate_input_model": false, + "host": "local_system", + "target": "local_system", + "cache_dir": "cache", + "output_name": "text_encoder", + "output_dir": "footprints", + "execution_providers": ["DmlExecutionProvider"] + } +} diff --git a/configs/olive/sd/unet.json b/configs/olive/sd/unet.json index 5b02f5535..1b02ce0c3 100644 --- a/configs/olive/sd/unet.json +++ b/configs/olive/sd/unet.json @@ -1 +1,128 @@ -{"input_model": {"type": "PyTorchModel", "config": {"model_path": "", "model_loader": "unet_load", "model_script": "modules/olive.py", "io_config": {"input_names": ["sample", "timestep", "encoder_hidden_states", "return_dict"], "output_names": ["out_sample"], "dynamic_axes": {"sample": {"0": "unet_sample_batch", "1": "unet_sample_channels", "2": "unet_sample_height", "3": "unet_sample_width"}, "timestep": {"0": "unet_time_batch"}, "encoder_hidden_states": {"0": "unet_hidden_batch", "1": "unet_hidden_sequence"}}}, "dummy_inputs_func": "unet_conversion_inputs"}}, "systems": {"local_system": {"type": "LocalSystem", "config": {"accelerators": ["gpu"]}}}, "evaluators": {"common_evaluator": {"metrics": [{"name": "latency", "type": "latency", "sub_types": [{"name": "avg"}], "user_config": {"user_script": "modules/olive.py", "dataloader_func": "unet_data_loader", "batch_size": 2}}]}}, "passes": {"optimize_DmlExecutionProvider": {"type": "OrtTransformersOptimization", "disable_search": true, "config": {"model_type": "unet", "opt_level": 0, "float16": true, "use_gpu": true, "keep_io_types": false, "optimization_options": {"enable_gelu": true, "enable_layer_norm": true, "enable_attention": true, "use_multi_head_attention": true, "enable_skip_layer_norm": false, "enable_embed_layer_norm": true, "enable_bias_skip_layer_norm": false, "enable_bias_gelu": true, "enable_gelu_approximation": false, "enable_qordered_matmul": false, "enable_shape_inference": true, "enable_gemm_fast_gelu": false, "enable_nhwc_conv": false, "enable_group_norm": true, "enable_bias_splitgelu": false, "enable_packed_qkv": true, "enable_packed_kv": true, "enable_bias_add": false, "group_norm_channels_last": false}, "force_fp32_ops": ["RandomNormalLike"], "force_fp16_inputs": {"GroupNorm": [0, 1, 2]}}}, "optimize_CUDAExecutionProvider": {"type": "OrtTransformersOptimization", "disable_search": true, "config": {"model_type": "unet", "opt_level": 0, "float16": true, "use_gpu": true, "keep_io_types": false}}, "optimize_ROCMExecutionProvider": {"type": "OrtTransformersOptimization", "disable_search": true, "config": {"model_type": "unet", "opt_level": 0, "float16": true, "use_gpu": true, "keep_io_types": false}}}, "pass_flows": [[]], "engine": {"search_strategy": {"execution_order": "joint", "search_algorithm": "exhaustive"}, "evaluator": "common_evaluator", "evaluate_input_model": false, "host": "local_system", "target": "local_system", "cache_dir": "cache", "output_name": "unet", "output_dir": "footprints", "execution_providers": ["DmlExecutionProvider"]}} \ No newline at end of file +{ + "input_model": { + "type": "PyTorchModel", + "config": { + "model_path": "", + "model_loader": "unet_load", + "model_script": "modules/olive.py", + "io_config": { + "input_names": [ + "sample", + "timestep", + "encoder_hidden_states", + "return_dict" + ], + "output_names": ["out_sample"], + "dynamic_axes": { + "sample": { + "0": "unet_sample_batch", + "1": "unet_sample_channels", + "2": "unet_sample_height", + "3": "unet_sample_width" + }, + "timestep": { "0": "unet_time_batch" }, + "encoder_hidden_states": { + "0": "unet_hidden_batch", + "1": "unet_hidden_sequence" + } + } + }, + "dummy_inputs_func": "unet_conversion_inputs" + } + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "config": { "accelerators": ["gpu"] } + } + }, + "evaluators": { + "common_evaluator": { + "metrics": [ + { + "name": "latency", + "type": "latency", + "sub_types": [{ "name": "avg" }], + "user_config": { + "user_script": "modules/olive.py", + "dataloader_func": "unet_data_loader", + "batch_size": 2 + } + } + ] + } + }, + "passes": { + "optimize_DmlExecutionProvider": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "unet", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false, + "optimization_options": { + "enable_gelu": true, + "enable_layer_norm": true, + "enable_attention": true, + "use_multi_head_attention": true, + "enable_skip_layer_norm": false, + "enable_embed_layer_norm": true, + "enable_bias_skip_layer_norm": false, + "enable_bias_gelu": true, + "enable_gelu_approximation": false, + "enable_qordered_matmul": false, + "enable_shape_inference": true, + "enable_gemm_fast_gelu": false, + "enable_nhwc_conv": false, + "enable_group_norm": true, + "enable_bias_splitgelu": false, + "enable_packed_qkv": true, + "enable_packed_kv": true, + "enable_bias_add": false, + "group_norm_channels_last": false + }, + "force_fp32_ops": ["RandomNormalLike"], + "force_fp16_inputs": { "GroupNorm": [0, 1, 2] } + } + }, + "optimize_CUDAExecutionProvider": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "unet", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false + } + }, + "optimize_ROCMExecutionProvider": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "unet", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false + } + } + }, + "pass_flows": [["optimize"]], + "engine": { + "search_strategy": { + "execution_order": "joint", + "search_algorithm": "exhaustive" + }, + "evaluator": "common_evaluator", + "evaluate_input_model": false, + "host": "local_system", + "target": "local_system", + "cache_dir": "cache", + "output_name": "unet", + "output_dir": "footprints", + "execution_providers": ["DmlExecutionProvider"] + } +} diff --git a/configs/olive/sd/vae_decoder.json b/configs/olive/sd/vae_decoder.json index 0601f40ad..3653ad422 100644 --- a/configs/olive/sd/vae_decoder.json +++ b/configs/olive/sd/vae_decoder.json @@ -1 +1,118 @@ -{"input_model": {"type": "PyTorchModel", "config": {"model_path": "", "model_loader": "vae_decoder_load", "model_script": "modules/olive.py", "io_config": {"input_names": ["latent_sample", "return_dict"], "output_names": ["sample"], "dynamic_axes": {"latent_sample": {"0": "batch", "1": "channels", "2": "height", "3": "width"}}}, "dummy_inputs_func": "vae_decoder_conversion_inputs"}}, "systems": {"local_system": {"type": "LocalSystem", "config": {"accelerators": ["gpu"]}}}, "evaluators": {"common_evaluator": {"metrics": [{"name": "latency", "type": "latency", "sub_types": [{"name": "avg"}], "user_config": {"user_script": "modules/olive.py", "dataloader_func": "vae_decoder_data_loader", "batch_size": 1}}]}}, "passes": {"optimize_DmlExecutionProvider": {"type": "OrtTransformersOptimization", "disable_search": true, "config": {"model_type": "vae", "opt_level": 0, "float16": true, "use_gpu": true, "keep_io_types": false, "optimization_options": {"enable_gelu": true, "enable_layer_norm": true, "enable_attention": true, "use_multi_head_attention": true, "enable_skip_layer_norm": false, "enable_embed_layer_norm": true, "enable_bias_skip_layer_norm": false, "enable_bias_gelu": true, "enable_gelu_approximation": false, "enable_qordered_matmul": false, "enable_shape_inference": true, "enable_gemm_fast_gelu": false, "enable_nhwc_conv": false, "enable_group_norm": true, "enable_bias_splitgelu": false, "enable_packed_qkv": true, "enable_packed_kv": true, "enable_bias_add": false, "group_norm_channels_last": false}, "force_fp32_ops": ["RandomNormalLike"], "force_fp16_inputs": {"GroupNorm": [0, 1, 2]}}}, "optimize_CUDAExecutionProvider": {"type": "OrtTransformersOptimization", "disable_search": true, "config": {"model_type": "vae", "opt_level": 0, "float16": true, "use_gpu": true, "keep_io_types": false}}, "optimize_ROCMExecutionProvider": {"type": "OrtTransformersOptimization", "disable_search": true, "config": {"model_type": "vae", "opt_level": 0, "float16": true, "use_gpu": true, "keep_io_types": false}}}, "pass_flows": [[]], "engine": {"search_strategy": {"execution_order": "joint", "search_algorithm": "exhaustive"}, "evaluator": "common_evaluator", "evaluate_input_model": false, "host": "local_system", "target": "local_system", "cache_dir": "cache", "output_name": "vae_decoder", "output_dir": "footprints", "execution_providers": ["DmlExecutionProvider"]}} \ No newline at end of file +{ + "input_model": { + "type": "PyTorchModel", + "config": { + "model_path": "", + "model_loader": "vae_decoder_load", + "model_script": "modules/olive.py", + "io_config": { + "input_names": ["latent_sample", "return_dict"], + "output_names": ["sample"], + "dynamic_axes": { + "latent_sample": { + "0": "batch", + "1": "channels", + "2": "height", + "3": "width" + } + } + }, + "dummy_inputs_func": "vae_decoder_conversion_inputs" + } + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "config": { "accelerators": ["gpu"] } + } + }, + "evaluators": { + "common_evaluator": { + "metrics": [ + { + "name": "latency", + "type": "latency", + "sub_types": [{ "name": "avg" }], + "user_config": { + "user_script": "modules/olive.py", + "dataloader_func": "vae_decoder_data_loader", + "batch_size": 1 + } + } + ] + } + }, + "passes": { + "optimize_DmlExecutionProvider": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "vae", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false, + "optimization_options": { + "enable_gelu": true, + "enable_layer_norm": true, + "enable_attention": true, + "use_multi_head_attention": true, + "enable_skip_layer_norm": false, + "enable_embed_layer_norm": true, + "enable_bias_skip_layer_norm": false, + "enable_bias_gelu": true, + "enable_gelu_approximation": false, + "enable_qordered_matmul": false, + "enable_shape_inference": true, + "enable_gemm_fast_gelu": false, + "enable_nhwc_conv": false, + "enable_group_norm": true, + "enable_bias_splitgelu": false, + "enable_packed_qkv": true, + "enable_packed_kv": true, + "enable_bias_add": false, + "group_norm_channels_last": false + }, + "force_fp32_ops": ["RandomNormalLike"], + "force_fp16_inputs": { "GroupNorm": [0, 1, 2] } + } + }, + "optimize_CUDAExecutionProvider": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "vae", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false + } + }, + "optimize_ROCMExecutionProvider": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "vae", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false + } + } + }, + "pass_flows": [["optimize"]], + "engine": { + "search_strategy": { + "execution_order": "joint", + "search_algorithm": "exhaustive" + }, + "evaluator": "common_evaluator", + "evaluate_input_model": false, + "host": "local_system", + "target": "local_system", + "cache_dir": "cache", + "output_name": "vae_decoder", + "output_dir": "footprints", + "execution_providers": ["DmlExecutionProvider"] + } +} diff --git a/configs/olive/sd/vae_encoder.json b/configs/olive/sd/vae_encoder.json index 30ff18bd0..52dc369f9 100644 --- a/configs/olive/sd/vae_encoder.json +++ b/configs/olive/sd/vae_encoder.json @@ -1 +1,118 @@ -{"input_model": {"type": "PyTorchModel", "config": {"model_path": "", "model_loader": "vae_encoder_load", "model_script": "modules/olive.py", "io_config": {"input_names": ["sample", "return_dict"], "output_names": ["latent_sample"], "dynamic_axes": {"sample": {"0": "batch", "1": "channels", "2": "height", "3": "width"}}}, "dummy_inputs_func": "vae_encoder_conversion_inputs"}}, "systems": {"local_system": {"type": "LocalSystem", "config": {"accelerators": ["gpu"]}}}, "evaluators": {"common_evaluator": {"metrics": [{"name": "latency", "type": "latency", "sub_types": [{"name": "avg"}], "user_config": {"user_script": "modules/olive.py", "dataloader_func": "vae_encoder_data_loader", "batch_size": 1}}]}}, "passes": {"optimize_DmlExecutionProvider": {"type": "OrtTransformersOptimization", "disable_search": true, "config": {"model_type": "vae", "opt_level": 0, "float16": true, "use_gpu": true, "keep_io_types": false, "optimization_options": {"enable_gelu": true, "enable_layer_norm": true, "enable_attention": true, "use_multi_head_attention": true, "enable_skip_layer_norm": false, "enable_embed_layer_norm": true, "enable_bias_skip_layer_norm": false, "enable_bias_gelu": true, "enable_gelu_approximation": false, "enable_qordered_matmul": false, "enable_shape_inference": true, "enable_gemm_fast_gelu": false, "enable_nhwc_conv": false, "enable_group_norm": true, "enable_bias_splitgelu": false, "enable_packed_qkv": true, "enable_packed_kv": true, "enable_bias_add": false, "group_norm_channels_last": false}, "force_fp32_ops": ["RandomNormalLike"], "force_fp16_inputs": {"GroupNorm": [0, 1, 2]}}}, "optimize_CUDAExecutionProvider": {"type": "OrtTransformersOptimization", "disable_search": true, "config": {"model_type": "vae", "opt_level": 0, "float16": true, "use_gpu": true, "keep_io_types": false}}, "optimize_ROCMExecutionProvider": {"type": "OrtTransformersOptimization", "disable_search": true, "config": {"model_type": "vae", "opt_level": 0, "float16": true, "use_gpu": true, "keep_io_types": false}}}, "pass_flows": [[]], "engine": {"search_strategy": {"execution_order": "joint", "search_algorithm": "exhaustive"}, "evaluator": "common_evaluator", "evaluate_input_model": false, "host": "local_system", "target": "local_system", "cache_dir": "cache", "output_name": "vae_encoder", "output_dir": "footprints", "execution_providers": ["DmlExecutionProvider"]}} \ No newline at end of file +{ + "input_model": { + "type": "PyTorchModel", + "config": { + "model_path": "", + "model_loader": "vae_encoder_load", + "model_script": "modules/olive.py", + "io_config": { + "input_names": ["sample", "return_dict"], + "output_names": ["latent_sample"], + "dynamic_axes": { + "sample": { + "0": "batch", + "1": "channels", + "2": "height", + "3": "width" + } + } + }, + "dummy_inputs_func": "vae_encoder_conversion_inputs" + } + }, + "systems": { + "local_system": { + "type": "LocalSystem", + "config": { "accelerators": ["gpu"] } + } + }, + "evaluators": { + "common_evaluator": { + "metrics": [ + { + "name": "latency", + "type": "latency", + "sub_types": [{ "name": "avg" }], + "user_config": { + "user_script": "modules/olive.py", + "dataloader_func": "vae_encoder_data_loader", + "batch_size": 1 + } + } + ] + } + }, + "passes": { + "optimize_DmlExecutionProvider": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "vae", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false, + "optimization_options": { + "enable_gelu": true, + "enable_layer_norm": true, + "enable_attention": true, + "use_multi_head_attention": true, + "enable_skip_layer_norm": false, + "enable_embed_layer_norm": true, + "enable_bias_skip_layer_norm": false, + "enable_bias_gelu": true, + "enable_gelu_approximation": false, + "enable_qordered_matmul": false, + "enable_shape_inference": true, + "enable_gemm_fast_gelu": false, + "enable_nhwc_conv": false, + "enable_group_norm": true, + "enable_bias_splitgelu": false, + "enable_packed_qkv": true, + "enable_packed_kv": true, + "enable_bias_add": false, + "group_norm_channels_last": false + }, + "force_fp32_ops": ["RandomNormalLike"], + "force_fp16_inputs": { "GroupNorm": [0, 1, 2] } + } + }, + "optimize_CUDAExecutionProvider": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "vae", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false + } + }, + "optimize_ROCMExecutionProvider": { + "type": "OrtTransformersOptimization", + "disable_search": true, + "config": { + "model_type": "vae", + "opt_level": 0, + "float16": true, + "use_gpu": true, + "keep_io_types": false + } + } + }, + "pass_flows": [["optimize"]], + "engine": { + "search_strategy": { + "execution_order": "joint", + "search_algorithm": "exhaustive" + }, + "evaluator": "common_evaluator", + "evaluate_input_model": false, + "host": "local_system", + "target": "local_system", + "cache_dir": "cache", + "output_name": "vae_encoder", + "output_dir": "footprints", + "execution_providers": ["DmlExecutionProvider"] + } +} diff --git a/configs/olive/sdxl/text_encoder.json b/configs/olive/sdxl/text_encoder.json index e8270e7df..e1fc1d672 100644 --- a/configs/olive/sdxl/text_encoder.json +++ b/configs/olive/sdxl/text_encoder.json @@ -130,7 +130,7 @@ } } }, - "pass_flows": [[]], + "pass_flows": [["optimize"]], "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 221d51605..479c48e4a 100644 --- a/configs/olive/sdxl/text_encoder_2.json +++ b/configs/olive/sdxl/text_encoder_2.json @@ -170,7 +170,7 @@ } } }, - "pass_flows": [[]], + "pass_flows": [["optimize"]], "engine": { "search_strategy": { "execution_order": "joint", diff --git a/configs/olive/sdxl/unet.json b/configs/olive/sdxl/unet.json index 82984a7e6..f5baf75f3 100644 --- a/configs/olive/sdxl/unet.json +++ b/configs/olive/sdxl/unet.json @@ -120,7 +120,7 @@ } } }, - "pass_flows": [[]], + "pass_flows": [["optimize"]], "engine": { "search_strategy": { "execution_order": "joint", diff --git a/configs/olive/sdxl/vae_decoder.json b/configs/olive/sdxl/vae_decoder.json index 33b2e549d..d5409a511 100644 --- a/configs/olive/sdxl/vae_decoder.json +++ b/configs/olive/sdxl/vae_decoder.json @@ -130,7 +130,7 @@ } } }, - "pass_flows": [[]], + "pass_flows": [["optimize"]], "engine": { "search_strategy": { "execution_order": "joint", diff --git a/configs/olive/sdxl/vae_encoder.json b/configs/olive/sdxl/vae_encoder.json index fd74a5c98..3e871d524 100644 --- a/configs/olive/sdxl/vae_encoder.json +++ b/configs/olive/sdxl/vae_encoder.json @@ -110,7 +110,7 @@ } } }, - "pass_flows": [[]], + "pass_flows": [["optimize"]], "engine": { "search_strategy": { "execution_order": "joint", diff --git a/modules/olive.py b/modules/olive.py index 95c52c100..41c46dce1 100644 --- a/modules/olive.py +++ b/modules/olive.py @@ -9,10 +9,12 @@ class ENVStore: __DESERIALIZER: Dict[Type, Callable[[str,], Any]] = { bool: lambda x: bool(int(x)), int: int, + str: lambda x: x, } __SERIALIZER: Dict[Type, Callable[[Any,], str]] = { bool: lambda x: str(int(x)), int: str, + str: lambda x: x, } def __getattr__(self, name: str): @@ -38,10 +40,12 @@ class ENVStore: class OliveOptimizerConfig(ENVStore): from_huggingface_cache: bool - use_fp16_fixed_vae: bool is_sdxl: bool + vae_id: str + vae_subfolder: str + width: int height: int batch_size: int @@ -220,9 +224,11 @@ def vae_encoder_inputs(_, torch_dtype): def vae_encoder_load(model_name): - if config.use_fp16_fixed_vae: - model_name = "madebyollin/sdxl-vae-fp16-fix" - model = diffusers.AutoencoderKL.from_pretrained(model_name, subfolder="vae_encoder" if os.path.isdir(os.path.join(model_name, "vae_encoder")) else "vae", **get_loader_arguments()) + subfolder = "vae_encoder" if os.path.isdir(os.path.join(model_name, "vae_encoder")) else "vae" + if config.vae_id is not None: + model_name = config.vae_id + subfolder = config.vae_subfolder + model = diffusers.AutoencoderKL.from_pretrained(model_name, subfolder=subfolder, **get_loader_arguments()) model.forward = lambda sample, return_dict: model.encode(sample, return_dict)[0].sample() return model @@ -248,9 +254,11 @@ def vae_decoder_inputs(_, torch_dtype): def vae_decoder_load(model_name): - if config.use_fp16_fixed_vae: - model_name = "madebyollin/sdxl-vae-fp16-fix" - model = diffusers.AutoencoderKL.from_pretrained(model_name, subfolder="vae_decoder" if os.path.isdir(os.path.join(model_name, "vae_decoder")) else "vae", **get_loader_arguments()) + subfolder = "vae_decoder" if os.path.isdir(os.path.join(model_name, "vae_decoder")) else "vae" + if config.vae_id is not None: + model_name = config.vae_id + subfolder = config.vae_subfolder + model = diffusers.AutoencoderKL.from_pretrained(model_name, subfolder=subfolder, **get_loader_arguments()) model.forward = model.decode return model diff --git a/modules/onnx_ep.py b/modules/onnx_ep.py index 9028c6a33..148d89840 100644 --- a/modules/onnx_ep.py +++ b/modules/onnx_ep.py @@ -74,7 +74,6 @@ def get_provider() -> Tuple: def install_execution_provider(ep: ExecutionProvider): from installer import pip, uninstall, installed - from modules.shared import log if installed("onnxruntime"): uninstall("onnxruntime") diff --git a/modules/onnx_pipelines.py b/modules/onnx_pipelines.py index 3339aae49..ce48fea9d 100644 --- a/modules/onnx_pipelines.py +++ b/modules/onnx_pipelines.py @@ -19,7 +19,6 @@ from installer import log from modules import shared from modules.paths import sd_configs_path from modules.sd_models import CheckpointInfo -from modules.sd_models_compile import CompiledModelState from modules.processing import StableDiffusionProcessing from modules.olive import config from modules.onnx import OnnxFakeModule, submodels_sd, submodels_sdxl, submodels_sdxl_refiner @@ -240,7 +239,10 @@ class OnnxRawPipeline(OnnxPipelineBase): 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["pass_flows"] = [[pass_key]] + for flow in olive_config["pass_flows"]: + for i in range(len(flow)): + if flow[i] == "optimize": + flow[i] = pass_key olive_config["input_model"]["config"]["model_path"] = os.path.abspath(os.path.join(in_dir, submodel, "model.onnx")) olive_config["passes"][pass_key]["config"]["float16"] = shared.opts.olive_float16 olive_config["engine"]["execution_providers"] = [shared.opts.onnx_execution_provider] @@ -327,7 +329,10 @@ class OnnxRawPipeline(OnnxPipelineBase): disable_classifier_free_guidance = p.cfg_scale < 0.01 config.from_huggingface_cache = self.from_huggingface_cache - config.use_fp16_fixed_vae = self._is_sdxl and not shared.opts.diffusers_vae_upcast + if self._is_sdxl and not shared.opts.diffusers_vae_upcast: + log.info("ONNX: VAE override set: id=madebyollin/sdxl-vae-fp16-fix, subfolder=") + config.vae_id = "madebyollin/sdxl-vae-fp16-fix" + config.vae_subfolder = "" config.is_sdxl = self._is_sdxl diff --git a/modules/ui_onnx.py b/modules/ui_onnx.py index e5603e807..2ea6b16ed 100644 --- a/modules/ui_onnx.py +++ b/modules/ui_onnx.py @@ -2,7 +2,6 @@ import os import json from typing import Dict, List, Union import gradio as gr -from olive.passes import REGISTRY def get_recursively(d: Union[Dict, List], *args): @@ -16,11 +15,14 @@ def create_ui(): from modules.shared import log, opts, cmd_opts from modules.paths import sd_configs_path from modules.onnx_ep import ExecutionProvider, install_execution_provider + from modules.olive import config as olive_config with gr.Blocks(analytics_enabled=False) as ui: with gr.Row(): with gr.Tabs(elem_id="tabs_onnx"): with gr.TabItem("Manage execution providers", id="onnxep"): + gr.Markdown("Uninstall existing execution provider and install another one.") + choices = [] for ep in ExecutionProvider: @@ -38,11 +40,32 @@ def create_ui(): ep_checkbox = gr.Radio(label="Execution provider", value=ep_default, choices=choices) ep_install = gr.Button(value="Install") - gr.Text("Warning! If you are trying to reinstall, it may not work due to permission issue.") + gr.Markdown("**Warning! If you are trying to reinstall, it may not work due to permission issue.**") ep_install.click(fn=install_execution_provider, inputs=ep_checkbox) + with gr.TabItem("Override VAE", id="force_vae"): + gr.Markdown("Ignore baked-in vae and replace it with what you want.") + + onnx_vae_id = gr.Textbox(label="Huggingface VAE ID", info="Leave empty for default (baked-in vae).", value="") + onnx_vae_subfolder = gr.Textbox(label="VAE subfolder", info="Leave empty for root. Default: vae", value="vae") + onnx_vae_apply_button = gr.Button(value="Apply") + + def onnx_vae_apply(id: str, subfolder: str): + olive_config.vae_id = id + olive_config.vae_subfolder = subfolder + if id == "": + log.info("ONNX: VAE override unset.") + del olive_config.vae_id + olive_config.vae_subfolder = "vae" + else: + log.info(f"ONNX: VAE override set: id={id}, subfolder={subfolder}") + + onnx_vae_apply_button.click(fn=onnx_vae_apply, inputs=[onnx_vae_id, onnx_vae_subfolder,]) + if opts.cuda_compile_backend == "olive-ai": + from olive.passes import REGISTRY + with gr.Tabs(elem_id="tabs_olive"): with gr.TabItem("Customize pass flow", id="pass_flow"): with gr.Tabs(elem_id="tabs_model_type"):