Remove background - onnx-community/BEN2-ONNX - bug when using webGPU

[pikilipita]

System Info

transformers 3.4.1
macOs 14.7.4
Chrome 136

Environment/Platform

• Website/web-app
• Browser extension
• Server-side (e.g., Node.js, Deno, Bun)
• Desktop app (e.g., Electron)
• Other (e.g., VSCode extension)

Description

While using onnx-community/BEN2-ONNX to remove background, the wasm version result is fine, but the webGPU result adds a 'ghost' on the mask:

the original image of this example be found here, but the ghosting issue is happening on all images I tested.
https://unsplash.com/photos/woman-standing-near-pink-concrete-wall-during-daytime-TT-ROxWj9nA

tested on transformers 3.4.1 with those settings:

var segmenter = await pipeline("background-removal", "onnx-community/BEN2-ONNX", { device:"webgpu" });

Reproduction

to reproduce, use those settings with any input image:
var segmenter = await pipeline("background-removal", "onnx-community/BEN2-ONNX", { device:"webgpu" });

[xenova]

cc @fs-eire in case you have any ideas.

[prathikr]

Looking into the issue...

[prathikr]

Hi @pikilipita do you have a full index.html that reproduces this issue? Something that generates the binary image and displays it on the page?

[pikilipita]

@prathikr
yes, I just made one for you:
https://pikimov.com/tr-test/

[prathikr]

Hi @pikilipita, thank you for your patience. @fs-eire and I looked at your model and noticed a peculiar pattern. In both the multifieldcrossatt and dec_blk modules there is a softmax that feeds to a cast-to-fp32 followed by a cast-to-fp16:

If I adjust it to follow the pattern cast-to-fp32 --> softmax --> cast-to-fp16 then the model works:

Can you please double check how the multifieldcrossatt and dec_blk modules are implemented to see if the second pattern can be followed?

[pikilipita]

Hello @prathikr , thank you for looking into the issue,
Where would you like me to look for modules? Not sure I can help, I'm just using recommended transformers pipeline code, I'm not an LLM expert at all.

[prathikr]

@xenova I see you uploaded the BEN2 model to HF model hub. Do you know where the model code is to check why the cast nodes are being exported incorrectly?

[xenova]

It's using the float16 quantization script: https://github.com/huggingface/transformers.js/blob/main/scripts/float16.py, which is adapted from a bunch of other scripts.

I'm pretty certain the bug is from there, for when there is already an existing cast node in the graph.

Would greatly appreciate help in fixing it @prathikr if you have time.

[prathikr]

@xenova do you use BEN2_Base.onnx as input to this script? I can see the softmax --> cast to fp32 --> cast to fp16 is in the base model so this script may not be the issue. It is likely an issue in the original model implementation. I have a version of onnx-community/BEN2-ONNX with the swapped nodes and the script that swaps the nodes. I can share both via email if that would be helpful.