Unsupported "InstanceNormalization" when importing vae_encoder.onnx of StableDiffusion

[Goody27]

I'm trying to get StableDiffusion working using AxonOnnx.

Reproduction steps

1. Download the file from https://huggingface.co/bes-dev/stable-diffusion-v1-4-onnx/raw/main/vae_encoder.onnx into models/onnx/vae_encoder.onnx.
2. Do {model, params} = AxonOnnx.import("models/onnx/vae_encoder.onnx") on iex or Livebook.
3. Then, you will get the following error:

** (ArgumentError) unsupported "InstanceNormalization"
    (axon_onnx 0.3.0) lib/axon_onnx/deserialize.ex:2144: AxonOnnx.Deserialize.recur_nodes/2
    (elixir 1.14.1) lib/enum.ex:2468: Enum."-reduce/3-lists^foldl/2-0-"/3
    (axon_onnx 0.3.0) lib/axon_onnx/deserialize.ex:44: AxonOnnx.Deserialize.graph_to_axon/2
    (axon_onnx 0.3.0) lib/axon_onnx/deserialize.ex:27: AxonOnnx.Deserialize.to_axon/2
    (stdlib 4.1.1) erl_eval.erl:744: :erl_eval.do_apply/7
    (stdlib 4.1.1) erl_eval.erl:492: :erl_eval.expr/6
    (elixir 1.14.1) src/elixir.erl:288: :elixir.eval_forms/3
    (elixir 1.14.1) lib/module/parallel_checker.ex:107: Module.ParallelChecker.verify/1

Environment

Erlang/OTP 25 [erts-13.1.1] [source] [64-bit] [smp:24:24] [ds:24:24:10] [async-threads:1] [jit:ns]

Elixir 1.14.1 (compiled with Erlang/OTP 23)

Mix.install(
  [
    {:nx, "~> 0.4"},
    {:axon, "~> 0.3"},
    {:exla, "~> 0.4"},
    {:axon_onnx, "~> 0.3"},
    {:stb_image, "~> 0.5"},
    {:kino, "~> 0.7.0"},
    {:table_rex, "~> 3.1.1"}
  ]
  # change to "cuda111" for Nvidia GPU
  # system_env: %{"XLA_TARGET" => xla_target}
)

• Windows 11 Pro + WSL
• 12th Gen Intel(R) Core(TM) i9-12900 2.40 GHz
• Memory: 64.0 GB

[seanmor5]

Thanks for opening! We support instance norm in Axon so this should be straightforward enough to implement, I will take a look this weekend

[seanmor5]

I've added instance norm on master, but importing the VAE doesn't work, there's some strange things going on that don't seem to respect the ONNX spec. I will look further into it

[Goody27]

Thank you for your research.