KeyError: 'unet_dataloader' occurs when optimizing unet in stable_diffusion_xl.py

[giocafe]

Describe the bug
Unable to conversion to onnx and latency optimization.
KeyError: 'unet_dataloader' occurs when optimizing unet in stable_diffusion_xl.py.

To Reproduce

• Pre-work (vae_encoder_gpu-dml_footprints.json file not found when converting stable diffusion xl base model #1202)
  Remove the below statement in following files: 1) config_vae_encoder.json, 2) config_vae_decoder.json, 3)config_unet.json
  "data_configs": [ { "name": "latency_data_config", "user_script": "user_script.py", "load_dataset_config": { "type": "simple_dataset" }, "dataloader_config": { "type": "vae_encoder_dataloader/vae_decoder_dataloader/unet_dataloader", "batch_size": 2 } } ], : "evaluator": "common_evaluator",

• Commands
  $ cd Olive/examples/directml/stable_diffusion_xl
  $ python stable_diffusion_xl.py --provider cuda/dml --model_id=stabilityai/stable-diffusion-xl-base-1.0 --optimize

Expected behavior
Enable to conversion to onnx and latency optimization.

Olive config
Add Olive configurations here.

Olive logs (--provider cuda)
(sdxl_cuda) C:\Users\gioca\Olive\examples\directml\stable_diffusion_xl>python stable_diffusion_xl.py --provider cuda --model_id=stabilityai/stable-diffusion-xl-base-1.0 --optimize
Download stable diffusion PyTorch pipeline...
Loading pipeline components...: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 11.18it/s]

Optimizing vae_encoder
[2024-08-21 15:17:52,280] [INFO] [run.py:138:run_engine] Running workflow default_workflow
[2024-08-21 15:17:52,280] [INFO] [cache.py:51:init] Using cache directory: C:\Users\gioca\Olive\examples\directml\stable_diffusion_xl\cache\default_workflow
[2024-08-21 15:17:52,280] [INFO] [engine.py:1012:save_olive_config] Saved Olive config to C:\Users\gioca\Olive\examples\directml\stable_diffusion_xl\cache\default_workflow\olive_config.json
[2024-08-21 15:17:52,280] [INFO] [accelerator_creator.py:224:create_accelerators] Running workflow on accelerator specs: gpu-cuda
[2024-08-21 15:17:52,280] [INFO] [engine.py:275:run] Running Olive on accelerator: gpu-cuda
[2024-08-21 15:17:52,280] [INFO] [engine.py:1109:_create_system] Creating target system ...
[2024-08-21 15:17:52,280] [INFO] [engine.py:1112:_create_system] Target system created in 0.000000 seconds
[2024-08-21 15:17:52,280] [INFO] [engine.py:1121:_create_system] Creating host system ...
[2024-08-21 15:17:52,280] [INFO] [engine.py:1124:_create_system] Host system created in 0.000000 seconds
[2024-08-21 15:17:52,297] [INFO] [engine.py:877:_run_pass] Running pass convert:OnnxConversion
[2024-08-21 15:17:52,297] [INFO] [engine.py:914:_run_pass] Loaded model from cache: 0_OnnxConversion-c79f256b-076cfb73
[2024-08-21 15:17:52,297] [INFO] [engine.py:877:_run_pass] Running pass optimize_cuda:OrtTransformersOptimization
[2024-08-21 15:17:52,297] [INFO] [engine.py:914:_run_pass] Loaded model from cache: 1_OrtTransformersOptimization-0-891c8d58-gpu-cuda
[2024-08-21 15:17:52,327] [INFO] [engine.py:374:run_accelerator] Save footprint to footprints\vae_encoder_gpu-cuda_footprints.json.
[2024-08-21 15:17:52,327] [INFO] [engine.py:292:run] Run history for gpu-cuda:
[2024-08-21 15:17:52,327] [INFO] [engine.py:587:dump_run_history] run history:
+---------------------------------------------------+------------------------------------+-----------------------------+----------------+-----------+
| model_id | parent_model_id | from_pass | duration_sec | metrics |
+===================================================+====================================+=============================+================+===========+
| c79f256b | | | | |
+---------------------------------------------------+------------------------------------+-----------------------------+----------------+-----------+
| 0_OnnxConversion-c79f256b-076cfb73 | c79f256b | OnnxConversion | 17.777 | |
+---------------------------------------------------+------------------------------------+-----------------------------+----------------+-----------+
| 1_OrtTransformersOptimization-0-891c8d58-gpu-cuda | 0_OnnxConversion-c79f256b-076cfb73 | OrtTransformersOptimization | 3.51828 | |
+---------------------------------------------------+------------------------------------+-----------------------------+----------------+-----------+
[2024-08-21 15:17:52,327] [INFO] [engine.py:307:run] No packaging config provided, skip packaging artifacts
Unoptimized Model : C:\Users\gioca\Olive\examples\directml\stable_diffusion_xl\cache\default_workflow\models\0_OnnxConversion-c79f256b-076cfb73\output_model\model.onnx
Optimized Model : C:\Users\gioca\Olive\examples\directml\stable_diffusion_xl\cache\default_workflow\models\1_OrtTransformersOptimization-0-891c8d58-gpu-cuda\output_model\model.onnx

