From 01f8cb44474e454903c11718e6a4f33dbde34bb8 Mon Sep 17 00:00:00 2001
From: Greendayle <Greendayle>
Date: Sat, 8 Oct 2022 18:02:56 +0200
Subject: made deepdanbooru optional, added to readme, automatic download of
 deepbooru model

---
 modules/deepbooru.py | 20 ++++++++++----------
 modules/shared.py    |  1 +
 modules/ui.py        | 19 ++++++++++++-------
 3 files changed, 23 insertions(+), 17 deletions(-)

(limited to 'modules')

diff --git a/modules/deepbooru.py b/modules/deepbooru.py
index 781b2249..7e3c0618 100644
--- a/modules/deepbooru.py
+++ b/modules/deepbooru.py
@@ -9,16 +9,16 @@ def _load_tf_and_return_tags(pil_image, threshold):
     import numpy as np
 
     this_folder = os.path.dirname(__file__)
-    model_path = os.path.join(this_folder, '..', 'models', 'deepbooru', 'deepdanbooru-v3-20211112-sgd-e28')
-
-    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")
+    model_path = os.path.abspath(os.path.join(this_folder, '..', 'models', 'deepbooru'))
+    if not os.path.exists(os.path.join(model_path, 'project.json')):
+        # there is no point importing these every time
+        import zipfile
+        from basicsr.utils.download_util import load_file_from_url
+        load_file_from_url(r"https://github.com/KichangKim/DeepDanbooru/releases/download/v3-20211112-sgd-e28/deepdanbooru-v3-20211112-sgd-e28.zip",
+                           model_path)
+        with zipfile.ZipFile(os.path.join(model_path, "deepdanbooru-v3-20211112-sgd-e28.zip"), "r") as zip_ref:
+            zip_ref.extractall(model_path)
+        os.remove(os.path.join(model_path, "deepdanbooru-v3-20211112-sgd-e28.zip"))
 
     tags = dd.project.load_tags_from_project(model_path)
     model = dd.project.load_model_from_project(
diff --git a/modules/shared.py b/modules/shared.py
index 02cb2722..c87b726e 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -44,6 +44,7 @@ parser.add_argument("--scunet-models-path", type=str, help="Path to directory wi
 parser.add_argument("--swinir-models-path", type=str, help="Path to directory with SwinIR model file(s).", default=os.path.join(models_path, 'SwinIR'))
 parser.add_argument("--ldsr-models-path", type=str, help="Path to directory with LDSR model file(s).", default=os.path.join(models_path, 'LDSR'))
 parser.add_argument("--xformers", action='store_true', help="enable xformers for cross attention layers")
+parser.add_argument("--deepdanbooru", action='store_true', help="enable deepdanbooru interrogator")
 parser.add_argument("--opt-split-attention", action='store_true', help="force-enables cross-attention layer optimization. By default, it's on for torch.cuda and off for other torch devices.")
 parser.add_argument("--disable-opt-split-attention", action='store_true', help="force-disables cross-attention layer optimization")
 parser.add_argument("--opt-split-attention-v1", action='store_true', help="enable older version of split attention optimization that does not consume all the VRAM it can find")
diff --git a/modules/ui.py b/modules/ui.py
index 30583fe9..c5c11c3c 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -23,9 +23,10 @@ import gradio.utils
 import gradio.routes
 
 from modules import sd_hijack
-from modules.deepbooru import get_deepbooru_tags
 from modules.paths import script_path
 from modules.shared import opts, cmd_opts
+if cmd_opts.deepdanbooru:
+    from modules.deepbooru import get_deepbooru_tags
 import modules.shared as shared
 from modules.sd_samplers import samplers, samplers_for_img2img
 from modules.sd_hijack import model_hijack
@@ -437,7 +438,10 @@ def create_toprow(is_img2img):
             with gr.Row(scale=1):
                 if is_img2img:
                     interrogate = gr.Button('Interrogate\nCLIP', elem_id="interrogate")
-                    deepbooru = gr.Button('Interrogate\nDeepBooru', elem_id="deepbooru")
+                    if cmd_opts.deepdanbooru:
+                        deepbooru = gr.Button('Interrogate\nDeepBooru', elem_id="deepbooru")
+                    else:
+                        deepbooru = None
                 else:
                     interrogate = None
                     deepbooru = None
@@ -782,11 +786,12 @@ def create_ui(wrap_gradio_gpu_call):
                 outputs=[img2img_prompt],
             )
 
-            img2img_deepbooru.click(
-                fn=interrogate_deepbooru,
-                inputs=[init_img],
-                outputs=[img2img_prompt],
-            )
+            if cmd_opts.deepdanbooru:
+                img2img_deepbooru.click(
+                    fn=interrogate_deepbooru,
+                    inputs=[init_img],
+                    outputs=[img2img_prompt],
+                )
 
             save.click(
                 fn=wrap_gradio_call(save_files),
-- 
cgit v1.2.1