[BUG] GPU doesn't work out of the box, workaround & benchmarks

[luke14free]

Describe the bug To run on the GPU you need to uninstall and re-install onnxruntime-gpu.

!pip install rembg[cli,gpu]
import onnxruntime as ort
ort.get_device()

Will return "CPU". But:

!pip install rembg[cli,gpu]
!pip uninstall onnxruntime-gpu --yes
!pip install onnxruntime-gpu

import onnxruntime as ort
ort.get_device()

Will return "GPU".

FYI to all of those interested, benchmarks on running on GPUs made us understand it's a bit useless to run on GPUs as improvements are quite marginal but costs will be quite higher (maybe parallel processing makes this more interesting, haven't tested much of this).

!wget https://raw.githubusercontent.com/danielgatis/rembg/master/examples/car-1.jpg
!mv car-1.jpg input.jpg

from rembg import remove

input_path = 'input.jpg'
output_path = 'output.png'

with open(input_path, 'rb') as i:
    with open(output_path, 'wb') as o:
        input = i.read()

then:

%%timeit
output = remove(input)

Results on Colab:

• CPU: 3.29s +/- 895 ms per loop
• GPU (T4): 3.79s +/- 56ms per loop
• GPU (A100): 2.85s +/- 11.3ms per loop

[github-actions[bot]]

This issue is stale because it has been open for 30 days with no activity.

[github-actions[bot]]

This issue was closed because it has been inactive for 14 days since being marked as stale.

[luke14free]

FYI, to those in need of better GPU speed -> https://github.com/nadermx/backgroundremover works much better with execution speed down to 0.4s on GPUs.