diff --git a/README.MD b/README.MD index 7637075..e6d659d 100644 --- a/README.MD +++ b/README.MD @@ -667,7 +667,7 @@ Remove background. compared to the similar background removal nodes, this node h This node combines the Alpha Matte node of Spacepxl's [ComfyUI-Image-Filters](https://github.com/spacepxl/ComfyUI-Image-Filters) and the functionality of ZHO-ZHO-ZHO's [ComfyUI-BRIA_AI-RMBG](https://github.com/ZHO-ZHO-ZHO/ComfyUI-BRIA_AI-RMBG). -*Download the [BRIA Background Removal v1.4](https://huggingface.co/briaai/RMBG-1.4) model file (model. pth) to the customer_modes/ComfyUI-LayerStyle/RMBG-1.4 floder. +*Download the [BRIA Background Removal v1.4](https://huggingface.co/briaai/RMBG-1.4) model file (model. pth) to the /ComfyUI/models/RMBG-1.4 folder. This model was developed by BRIA AI and can be used as an open-source model for non-commercial purposes. ![image](image/rembg_ultra_example.png) diff --git a/README_CN.MD b/README_CN.MD index bcb47fa..8e0b442 100644 --- a/README_CN.MD +++ b/README_CN.MD @@ -653,7 +653,7 @@ ### RemBgUltra 去除背景。与类似的背景移除节点相比,这个节点具有超高的边缘细节。 本节点结合了spacepxl的[ComfyUI-Image-Filters](https://github.com/spacepxl/ComfyUI-Image-Filters)的Alpha Matte节点,以及ZHO-ZHO-ZHO的[ComfyUI-BRIA_AI-RMBG](https://github.com/ZHO-ZHO-ZHO/ComfyUI-BRIA_AI-RMBG)的功能。 -*将[BRIA Background Removal v1.4](https://huggingface.co/briaai/RMBG-1.4)模型文件(model.pth)下载至/custom_nodes/ComfyUI_LayerStyle/RMBG-1.4。 +*将[BRIA Background Removal v1.4](https://huggingface.co/briaai/RMBG-1.4)模型文件(model.pth)下载至/ComfyUI/models/RMBG-1.4文件夹。 该模型由 BRIA AI 开发,可作为非商业用途的开源模型。 ![image](image/rembg_ultra_example.png) diff --git a/RMBG-1.4/put model here.txt b/RMBG-1.4/put model here.txt deleted file mode 100644 index 8b13789..0000000 --- a/RMBG-1.4/put model here.txt +++ /dev/null @@ -1 +0,0 @@ - diff --git a/py/imagefunc.py b/py/imagefunc.py index 2ea9b80..6224dd7 100644 --- a/py/imagefunc.py +++ b/py/imagefunc.py @@ -22,6 +22,7 @@ import torch.nn.functional as F import colorsys from colour.io.luts.iridas_cube import read_LUT_IridasCube, LUT3D, LUT3x1D from typing import Union +from main import folder_paths from .briarmbg import BriaRMBG def log(message:str, message_type:str='info'): @@ -682,7 +683,10 @@ def load_RMBG_model(): current_directory = os.path.dirname(os.path.abspath(__file__)) device = "cuda" if torch.cuda.is_available() else "cpu" net = BriaRMBG() - model_path = os.path.join(os.path.dirname(current_directory), "RMBG-1.4/model.pth") + _path = os.path.join("RMBG-1.4", "model.pth") + model_path = os.path.join(folder_paths.models_dir, _path) + if not os.path.exists(model_path): + model_path = os.path.join(os.path.dirname(current_directory), _path) net.load_state_dict(torch.load(model_path, map_location=device)) net.to(device) net.eval()