1580158
/
automatic
mirror of https://github.com/vladmandic/automatic
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

OnnxConversion pass (temp)

pull/2784/head
Seunghoon Lee 2024-01-25 21:07:01 +09:00
parent fc465f94bc
commit 8bcd60065d
No known key found for this signature in database
GPG Key ID: 436E38F4E70BD152
10 changed files with 21 additions and 760 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

63
configs/onnx/sd_text_encoder.json
View File

@ -1,63 +0,0 @@
{
"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": {
"convert": {
"type": "OnnxConversion",
"config": {
"target_opset": 14
}
}
},
"pass_flows": [["convert"]],
"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"]
}
}

83
configs/onnx/sd_unet.json
View File

@ -1,83 +0,0 @@
{
"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": {
"convert": {
"type": "OnnxConversion",
"config": {
"target_opset": 14,
"save_as_external_data": true,
"all_tensors_to_one_file": true,
"external_data_name": "weights.pb"
}
}
},
"pass_flows": [["convert"]],
"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"]
}
}

70
configs/onnx/sd_vae_decoder.json
View File

@ -1,70 +0,0 @@
{
"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": {
"convert": {
"type": "OnnxConversion",
"config": {
"target_opset": 14
}
}
},
"pass_flows": [["convert"]],
"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"]
}
}

70
configs/onnx/sd_vae_encoder.json
View File

@ -1,70 +0,0 @@
{
"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": {
"convert": {
"type": "OnnxConversion",
"config": {
"target_opset": 14
}
}
},
"pass_flows": [["convert"]],
"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"]
}
}

96
configs/onnx/sdxl_text_encoder.json
View File

@ -1,96 +0,0 @@
{
"input_model": {
"type": "PyTorchModel",
"config": {
"model_path": "",
"model_loader": "text_encoder_load",
"model_script": "modules/olive.py",
"io_config": {
"input_names": ["input_ids", "output_hidden_states"],
"output_names": [
"last_hidden_state",
"pooler_output",
"hidden_states.0",
"hidden_states.1",
"hidden_states.2",
"hidden_states.3",
"hidden_states.4",
"hidden_states.5",
"hidden_states.6",
"hidden_states.7",
"hidden_states.8",
"hidden_states.9",
"hidden_states.10",
"hidden_states.11",
"hidden_states.12"
],
"dynamic_axes": {
"input_ids": { "0": "batch_size", "1": "sequence_length" },
"last_hidden_state": { "0": "batch_size", "1": "sequence_length" },
"pooler_output": { "0": "batch_size" },
"hidden_states.0": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.1": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.2": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.3": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.4": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.5": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.6": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.7": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.8": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.9": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.10": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.11": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.12": { "0": "batch_size", "1": "sequence_length" }
}
},
"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": {
"convert": {
"type": "OnnxConversion",
"config": {
"target_opset": 14
}
}
},
"pass_flows": [["convert"]],
"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"]
}
}

136
configs/onnx/sdxl_text_encoder_2.json
View File

@ -1,136 +0,0 @@
{
"input_model": {
"type": "PyTorchModel",
"config": {
"model_path": "",
"model_loader": "text_encoder_2_load",
"model_script": "modules/olive.py",
"io_config": {
"input_names": ["input_ids", "output_hidden_states"],
"output_names": [
"text_embeds",
"last_hidden_state",
"hidden_states.0",
"hidden_states.1",
"hidden_states.2",
"hidden_states.3",
"hidden_states.4",
"hidden_states.5",
"hidden_states.6",
"hidden_states.7",
"hidden_states.8",
"hidden_states.9",
"hidden_states.10",
"hidden_states.11",
"hidden_states.12",
"hidden_states.13",
"hidden_states.14",
"hidden_states.15",
"hidden_states.16",
"hidden_states.17",
"hidden_states.18",
"hidden_states.19",
"hidden_states.20",
"hidden_states.21",
"hidden_states.22",
"hidden_states.23",
"hidden_states.24",
"hidden_states.25",
"hidden_states.26",
"hidden_states.27",
"hidden_states.28",
"hidden_states.29",
"hidden_states.30",
"hidden_states.31",
"hidden_states.32"
],
"dynamic_axes": {
"input_ids": { "0": "batch_size", "1": "sequence_length" },
"text_embeds": { "0": "batch_size", "1": "sequence_length" },
"last_hidden_state": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.0": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.1": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.2": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.3": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.4": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.5": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.6": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.7": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.8": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.9": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.10": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.11": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.12": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.13": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.14": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.15": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.16": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.17": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.18": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.19": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.20": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.21": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.22": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.23": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.24": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.25": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.26": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.27": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.28": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.29": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.30": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.31": { "0": "batch_size", "1": "sequence_length" },
"hidden_states.32": { "0": "batch_size", "1": "sequence_length" }
}
},
"dummy_inputs_func": "text_encoder_2_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_2_data_loader",
"batch_size": 1
}
}
]
}
},
"passes": {
"convert": {
"type": "OnnxConversion",
"config": {
"target_opset": 14
}
}
},
"pass_flows": [["convert"]],
"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_2",
"output_dir": "footprints",
"execution_providers": ["DmlExecutionProvider"]
}
}

