peterheb/rembg-lambda

Demo of rembg running in AWS Lambda

rembg-lambda's Issues

[Suggestion] update documentation to add a link to `rembg-aws-lambda`

Hello, this is just a suggestion in regards to the following line in the Dockerfile:

RUN mkdir .u2net && curl -L https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net.onnx -o .u2net/u2net.onnx

The good news is that this download step can be circumvented, if such is desired -- keep in mind the following approach is pure Python and has no CLI support, however.

In case it is helpful for other Python developers, I've created an almost identical project for AWS environments, rembg-aws-lambda. It strips down some code, such as alpha matting, dependencies, CLI, and downloading the *.onnx files. The PyPI package contains u2net.onnx already, so installing it or setting env variables as NUMBA_CACHE_DIR and U2NET_HOME is not required anymore.

The project can be installed with:

pip install rembg-aws-lambda

Re. any cold start issues, they should be eliminated by this variant package. I have done extensive testing before this previously, and determined that the cold start issue(s) to be primarly due to:

1. downloading the *.onnx files into the /tmp directory each time
2. import statement for pymatting - as this package heavily utilizes decorators such as @njit, which uses numba to cache function optimizations. As this saves it into NUMBA_CACHE_DIR which should ideally point to /tmp for AWS Lambda Docker images (the only writable directory on the container), this takes a lot of time on each cold start as /tmp directory is clean (empty) at this point.

I have tested this out and determined that an AWS Lambda function (deployed w/ Docker, see a sample Dockerfile) using an X86_64 architecture and 1769 MB memory to be optimal. Remember to also set OMP_NUM_THREADS=2 in the environment, to take advantage of multiprocessing/threads as well.

The initial run on AWS Lambda with the above specs appears to take ~45 seconds on an initial run, but is much faster (< 10 seconds) on subsequent executions.

Check it out, I'd be curious to hear your thoughts. I can also create a PR to add a note in the documentation to mention this approach, as needed.