Optimizing vae_decoder
[2024-08-21 15:17:52,343] [INFO] [run.py:138:run_engine] Running workflow default_workflow
[2024-08-21 15:17:52,343] [INFO] [cache.py:51:init] Using cache directory: C:\Users\gioca\Olive\examples\directml\stable_diffusion_xl\cache\default_workflow
[2024-08-21 15:17:52,349] [INFO] [engine.py:1012:save_olive_config] Saved Olive config to C:\Users\gioca\Olive\examples\directml\stable_diffusion_xl\cache\default_workflow\olive_config.json
[2024-08-21 15:17:52,349] [INFO] [accelerator_creator.py:224:create_accelerators] Running workflow on accelerator specs: gpu-cuda
[2024-08-21 15:17:52,349] [INFO] [engine.py:275:run] Running Olive on accelerator: gpu-cuda
[2024-08-21 15:17:52,359] [INFO] [engine.py:1109:_create_system] Creating target system ...
[2024-08-21 15:17:52,359] [INFO] [engine.py:1112:_create_system] Target system created in 0.000000 seconds
[2024-08-21 15:17:52,359] [INFO] [engine.py:1121:_create_system] Creating host system ...
[2024-08-21 15:17:52,359] [INFO] [engine.py:1124:_create_system] Host system created in 0.000000 seconds
[2024-08-21 15:17:52,376] [INFO] [engine.py:877:_run_pass] Running pass convert:OnnxConversion
[2024-08-21 15:17:52,376] [INFO] [engine.py:914:_run_pass] Loaded model from cache: 2_OnnxConversion-958bc3c7-076cfb73
[2024-08-21 15:17:52,376] [INFO] [engine.py:877:_run_pass] Running pass optimize_cuda:OrtTransformersOptimization
[2024-08-21 15:17:52,376] [INFO] [engine.py:914:_run_pass] Loaded model from cache: 3_OrtTransformersOptimization-2-900e718a-gpu-cuda
[2024-08-21 15:17:52,438] [INFO] [engine.py:374:run_accelerator] Save footprint to footprints\vae_decoder_gpu-cuda_footprints.json.
[2024-08-21 15:17:52,438] [INFO] [engine.py:292:run] Run history for gpu-cuda:
[2024-08-21 15:17:52,453] [INFO] [engine.py:587:dump_run_history] run history:
+---------------------------------------------------+------------------------------------+-----------------------------+----------------+-----------+
| model_id | parent_model_id | from_pass | duration_sec | metrics |
+===================================================+====================================+=============================+================+===========+
| 958bc3c7 | | | | |
+---------------------------------------------------+------------------------------------+-----------------------------+----------------+-----------+
| 2_OnnxConversion-958bc3c7-076cfb73 | 958bc3c7 | OnnxConversion | 29.7135 | |
+---------------------------------------------------+------------------------------------+-----------------------------+----------------+-----------+
| 3_OrtTransformersOptimization-2-900e718a-gpu-cuda | 2_OnnxConversion-958bc3c7-076cfb73 | OrtTransformersOptimization | 3.60028 | |
+---------------------------------------------------+------------------------------------+-----------------------------+----------------+-----------+
[2024-08-21 15:17:52,453] [INFO] [engine.py:307:run] No packaging config provided, skip packaging artifacts
Unoptimized Model : C:\Users\gioca\Olive\examples\directml\stable_diffusion_xl\cache\default_workflow\models\2_OnnxConversion-958bc3c7-076cfb73\output_model\model.onnx
Optimized Model : C:\Users\gioca\Olive\examples\directml\stable_diffusion_xl\cache\default_workflow\models\3_OrtTransformersOptimization-2-900e718a-gpu-cuda\output_model\model.onnx

