From bed6ab1df19eaf4bdf28501525958b0ef5d87517 Mon Sep 17 00:00:00 2001 From: yolain Date: Sat, 30 Nov 2024 16:49:50 +0800 Subject: [PATCH] Fix download judgement no longer requested when rmbg-2.0 model already exists --- py/image.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/py/image.py b/py/image.py index 5ba74fc..5625917 100644 --- a/py/image.py +++ b/py/image.py @@ -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")