cleanup rembg

Signed-off-by: vladmandic <mandic00@live.com>

modules/rembg/ben2.py

@@ -0,0 +1,24 @@
+model = None
+
+
+def remove(image, refine: bool = True):
+    global model # pylint: disable=global-statement
+    from modules import shared, devices
+
+    if model is None:
+        from huggingface_hub import hf_hub_download
+        from modules.rembg.ben2_model import BEN_Base
+        model = BEN_Base()
+        model_file = hf_hub_download(
+            repo_id='PramaLLC/BEN2',
+            filename='BEN2_Base.pth',
+            cache_dir=shared.opts.hfcache_dir)
+        model.loadcheckpoints(model_file)
+        model = model.to(device=devices.device, dtype=devices.dtype).eval()
+
+    model = model.to(device=devices.device)
+    foreground = model.inference(image, refine_foreground=refine)
+    model = model.to(device=devices.cpu)
+    if foreground is None:
+        return image
+    return foreground

modules/rembg/rembg_api.py

@@ -0,0 +1,44 @@
+from fastapi import Body
+
+
+async def post_rembg(
+    input_image: str = Body("", title='rembg input image'),
+    model: str = Body("u2net", title='rembg model'),
+    return_mask: bool = Body(False, title='return mask'),
+    alpha_matting: bool = Body(False, title='alpha matting'),
+    alpha_matting_foreground_threshold: int = Body(240, title='alpha matting foreground threshold'),
+    alpha_matting_background_threshold: int = Body(10, title='alpha matting background threshold'),
+    alpha_matting_erode_size: int = Body(10, title='alpha matting erode size'),
+    refine: bool = Body(False, title="refine foreground (ben2 only)")
+):
+    if not model or model == "None":
+        return {}
+
+    from modules.api import api
+    input_image = api.decode_base64_to_image(input_image)
+    if input_image is None:
+        return {}
+
+    if model == "ben2":
+        from modules.rembg import ben2
+        image = ben2.remove(input_image, refine=refine)
+    else:
+        from installer import install
+        for pkg in ["dctorch==0.1.2", "pymatting", "pooch", "rembg"]:
+            install(pkg, no_deps=True, ignore=False)
+        import rembg
+        image = rembg.remove(
+            input_image,
+            session=rembg.new_session(model),
+            only_mask=return_mask,
+            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,
+        )
+    return {"image": api.encode_pil_to_base64(image).decode("utf-8")}
+
+
+def register_api(app):
+    print('HERE')
+    app.add_api_route("/sdapi/v1/rembg", post_rembg, methods=["POST"], tags=["REMBG"])