Optimizing unet
Traceback (most recent call last):
File "C:\Users\gioca\Olive\examples\directml\stable_diffusion_xl\stable_diffusion_xl.py", line 633, in
main()
File "C:\Users\gioca\Olive\examples\directml\stable_diffusion_xl\stable_diffusion_xl.py", line 599, in main
optimize(
File "C:\Users\gioca\Olive\examples\directml\stable_diffusion_xl\stable_diffusion_xl.py", line 367, in optimize
olive_run(olive_config)
File "C:\Users\gioca\anaconda3\envs\sdxl_cuda\Lib\site-packages\olive\workflows\run\run.py", line 296, in run
run_config: RunConfig = RunConfig.parse_file_or_obj(run_config)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\gioca\anaconda3\envs\sdxl_cuda\Lib\site-packages\olive\common\config_utils.py", line 119, in parse_file_or_obj
return cls.parse_obj(file_or_obj)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\gioca\anaconda3\envs\sdxl_cuda\Lib\site-packages\pydantic\v1\main.py", line 526, in parse_obj
return cls(**obj)
^^^^^^^^^^
File "C:\Users\gioca\anaconda3\envs\sdxl_cuda\Lib\site-packages\pydantic\v1\main.py", line 339, in init
values, fields_set, validation_error = validate_model(pydantic_self.class, data)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\gioca\anaconda3\envs\sdxl_cuda\Lib\site-packages\pydantic\v1\main.py", line 1074, in validate_model
v_, errors_ = field.validate(value, values, loc=field.alias, cls=cls_)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\gioca\anaconda3\envs\sdxl_cuda\Lib\site-packages\pydantic\v1\fields.py", line 864, in validate
v, errors = self._apply_validators(v, values, loc, cls, self.pre_validators)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\gioca\anaconda3\envs\sdxl_cuda\Lib\site-packages\pydantic\v1\fields.py", line 1154, in _apply_validators
v = validator(cls, v, values, self, self.model_config)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\gioca\anaconda3\envs\sdxl_cuda\Lib\site-packages\pydantic\v1\class_validators.py", line 304, in
return lambda cls, v, values, field, config: validator(cls, v)
^^^^^^^^^^^^^^^^^
File "C:\Users\gioca\anaconda3\envs\sdxl_cuda\Lib\site-packages\olive\workflows\run\config.py", line 181, in validate_data_config_names
data_config_obj = validate_config(data_config, DataConfig)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\gioca\anaconda3\envs\sdxl_cuda\Lib\site-packages\olive\common\config_utils.py", line 329, in validate_config
config = instance_class(**config)
^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\gioca\anaconda3\envs\sdxl_cuda\Lib\site-packages\olive\data\config.py", line 77, in init
self._fill_in_params()
File "C:\Users\gioca\anaconda3\envs\sdxl_cuda\Lib\site-packages\olive\data\config.py", line 132, in _fill_in_params
component = Registry.get_component(k, v.type)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "C:\Users\gioca\anaconda3\envs\sdxl_cuda\Lib\site-packages\olive\data\registry.py", line 178, in get_component
return cls._REGISTRY[component][name.lower()]
~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^
KeyError: 'unet_dataloader'

Other information

• OS : Windows 11
• Python : 3.12.4
• Packages
  accelerate 0.33.0
  aiohappyeyeballs 2.4.0
  aiohttp 3.10.5
  aiosignal 1.3.1
  alembic 1.13.2
  annotated-types 0.7.0
  attrs 24.2.0
  certifi 2024.7.4
  charset-normalizer 3.3.2
  colorama 0.4.6
  coloredlogs 15.0.1
  colorlog 6.8.2
  datasets 2.21.0
  diffusers 0.29.2
  dill 0.3.8
  filelock 3.15.4
  flatbuffers 24.3.25
  frozenlist 1.4.1
  fsspec 2024.6.1
  greenlet 3.0.3
  huggingface-hub 0.24.6
  humanfriendly 10.0
  idna 3.7
  importlib_metadata 8.3.0
  intel-openmp 2021.4.0
  invisible-watermark 0.2.0
  Jinja2 3.1.4
  lightning-utilities 0.11.6
  Mako 1.3.5
  MarkupSafe 2.1.5
  mkl 2021.4.0
  mpmath 1.3.0
  multidict 6.0.5
  multiprocess 0.70.16
  networkx 3.3
  numpy 1.26.4
  olive-ai 0.7.0
  onnx 1.16.1
  onnxruntime-gpu/directml 1.18.0
  opencv-python 4.10.0.84
  optimum 1.21.4
  optuna 3.6.1
  packaging 24.1
  pandas 2.2.2
  pillow 10.4.0
  pip 24.2
  protobuf 3.20.3
  psutil 6.0.0
  pyarrow 17.0.0
  pydantic 2.8.2
  pydantic_core 2.20.1
  pyreadline3 3.4.1
  python-dateutil 2.9.0.post0
  pytz 2024.1
  PyWavelets 1.7.0
  PyYAML 6.0.2
  regex 2024.7.24
  requests 2.32.3
  safetensors 0.4.4
  sentencepiece 0.2.0
  setuptools 72.1.0
  six 1.16.0
  SQLAlchemy 2.0.32
  sympy 1.13.2
  tabulate 0.9.0
  tbb 2021.13.1
  tokenizers 0.19.1
  torch 2.3.1 (←downgrade from 2.4.0)
  torchmetrics 1.4.1
  tqdm 4.66.5
  transformers 4.42.4
  typing_extensions 4.12.2
  tzdata 2024.1
  urllib3 2.2.2
  wheel 0.43.0
  xxhash 3.5.0
  yarl 1.9.4
  zipp 3.20.0

[giocafe]

After removing the below statement in config_unet.json file, conversion to onnx and latency optimization are done.
"data_configs": [ { "name": "latency_data_config", "user_script": "user_script.py", "load_dataset_config": { "type": "simple_dataset" }, "dataloader_config": { "type": "unet_dataloader", "batch_size": 2 } } ],

And then, I performed an inference test.
However, the image window is created, but no image appears and a black window appears.
At this time, the GPU usage appears to be 100% for a moment and then returns to 0%,
and the GPU memory usage increases.
And the prompt just stops without leaving any logs.

And when I close the created image window, the prompt returns to the command line.
At this time, the GPU memory usage decreases again.

Judging from the increase in GPU usage, the image creation task via the GPU seems to have been performed normally,
but since the resulting image is a black screen, I doubt whether this is normal operation.