When trying to convert either the Decoder or the UNET with flexible shapes, coremltools fails.

The first op to blame is the group_norm op.
I found a related issue: apple/coremltools#1303 where they provided a work around for the group_norm op - which worked, but then another layer/op throws and error. It might be add element wise or something else. Maybe there is another work around for other ops?

I thought that maybe I could use the neural_network.flexible_shape_utils to modify the network after the fact, but that apparently only works with the old neural network type of CoreML and not the MLProgram type that probably is maybe required for the ANE and thus faster initialization. Maybe could test that with by omitting convert_to="mlprogram", but I'm sure that's probably a lot slower.

I did investigate two UNET models generated with ORIGINAL attention. One with a latent height of 96 and the other the normal 64x64. I diffed the weights and there is no difference. Seems like the main difference is in the architecture part of the MLPackage. deploying such large files for each aspect ratios made available to users sounds insane. Potentially, maybe there is a way to compile these on device and switching out the architecture portions of the MLPackage files/folder?

I've not had any luck with the SPLIT_EINSUM setting with regards to different aspect ratios. Because when testing them, I'll get kernel panics... see: SPLIT_EINSUM - Kernel Panic when testing UNET created with height 96 and width 64 and SPLIT_EINSUM