Mirror/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2024-06-07 21:20:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

added deepbooru settings (threshold and sort by alpha or likelyhood)

Browse Source
This commit is contained in:
JC_Array 2022-10-10 18:01:49 -05:00
parent b980e7188c
commit 76ef3d75f6
2 changed files with 31 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
modules/deepbooru.py
View File

@ -3,31 +3,32 @@ from concurrent.futures import ProcessPoolExecutor
import multiprocessing
import time
def get_deepbooru_tags(pil_image, threshold=0.5):
def get_deepbooru_tags(pil_image):
"""
This method is for running only one image at a time for simple use. Used to the img2img interrogate.
"""
from modules import shared # prevents circular reference
create_deepbooru_process(threshold)
create_deepbooru_process(shared.opts.deepbooru_threshold, shared.opts.deepbooru_sort_alpha)
shared.deepbooru_process_return["value"] = -1
shared.deepbooru_process_queue.put(pil_image)
while shared.deepbooru_process_return["value"] == -1:
time.sleep(0.2)
tags = shared.deepbooru_process_return["value"]
release_process()
return tags
def deepbooru_process(queue, deepbooru_process_return, threshold):
def deepbooru_process(queue, deepbooru_process_return, threshold, alpha_sort):
model, tags = get_deepbooru_tags_model()
while True: # while process is running, keep monitoring queue for new image
pil_image = queue.get()
if pil_image == "QUIT":
break
else:
deepbooru_process_return["value"] = get_deepbooru_tags_from_model(model, tags, pil_image, threshold)
deepbooru_process_return["value"] = get_deepbooru_tags_from_model(model, tags, pil_image, threshold, alpha_sort)
def create_deepbooru_process(threshold=0.5):
def create_deepbooru_process(threshold, alpha_sort):
"""
Creates deepbooru process. A queue is created to send images into the process. This enables multiple images
to be processed in a row without reloading the model or creating a new process. To return the data, a shared
@ -40,7 +41,7 @@ def create_deepbooru_process(threshold=0.5):
shared.deepbooru_process_queue = shared.deepbooru_process_manager.Queue()
shared.deepbooru_process_return = shared.deepbooru_process_manager.dict()
shared.deepbooru_process_return["value"] = -1
shared.deepbooru_process = multiprocessing.Process(target=deepbooru_process, args=(shared.deepbooru_process_queue, shared.deepbooru_process_return, threshold))
shared.deepbooru_process = multiprocessing.Process(target=deepbooru_process, args=(shared.deepbooru_process_queue, shared.deepbooru_process_return, threshold, alpha_sort))
shared.deepbooru_process.start()
@ -80,7 +81,7 @@ def get_deepbooru_tags_model():
return model, tags
def get_deepbooru_tags_from_model(model, tags, pil_image, threshold=0.5):
def get_deepbooru_tags_from_model(model, tags, pil_image, threshold, alpha_sort):
import deepdanbooru as dd
import tensorflow as tf
import numpy as np
@ -105,15 +106,28 @@ def get_deepbooru_tags_from_model(model, tags, pil_image, threshold=0.5):
for i, tag in enumerate(tags):
result_dict[tag] = y[i]
result_tags_out = []
unsorted_tags_in_theshold = []
result_tags_print = []
for tag in tags:
if result_dict[tag] >= threshold:
if tag.startswith("rating:"):
continue
result_tags_out.append(tag)
unsorted_tags_in_theshold.append((result_dict[tag], tag))
result_tags_print.append(f'{result_dict[tag]} {tag}')
# sort tags
result_tags_out = []
sort_ndx = 0
print(alpha_sort)
if alpha_sort:
sort_ndx = 1
# sort by reverse by likelihood and normal for alpha
unsorted_tags_in_theshold.sort(key=lambda y: y[sort_ndx], reverse=(not alpha_sort))
for weight, tag in unsorted_tags_in_theshold:
result_tags_out.append(tag)
print('\n'.join(sorted(result_tags_print, reverse=True)))
return ', '.join(result_tags_out).replace('_', ' ').replace(':', ' ')

6
modules/shared.py
View File

@ -261,6 +261,12 @@ options_templates.update(options_section(('sampler-params', "Sampler parameters"
's_noise': OptionInfo(1.0, "sigma noise", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
}))
if cmd_opts.deepdanbooru:
options_templates.update(options_section(('deepbooru-params', "DeepBooru parameters"), {
"deepbooru_sort_alpha": OptionInfo(True, "Sort Alphabetical", gr.Checkbox),
'deepbooru_threshold': OptionInfo(0.5, "Threshold", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
}))
class Options:
data = None