AUTOMATIC1111/stable-diffusion-webui-rembg
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui-rembg.git synced 2026-05-01 03:31:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Changed the default download location of u2net models to <stable-diffusion-webui>/models/u2net

Browse Source
This commit is contained in:
h3rmit
2023-08-28 01:45:27 +03:00
parent 3d9eedbbf0
commit 1c4a698cc8
1 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
scripts/postprocessing_rembg.py
View File

@@ -2,7 +2,9 @@ from modules import scripts_postprocessing
import gradio as gr import gradio as gr
from modules.ui_components import FormRow from modules.ui_components import FormRow
from modules.paths_internal import models_path
import rembg import rembg
import os
models = [ models = [
"None", "None",
@@ -50,6 +52,9 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
if not model or model == "None": if not model or model == "None":
return return
if "U2NET_HOME" not in os.environ:
os.environ["U2NET_HOME"] = os.path.join(models_path, "u2net")
pp.image = rembg.remove( pp.image = rembg.remove(
pp.image, pp.image,
session=rembg.new_session(model), session=rembg.new_session(model),