AI/rembg
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg synced 2025-08-16 23:25:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #755 from divinity76/patch-2

Browse Source
This commit is contained in:
Daniel Gatis 2025-05-11 10:38:09 -03:00 committed by GitHub
commit dbffeff61d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 10 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
rembg/session_factory.py
View File

@ -22,20 +22,27 @@ def new_session(model_name: str = "u2net", *args, **kwargs) -> BaseSession:
*args: Additional positional arguments.
**kwargs: Additional keyword arguments.
Raises:
ValueError: If no session class with the given `model_name` is found.
Returns:
BaseSession: The created session object.
"""
session_class: Type[BaseSession] = U2netSession
session_class: Optional[Type[BaseSession]] = None
for sc in sessions_class:
if sc.name() == model_name:
session_class = sc
break
if session_class is None:
raise ValueError(f"No session class found for model '{model_name}'")
sess_opts = ort.SessionOptions()
if "OMP_NUM_THREADS" in os.environ:
sess_opts.inter_op_num_threads = int(os.environ["OMP_NUM_THREADS"])
sess_opts.intra_op_num_threads = int(os.environ["OMP_NUM_THREADS"])
threads = int(os.environ["OMP_NUM_THREADS"])
sess_opts.inter_op_num_threads = threads
sess_opts.intra_op_num_threads = threads
return session_class(model_name, sess_opts, *args, **kwargs)