[HF][Optimum] Compiling unet in stable diffusion XL pipeline failed since Neuron SDK 2.18 about aws-neuron-sdk

[aws-bhegedus]

Thanks Jingya, we will have a fix in a future release to allow enabling the caching.
Is this problem only there for SDXL-random-tiny? Curious about SDXL-base, which I believe is larger and takes longer to compile so may be a bigger problem.

[JingyaHuang]

Also tried with PyTorch 2.1.2 setup, not working neither.

[aws-bhegedus]

Hi Jingya, I'm trying to reproduce the problem. I installed optimum and optimum-neuron with
pip install "optimum[neuronx, diffusers]"
based on https://huggingface.co/docs/optimum-neuron/tutorials/stable_diffusion.
However this seems to get v0.0.3 which doesn't find NeuronStableDiffusionXLPipeline. I also tried downgrading to 0.0.2 which has another problem. Is this expected with these versions, and is there a way to get 0.021? Thanks.

[JingyaHuang]

The installation with neuronx extra is what we are going to fix with the 0.0.21 optimum-neuron release. Fow now, to install the latest optimum-neuron release(0.0.20), could you try with:

pip install optimum==1.18.0
pip install optimum-neuron==0.0.20

Or the 0.0.21 dev version can be installed from source:

pip install git+https://github.com/huggingface/optimum-neuron

Then you could install pip install diffusers.

[aws-bhegedus]

Thanks Jingya, I updated optimum-neuron and diffusers and now I can reproduce the issue.

[aws-bhegedus]

Hi Jingya, I found that the issue can be prevented if we set inline_weights_to_neff=True when tracing the UNet. Would that be a sufficient workaround for now? I will also look into the root cause but that may take some time.

[JingyaHuang]

Hi @aws-bhegedus, thanks for investigating it!

Optimum Neuron could force setting inline_weights_to_neff=False for sdxl models for now. But given that our caching mechanism relies on the neff weights separation, we won't be able to cache and load sdxl models (which takes time for the compilation).

[JingyaHuang]

Thanks @aws-bhegedus, that will be awesome!

tiny-random-stable-diffusion-xl is a smaller version (fewer layers) of sdxl models in the pipe with random weights that we built to shorten the testing time, if the compilation fails for the tiny version, it's very unlikely that it could work for the larger pretrained checkpoint. And since the compilation of all sdxl components takes more than an hour, without the caching it could be a little bit discouraging for first-time users.