Use shared models path

pull/55/head
d8ahazard 2023-11-16 13:19:11 -06:00
parent 91e16d32b1
commit 7741d8c08a
1 changed files with 7 additions and 2 deletions

View File

@ -21,8 +21,13 @@ from collections import OrderedDict
from PIL import Image
model_cache = OrderedDict()
sam_model_dir = os.path.join(
extensions_dir, "PBRemTools/models/")
models_path = shared.models_path
sams_dir = os.path.join(models_path, "sam")
if os.path.exists(sams_dir):
sam_model_dir = sams_dir
else:
sam_model_dir = os.path.join(
extensions_dir, "PBRemTools/models/")
model_list = [f for f in os.listdir(sam_model_dir) if os.path.isfile(
os.path.join(sam_model_dir, f)) and f.split('.')[-1] != 'txt']