AUTOMATIC1111/stable-diffusion-webui-rembg
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui-rembg.git synced 2026-04-30 19:21:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #28 from h3rmit-git/master

Browse Source 
Changed the default download location of u2net models
This commit is contained in:
AUTOMATIC1111
2023-12-30 16:02:29 +03:00
committed by GitHub
parent 47165253bc 1c4a698cc8
commit 688564de22
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, ui_components
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",
@@ -55,6 +57,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),