Add the option to choose cache directory

[alexblattner]

Is your feature request related to a problem? Please describe. I'd like to deploy on replicate perfectly, but without having it download the models every time.

Describe the solution you'd like Using a cache directory of your choice to push the cache to replicate

Describe alternatives you've considered Have it download again and again every time

Additional context Add any other context or screenshots about the feature request here.

[danielgatis]

Try to set the env var U2NET_HOME. take a look at: https://github.com/danielgatis/rembg/blob/main/rembg/sessions/base.py#L78

[alexblattner]

@danielgatis what if I don't want to change env? I mean it is definitely a solution, but feels awkward tbh

[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.

[alexblattner]

solution, this saves to the local directory named 'insightface':

app = FaceAnalysis(name="buffalo_l",root='insightface', providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])

[MatthewDegtyar]

For anyone else on this journey I got my runtimes times down from ~9s to ~2s (depending on image size).

Notes: My .env file for caching: U2NET_HOME = "u2net"

This was run on an M2 (cpu). I have yet to see if this will work on AMD64 Linux as intended.

BTW I run a product image startup: Deggie AI

Solution

import onnxruntime as ort from rembg.sessions import U2netSession from rembg import remove

image = None # load your image here

sess_opts = ort.SessionOptions() providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']

session = U2netSession( model_name='u2net_model', sess_opts=sess_opts, root_dir='./session-cache', providers=providers )

model_type: str = 'u2net' # Default model type use_gpu:bool = False # Option to use GPU alpha_matting: bool = False alpha_matting_foreground_threshold: int = 240 alpha_matting_background_threshold: int = 10 alpha_matting_erode_size: int = 5 post_process_mask: bool = True always_return_png: bool = True

remove_options = { 'alpha_matting': alpha_matting, 'alpha_matting_foreground_threshold': alpha_matting_foreground_threshold, 'alpha_matting_background_threshold': alpha_matting_background_threshold, 'alpha_matting_erode_size': alpha_matting_erode_size, 'session': session, 'post_process_mask': post_process_mask }

output_image = remove(image, **remove_options) output_image = output_image.convert("RGBA")

Logs

With suggested code:

• Total time session init: 0.34816789627075195
• Total time rm_operation: 1.2982029914855957
• Total time rembg: 1.64644193649292
• Total time success DEV: 2.1654109954833984

W/o suggested code:

• Total time session init: 7.7538909912109375
• Total time rm_operation: 1.1261260509490967
• Total time rembg: 8.880061149597168
• Total time success DEV: 9.620557069778442

Image output:

Image input: