├── .gitignore ├── LICENSE ├── README.md ├── install.py ├── preview.png └── scripts ├── api.py └── postprocessing_rembg.py /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 AUTOMATIC1111 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Rembg 2 | 3 | Extension for [webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui). Removes backgrounds from pictures. 4 | 5 | ![](preview.png) 6 | 7 | Find the UI for rembg in the Extras tab after installing the extension. 8 | 9 | # Installation 10 | 11 | Install from webui's Extensions tab. 12 | 13 | # Credits 14 | 15 | * rembg library that does all the work: https://github.com/danielgatis/rembg 16 | -------------------------------------------------------------------------------- /install.py: -------------------------------------------------------------------------------- 1 | import launch 2 | 3 | if not launch.is_installed("rembg"): 4 | launch.run_pip("install rembg==2.0.50 --no-deps", "rembg") 5 | 6 | for dep in ['onnxruntime', 'pymatting', 'pooch']: 7 | if not launch.is_installed(dep): 8 | launch.run_pip(f"install {dep}", f"{dep} for REMBG extension") 9 | -------------------------------------------------------------------------------- /preview.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AUTOMATIC1111/stable-diffusion-webui-rembg/a4c07b857e73f3035f759876797fa6de986def3d/preview.png -------------------------------------------------------------------------------- /scripts/api.py: -------------------------------------------------------------------------------- 1 | from fastapi import FastAPI, Body 2 | 3 | from modules.api.models import * 4 | from modules.api import api 5 | import gradio as gr 6 | 7 | import rembg 8 | 9 | # models = [ 10 | # "None", 11 | # "u2net", 12 | # "u2netp", 13 | # "u2net_human_seg", 14 | # "u2net_cloth_seg", 15 | # "silueta", 16 | # ] 17 | 18 | 19 | def rembg_api(_: gr.Blocks, app: FastAPI): 20 | @app.post("/rembg") 21 | async def rembg_remove( 22 | input_image: str = Body("", title='rembg input image'), 23 | model: str = Body("u2net", title='rembg model'), 24 | return_mask: bool = Body(False, title='return mask'), 25 | alpha_matting: bool = Body(False, title='alpha matting'), 26 | alpha_matting_foreground_threshold: int = Body(240, title='alpha matting foreground threshold'), 27 | alpha_matting_background_threshold: int = Body(10, title='alpha matting background threshold'), 28 | alpha_matting_erode_size: int = Body(10, title='alpha matting erode size') 29 | ): 30 | if not model or model == "None": 31 | return 32 | 33 | input_image = api.decode_base64_to_image(input_image) 34 | 35 | image = rembg.remove( 36 | input_image, 37 | session=rembg.new_session(model), 38 | only_mask=return_mask, 39 | alpha_matting=alpha_matting, 40 | alpha_matting_foreground_threshold=alpha_matting_foreground_threshold, 41 | alpha_matting_background_threshold=alpha_matting_background_threshold, 42 | alpha_matting_erode_size=alpha_matting_erode_size, 43 | ) 44 | 45 | return {"image": api.encode_pil_to_base64(image).decode("utf-8")} 46 | 47 | try: 48 | import modules.script_callbacks as script_callbacks 49 | 50 | script_callbacks.on_app_started(rembg_api) 51 | except: 52 | pass 53 | -------------------------------------------------------------------------------- /scripts/postprocessing_rembg.py: -------------------------------------------------------------------------------- 1 | from modules import scripts_postprocessing, ui_components 2 | import gradio as gr 3 | 4 | from modules.ui_components import FormRow 5 | from modules.paths_internal import models_path 6 | import rembg 7 | import os 8 | 9 | models = [ 10 | "None", 11 | "isnet-general-use", 12 | "u2net", 13 | "u2netp", 14 | "u2net_human_seg", 15 | "u2net_cloth_seg", 16 | "silueta", 17 | "isnet-general-use", 18 | "isnet-anime", 19 | ] 20 | 21 | class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing): 22 | name = "Rembg" 23 | order = 20000 24 | model = None 25 | 26 | def ui(self): 27 | with ui_components.InputAccordion(False, label="Remove background") as enable: 28 | with gr.Row(): 29 | model = gr.Dropdown(label="Remove background", choices=models, value="None") 30 | return_mask = gr.Checkbox(label="Return mask", value=False) 31 | alpha_matting = gr.Checkbox(label="Alpha matting", value=False) 32 | 33 | with gr.Row(visible=False) as alpha_mask_row: 34 | alpha_matting_erode_size = gr.Slider(label="Erode size", minimum=0, maximum=40, step=1, value=10) 35 | alpha_matting_foreground_threshold = gr.Slider(label="Foreground threshold", minimum=0, maximum=255, step=1, value=240) 36 | alpha_matting_background_threshold = gr.Slider(label="Background threshold", minimum=0, maximum=255, step=1, value=10) 37 | 38 | alpha_matting.change( 39 | fn=lambda x: gr.update(visible=x), 40 | inputs=[alpha_matting], 41 | outputs=[alpha_mask_row], 42 | ) 43 | 44 | return { 45 | "enable": enable, 46 | "model": model, 47 | "return_mask": return_mask, 48 | "alpha_matting": alpha_matting, 49 | "alpha_matting_foreground_threshold": alpha_matting_foreground_threshold, 50 | "alpha_matting_background_threshold": alpha_matting_background_threshold, 51 | "alpha_matting_erode_size": alpha_matting_erode_size, 52 | } 53 | 54 | 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): 55 | if not enable: 56 | return 57 | 58 | if not model or model == "None": 59 | return 60 | 61 | if "U2NET_HOME" not in os.environ: 62 | os.environ["U2NET_HOME"] = os.path.join(models_path, "u2net") 63 | 64 | pp.image = rembg.remove( 65 | pp.image, 66 | session=rembg.new_session(model), 67 | only_mask=return_mask, 68 | alpha_matting=alpha_matting, 69 | alpha_matting_foreground_threshold=alpha_matting_foreground_threshold, 70 | alpha_matting_background_threshold=alpha_matting_background_threshold, 71 | alpha_matting_erode_size=alpha_matting_erode_size, 72 | ) 73 | 74 | pp.info["Rembg"] = model 75 | --------------------------------------------------------------------------------