vae override for onnx

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"]}}
+{
+  "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"]
+  }
+}

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"]}}
+{
+  "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"]
+  }
+}

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"]}}
+{
+  "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"]
+  }
+}

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"]}}
+{
+  "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"]
+  }
+}

configs/olive/sdxl/text_encoder.json

@@ -130,7 +130,7 @@
       }
     }
   },
-  "pass_flows": [[]],
+  "pass_flows": [["optimize"]],
   "engine": {
     "search_strategy": {
       "execution_order": "joint",

configs/olive/sdxl/text_encoder_2.json

@@ -170,7 +170,7 @@
       }
     }
   },
-  "pass_flows": [[]],
+  "pass_flows": [["optimize"]],
   "engine": {
     "search_strategy": {
       "execution_order": "joint",

configs/olive/sdxl/unet.json

@@ -120,7 +120,7 @@
       }
     }
   },
-  "pass_flows": [[]],
+  "pass_flows": [["optimize"]],
   "engine": {
     "search_strategy": {
       "execution_order": "joint",

configs/olive/sdxl/vae_decoder.json

@@ -130,7 +130,7 @@
       }
     }
   },
-  "pass_flows": [[]],
+  "pass_flows": [["optimize"]],
   "engine": {
     "search_strategy": {
       "execution_order": "joint",

configs/olive/sdxl/vae_encoder.json

@@ -110,7 +110,7 @@
       }
     }
   },
-  "pass_flows": [[]],
+  "pass_flows": [["optimize"]],
   "engine": {
     "search_strategy": {
       "execution_order": "joint",

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
 

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")

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
 

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"):