diff --git a/install.py b/install.py index 8ec2c14..1a73e3c 100644 --- a/install.py +++ b/install.py @@ -1,7 +1,7 @@ import launch if not launch.is_installed("rembg"): - launch.run_pip("install rembg==2.0.36 --no-deps", "rembg") + launch.run_pip("install rembg==2.0.38 --no-deps", "rembg") for dep in ['onnxruntime', 'pymatting', 'pooch']: if not launch.is_installed(dep): diff --git a/scripts/api.py b/scripts/api.py new file mode 100644 index 0000000..0a0abd3 --- /dev/null +++ b/scripts/api.py @@ -0,0 +1,52 @@ +from fastapi import FastAPI, Body + +from modules.api.models import * +from modules.api import api +import gradio as gr + +import rembg + +# models = [ +# "None", +# "u2net", +# "u2netp", +# "u2net_human_seg", +# "u2net_cloth_seg", +# "silueta", +# ] + + +def rembg_api(_: gr.Blocks, app: FastAPI): + @app.post("/rembg") + async def rembg_remove( + 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') + ): + if not model or model == "None": + return + + input_image = api.decode_base64_to_image(input_image) + + 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")} + +try: + import modules.script_callbacks as script_callbacks + + script_callbacks.on_app_started(rembg_api) +except: + pass diff --git a/scripts/postprocessing_rembg.py b/scripts/postprocessing_rembg.py index 3e1417f..d50c51e 100644 --- a/scripts/postprocessing_rembg.py +++ b/scripts/postprocessing_rembg.py @@ -1,8 +1,10 @@ -from modules import scripts_postprocessing +from modules import scripts_postprocessing, ui_components import gradio as gr from modules.ui_components import FormRow +from modules.paths_internal import models_path import rembg +import os models = [ "None", @@ -12,6 +14,8 @@ models = [ "u2net_human_seg", "u2net_cloth_seg", "silueta", + "isnet-general-use", + "isnet-anime", ] class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing): @@ -20,23 +24,25 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing): model = None def ui(self): - with FormRow(): - model = gr.Dropdown(label="Remove background", choices=models, value="None") - return_mask = gr.Checkbox(label="Return mask", value=False) - alpha_matting = gr.Checkbox(label="Alpha matting", value=False) + with ui_components.InputAccordion(False, label="Remove background") as enable: + with gr.Row(): + model = gr.Dropdown(label="Remove background", choices=models, value="None") + return_mask = gr.Checkbox(label="Return mask", value=False) + alpha_matting = gr.Checkbox(label="Alpha matting", value=False) - with FormRow(visible=False) as alpha_mask_row: - alpha_matting_erode_size = gr.Slider(label="Erode size", minimum=0, maximum=40, step=1, value=10) - alpha_matting_foreground_threshold = gr.Slider(label="Foreground threshold", minimum=0, maximum=255, step=1, value=240) - alpha_matting_background_threshold = gr.Slider(label="Background threshold", minimum=0, maximum=255, step=1, value=10) + with gr.Row(visible=False) as alpha_mask_row: + alpha_matting_erode_size = gr.Slider(label="Erode size", minimum=0, maximum=40, step=1, value=10) + alpha_matting_foreground_threshold = gr.Slider(label="Foreground threshold", minimum=0, maximum=255, step=1, value=240) + alpha_matting_background_threshold = gr.Slider(label="Background threshold", minimum=0, maximum=255, step=1, value=10) - alpha_matting.change( - fn=lambda x: gr.update(visible=x), - inputs=[alpha_matting], - outputs=[alpha_mask_row], - ) + alpha_matting.change( + fn=lambda x: gr.update(visible=x), + inputs=[alpha_matting], + outputs=[alpha_mask_row], + ) return { + "enable": enable, "model": model, "return_mask": return_mask, "alpha_matting": alpha_matting, @@ -45,10 +51,16 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing): "alpha_matting_erode_size": alpha_matting_erode_size, } - def process(self, pp: scripts_postprocessing.PostprocessedImage, model, return_mask, alpha_matting, alpha_matting_foreground_threshold, alpha_matting_background_threshold, alpha_matting_erode_size): - if model == "None": + def process(self, pp: scripts_postprocessing.PostprocessedImage, enable, model, return_mask, alpha_matting, alpha_matting_foreground_threshold, alpha_matting_background_threshold, alpha_matting_erode_size): + if not enable: return + if not model or model == "None": + return + + if "U2NET_HOME" not in os.environ: + os.environ["U2NET_HOME"] = os.path.join(models_path, "u2net") + pp.image = rembg.remove( pp.image, session=rembg.new_session(model),