89
configs/onnx/sdxl_unet.json
View File

@ -1,89 +0,0 @@
{
"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",
"text_embeds",
"time_ids"
],
"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"
},
"text_embeds": {
"0": "unet_text_embeds_batch",
"1": "unet_text_embeds_size"
},
"time_ids": { "0": "unet_time_ids_batch", "1": "unet_time_ids_size" }
}
},
"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": {
"convert": {
"type": "OnnxConversion",
"config": {
"target_opset": 14,
"save_as_external_data": true,
"all_tensors_to_one_file": true,
"external_data_name": "weights.pb"
}
}
},
"pass_flows": [["convert"]],
"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"]
}
}

76
configs/onnx/sdxl_vae_decoder.json
View File

@ -1,76 +0,0 @@
{
"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_size",
"1": "num_channels_latent",
"2": "height_latent",
"3": "width_latent"
},
"sample": {
"0": "batch_size",
"1": "num_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": {
"convert": {
"type": "OnnxConversion",
"config": {
"target_opset": 14
}
}
},
"pass_flows": [["convert"]],
"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"]
}
}

76
configs/onnx/sdxl_vae_encoder.json
View File

@ -1,76 +0,0 @@
{
"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_size",
"1": "num_channels",
"2": "height",
"3": "width"
},
"latent_sample": {
"0": "batch_size",
"1": "num_channels_latent",
"2": "height_latent",
"3": "width_latent"
}
}
},
"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": {
"convert": {
"type": "OnnxConversion",
"config": {
"target_opset": 14
}
}
},
"pass_flows": [["convert"]],
"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"]
}
}

22
modules/onnx_pipelines.py
View File

@ -28,6 +28,22 @@ SUBMODELS_SD = ("text_encoder", "unet", "vae_encoder", "vae_decoder",)
SUBMODELS_SDXL = ("text_encoder", "text_encoder_2", "unet", "vae_encoder", "vae_decoder",)
SUBMODELS_SDXL_REFINER = ("text_encoder_2", "unet", "vae_encoder", "vae_decoder",)
CONVERSION_PASS = {
"type": "OnnxConversion",
"config": {
"target_opset": 14,
},
}
CONVERSION_PASS_UNET = {
"type": "OnnxConversion",
"config": {
"target_opset": 14,
"save_as_external_data": True,
"all_tensors_to_one_file": True,
"external_data_name": "weights.pb",
},
}
class OnnxPipelineBase(OnnxFakeModule, diffusers.DiffusionPipeline, metaclass=ABCMeta):
model_type: str
@ -167,9 +183,13 @@ class OnnxRawPipeline(OnnxPipelineBase):
for submodel in submodels:
log.info(f"\nConverting {submodel}")
with open(os.path.join(sd_configs_path, "onnx", f"{'sdxl' if self._is_sdxl else 'sd'}_{submodel}.json"), "r") as config_file:
with open(os.path.join(sd_configs_path, "olive", 'sdxl' if self._is_sdxl else 'sd', f"{submodel}.json"), "r") as config_file:
conversion_config = json.load(config_file)
conversion_config["input_model"]["config"]["model_path"] = os.path.abspath(in_dir)
conversion_config["passes"] = {
"_conversion": CONVERSION_PASS_UNET if submodel == "unet" else CONVERSION_PASS,
}
conversion_config["pass_flows"] = [["_conversion"]]
conversion_config["engine"]["execution_providers"] = [shared.opts.onnx_execution_provider]
run(conversion_config)