nodes/yolain-ComfyUI-Easy-Use
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix download judgement no longer requested when rmbg-2.0 model already exists

Browse Source
This commit is contained in:
yolain
2024-11-30 16:49:50 +08:00
parent ff8ba6b209
commit bed6ab1df1
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/image.py
+3 -2
View File
@@ -822,9 +822,10 @@ class imageRemBg:
if rem_mode == "RMBG-2.0":
repo_id = REMBG_MODELS[rem_mode]['model_url']
model_path = os.path.join(REMBG_DIR, 'RMBG-2.0')
from huggingface_hub import snapshot_download
if not os.path.exists(model_path):
from huggingface_hub import snapshot_download
snapshot_download(repo_id=repo_id, local_dir=model_path, ignore_patterns=["*.md", "*.txt"])
from transformers import AutoModelForImageSegmentation
snapshot_download(repo_id=repo_id, local_dir=model_path, ignore_patterns=["*.md", "*.txt"])
model = AutoModelForImageSegmentation.from_pretrained(model_path, trust_remote_code=True)
torch.set_float32_matmul_precision('high')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")