From 17a99baf0c929e5df4dfc4b2a96aa3890a141112 Mon Sep 17 00:00:00 2001
From: Greendayle <Greendayle>
Date: Wed, 5 Oct 2022 22:05:24 +0200
Subject: [PATCH] better model search

---
 modules/deepbooru.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/modules/deepbooru.py b/modules/deepbooru.py
index 841cb9c5f..a64fd9cd1 100644
--- a/modules/deepbooru.py
+++ b/modules/deepbooru.py
@@ -9,8 +9,15 @@ import tensorflow as tf
 def _load_tf_and_return_tags(pil_image, threshold):
     this_folder = os.path.dirname(__file__)
     model_path = os.path.join(this_folder, '..', 'models', 'deepbooru', 'deepdanbooru-v3-20211112-sgd-e28')
-    if not os.path.exists(model_path):
-        return "Download https://github.com/KichangKim/DeepDanbooru/releases/download/v3-20211112-sgd-e28/deepdanbooru-v3-20211112-sgd-e28.zip unpack and put into models/deepbooru"
+
+    model_good = False
+    for path_candidate in [model_path, os.path.dirname(model_path)]:
+        if os.path.exists(os.path.join(path_candidate, 'project.json')):
+            model_path = path_candidate
+            model_good = True
+    if not model_good:
+        return ("Download https://github.com/KichangKim/DeepDanbooru/releases/download/v3-20211112-sgd-e28/"
+                "deepdanbooru-v3-20211112-sgd-e28.zip unpack and put into models/deepbooru")
 
     tags = dd.project.load_tags_from_project(model_path)
     model = dd.project.load_model_from_project(
-- 
GitLab