1580158
/
stable-diffusion-webui-dataset-tag-editor
mirror of https://github.com/toshiaki1729/stable-diffusion-webui-dataset-tag-editor.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Remove cmd args requirement for deepdanbooru

pull/10/head
dtlnor 2022-11-22 13:01:06 +09:00
parent 04a41d0234
commit e13cfe1aeb
2 changed files with 11 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
scripts/dataset_tag_editor/dataset_tag_editor.py
View File

@ -6,9 +6,7 @@ from modules.shared import opts, cmd_opts
from modules.textual_inversion.dataset import re_numbers_at_start
from PIL import Image
from enum import Enum
if cmd_opts.deepdanbooru:
import modules.deepbooru as deepbooru
import modules.deepbooru as deepbooru
re_tags = re.compile(r'^(.+) \[\d+\]$')
@ -193,9 +191,6 @@ class DatasetTagEditor:
self.dataset_dir = img_dir
if use_booru and not cmd_opts.deepdanbooru:
print('Cannot use DeepDanbooru without --deepdanbooru commandline option.')
print(f'Total {len(filepath_set)} files under the directory including not image files.')
def load_images(filepath_set: Set[str]):
@ -237,7 +232,7 @@ class DatasetTagEditor:
if use_clip:
interrogate_text += shared.interrogator.generate_caption(img)
if use_booru and cmd_opts.deepdanbooru:
if use_booru:
tmp = deepbooru.get_tags_from_process(img)
interrogate_text += (', ' if interrogate_text and tmp else '') + tmp
@ -255,7 +250,7 @@ class DatasetTagEditor:
if interrogate_method != InterrogateMethod.NONE:
if use_clip:
shared.interrogator.load()
if use_booru and cmd_opts.deepdanbooru:
if use_booru:
db_opts = deepbooru.create_deepbooru_opts()
db_opts[deepbooru.OPT_INCLUDE_RANKS] = False
deepbooru.create_deepbooru_process(opts.interrogate_deepbooru_score_threshold, db_opts)
@ -266,7 +261,7 @@ class DatasetTagEditor:
if interrogate_method != InterrogateMethod.NONE:
if use_clip:
shared.interrogator.send_blip_to_ram()
if use_booru and cmd_opts.deepdanbooru:
if use_booru:
deepbooru.release_process()
self.construct_tag_counts()

19
scripts/main.py
View File

@ -290,10 +290,7 @@ def on_ui_tabs():
rb_use_interrogator = gr.Radio(choices=['No', 'If Empty', 'Overwrite', 'Prepend', 'Append'], value='No', label='Use Interrogator Caption')
with gr.Row():
cb_use_clip_to_prefill = gr.Checkbox(value=False, label='Use BLIP')
if cmd_opts.deepdanbooru:
cb_use_booru_to_prefill = gr.Checkbox(value=False, label='Use DeepDanbooru')
else:
cb_use_booru_to_prefill = gr.Checkbox(value=False, label='Use DeepDanbooru', interactive=False)
cb_use_booru_to_prefill = gr.Checkbox(value=False, label='Use DeepDanbooru')
gl_dataset_images = gr.Gallery(label='Dataset Images', elem_id="dataset_tag_editor_dataset_gallery").style(grid=opts.dataset_editor_image_columns)
txt_filter = gr.HTML(value=get_current_txt_filter())
@ -366,8 +363,7 @@ def on_ui_tabs():
with gr.Tab(label='Interrogate Selected Image'):
with gr.Row():
btn_interrogate_clip = gr.Button(value='Interrogate with BLIP')
if cmd_opts.deepdanbooru:
btn_interrogate_booru = gr.Button(value='Interrogate with DeepDanbooru')
btn_interrogate_booru = gr.Button(value='Interrogate with DeepDanbooru')
tb_interrogate_selected_image = gr.Textbox(label='Interrogate Result', interactive=True, lines=6)
with gr.Row():
btn_copy_interrogate = gr.Button(value='Copy and Overwrite')
@ -522,12 +518,11 @@ def on_ui_tabs():
outputs=[tb_interrogate_selected_image]
)
if cmd_opts.deepdanbooru:
btn_interrogate_booru.click(
fn=interrogate_selected_image_booru,
inputs=None,
outputs=[tb_interrogate_selected_image]
)
btn_interrogate_booru.click(
fn=interrogate_selected_image_booru,
inputs=None,
outputs=[tb_interrogate_selected_image]
)
btn_copy_interrogate.click(
fn=lambda a:a,