mirror of
https://github.com/openvinotoolkit/stable-diffusion-webui.git
synced 2024-12-14 22:53:25 +03:00
deepbooru: added option to use spaces or underscores
deepbooru: added option to quote (\) in tags deepbooru/BLIP: write caption to file instead of image filename deepbooru/BLIP: now possible to use both for captions deepbooru: process is stopped even if an exception occurs
This commit is contained in:
parent
c3c8eef9fd
commit
698d303b04
@ -2,33 +2,44 @@ import os.path
|
|||||||
from concurrent.futures import ProcessPoolExecutor
|
from concurrent.futures import ProcessPoolExecutor
|
||||||
import multiprocessing
|
import multiprocessing
|
||||||
import time
|
import time
|
||||||
|
import re
|
||||||
|
|
||||||
|
re_special = re.compile(r'([\\()])')
|
||||||
|
|
||||||
def get_deepbooru_tags(pil_image):
|
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.
|
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
|
from modules import shared # prevents circular reference
|
||||||
create_deepbooru_process(shared.opts.interrogate_deepbooru_score_threshold, shared.opts.deepbooru_sort_alpha)
|
|
||||||
shared.deepbooru_process_return["value"] = -1
|
try:
|
||||||
shared.deepbooru_process_queue.put(pil_image)
|
create_deepbooru_process(shared.opts.interrogate_deepbooru_score_threshold, create_deepbooru_opts())
|
||||||
while shared.deepbooru_process_return["value"] == -1:
|
return get_tags_from_process(pil_image)
|
||||||
time.sleep(0.2)
|
finally:
|
||||||
tags = shared.deepbooru_process_return["value"]
|
|
||||||
release_process()
|
release_process()
|
||||||
return tags
|
|
||||||
|
|
||||||
|
|
||||||
def deepbooru_process(queue, deepbooru_process_return, threshold, alpha_sort):
|
def create_deepbooru_opts():
|
||||||
|
from modules import shared
|
||||||
|
|
||||||
|
return {
|
||||||
|
"use_spaces": shared.opts.deepbooru_use_spaces,
|
||||||
|
"use_escape": shared.opts.deepbooru_escape,
|
||||||
|
"alpha_sort": shared.opts.deepbooru_sort_alpha,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def deepbooru_process(queue, deepbooru_process_return, threshold, deepbooru_opts):
|
||||||
model, tags = get_deepbooru_tags_model()
|
model, tags = get_deepbooru_tags_model()
|
||||||
while True: # while process is running, keep monitoring queue for new image
|
while True: # while process is running, keep monitoring queue for new image
|
||||||
pil_image = queue.get()
|
pil_image = queue.get()
|
||||||
if pil_image == "QUIT":
|
if pil_image == "QUIT":
|
||||||
break
|
break
|
||||||
else:
|
else:
|
||||||
deepbooru_process_return["value"] = get_deepbooru_tags_from_model(model, tags, pil_image, threshold, alpha_sort)
|
deepbooru_process_return["value"] = get_deepbooru_tags_from_model(model, tags, pil_image, threshold, deepbooru_opts)
|
||||||
|
|
||||||
|
|
||||||
def create_deepbooru_process(threshold, alpha_sort):
|
def create_deepbooru_process(threshold, deepbooru_opts):
|
||||||
"""
|
"""
|
||||||
Creates deepbooru process. A queue is created to send images into the process. This enables multiple images
|
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
|
to be processed in a row without reloading the model or creating a new process. To return the data, a shared
|
||||||
@ -41,10 +52,23 @@ def create_deepbooru_process(threshold, alpha_sort):
|
|||||||
shared.deepbooru_process_queue = shared.deepbooru_process_manager.Queue()
|
shared.deepbooru_process_queue = shared.deepbooru_process_manager.Queue()
|
||||||
shared.deepbooru_process_return = shared.deepbooru_process_manager.dict()
|
shared.deepbooru_process_return = shared.deepbooru_process_manager.dict()
|
||||||
shared.deepbooru_process_return["value"] = -1
|
shared.deepbooru_process_return["value"] = -1
|
||||||
shared.deepbooru_process = multiprocessing.Process(target=deepbooru_process, args=(shared.deepbooru_process_queue, shared.deepbooru_process_return, threshold, alpha_sort))
|
shared.deepbooru_process = multiprocessing.Process(target=deepbooru_process, args=(shared.deepbooru_process_queue, shared.deepbooru_process_return, threshold, deepbooru_opts))
|
||||||
shared.deepbooru_process.start()
|
shared.deepbooru_process.start()
|
||||||
|
|
||||||
|
|
||||||
|
def get_tags_from_process(image):
|
||||||
|
from modules import shared
|
||||||
|
|
||||||
|
shared.deepbooru_process_return["value"] = -1
|
||||||
|
shared.deepbooru_process_queue.put(image)
|
||||||
|
while shared.deepbooru_process_return["value"] == -1:
|
||||||
|
time.sleep(0.2)
|
||||||
|
caption = shared.deepbooru_process_return["value"]
|
||||||
|
shared.deepbooru_process_return["value"] = -1
|
||||||
|
|
||||||
|
return caption
|
||||||
|
|
||||||
|
|
||||||
def release_process():
|
def release_process():
|
||||||
"""
|
"""
|
||||||
Stops the deepbooru process to return used memory
|
Stops the deepbooru process to return used memory
|
||||||
@ -81,10 +105,15 @@ def get_deepbooru_tags_model():
|
|||||||
return model, tags
|
return model, tags
|
||||||
|
|
||||||
|
|
||||||
def get_deepbooru_tags_from_model(model, tags, pil_image, threshold, alpha_sort):
|
def get_deepbooru_tags_from_model(model, tags, pil_image, threshold, deepbooru_opts):
|
||||||
import deepdanbooru as dd
|
import deepdanbooru as dd
|
||||||
import tensorflow as tf
|
import tensorflow as tf
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
|
alpha_sort = deepbooru_opts['alpha_sort']
|
||||||
|
use_spaces = deepbooru_opts['use_spaces']
|
||||||
|
use_escape = deepbooru_opts['use_escape']
|
||||||
|
|
||||||
width = model.input_shape[2]
|
width = model.input_shape[2]
|
||||||
height = model.input_shape[1]
|
height = model.input_shape[1]
|
||||||
image = np.array(pil_image)
|
image = np.array(pil_image)
|
||||||
@ -129,4 +158,12 @@ def get_deepbooru_tags_from_model(model, tags, pil_image, threshold, alpha_sort)
|
|||||||
|
|
||||||
print('\n'.join(sorted(result_tags_print, reverse=True)))
|
print('\n'.join(sorted(result_tags_print, reverse=True)))
|
||||||
|
|
||||||
return ', '.join(result_tags_out).replace('_', ' ').replace(':', ' ')
|
tags_text = ', '.join(result_tags_out)
|
||||||
|
|
||||||
|
if use_spaces:
|
||||||
|
tags_text = tags_text.replace('_', ' ')
|
||||||
|
|
||||||
|
if use_escape:
|
||||||
|
tags_text = re.sub(re_special, r'\\\1', tags_text)
|
||||||
|
|
||||||
|
return tags_text.replace(':', ' ')
|
||||||
|
@ -260,6 +260,8 @@ options_templates.update(options_section(('interrogate', "Interrogate Options"),
|
|||||||
"interrogate_clip_max_length": OptionInfo(48, "Interrogate: maximum description length", gr.Slider, {"minimum": 1, "maximum": 256, "step": 1}),
|
"interrogate_clip_max_length": OptionInfo(48, "Interrogate: maximum description length", gr.Slider, {"minimum": 1, "maximum": 256, "step": 1}),
|
||||||
"interrogate_deepbooru_score_threshold": OptionInfo(0.5, "Interrogate: deepbooru score threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}),
|
"interrogate_deepbooru_score_threshold": OptionInfo(0.5, "Interrogate: deepbooru score threshold", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01}),
|
||||||
"deepbooru_sort_alpha": OptionInfo(True, "Interrogate: deepbooru sort alphabetically"),
|
"deepbooru_sort_alpha": OptionInfo(True, "Interrogate: deepbooru sort alphabetically"),
|
||||||
|
"deepbooru_use_spaces": OptionInfo(False, "use spaces for tags in deepbooru"),
|
||||||
|
"deepbooru_escape": OptionInfo(True, "escape (\\) brackets in deepbooru (so they are used as literal brackets and not for emphasis)"),
|
||||||
}))
|
}))
|
||||||
|
|
||||||
options_templates.update(options_section(('ui', "User interface"), {
|
options_templates.update(options_section(('ui', "User interface"), {
|
||||||
|
@ -10,7 +10,28 @@ from modules.shared import opts, cmd_opts
|
|||||||
if cmd_opts.deepdanbooru:
|
if cmd_opts.deepdanbooru:
|
||||||
import modules.deepbooru as deepbooru
|
import modules.deepbooru as deepbooru
|
||||||
|
|
||||||
|
|
||||||
def preprocess(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru=False):
|
def preprocess(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru=False):
|
||||||
|
try:
|
||||||
|
if process_caption:
|
||||||
|
shared.interrogator.load()
|
||||||
|
|
||||||
|
if process_caption_deepbooru:
|
||||||
|
deepbooru.create_deepbooru_process(opts.interrogate_deepbooru_score_threshold, deepbooru.create_deepbooru_opts())
|
||||||
|
|
||||||
|
preprocess_work(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru)
|
||||||
|
|
||||||
|
finally:
|
||||||
|
|
||||||
|
if process_caption:
|
||||||
|
shared.interrogator.send_blip_to_ram()
|
||||||
|
|
||||||
|
if process_caption_deepbooru:
|
||||||
|
deepbooru.release_process()
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def preprocess_work(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru=False):
|
||||||
width = process_width
|
width = process_width
|
||||||
height = process_height
|
height = process_height
|
||||||
src = os.path.abspath(process_src)
|
src = os.path.abspath(process_src)
|
||||||
@ -25,30 +46,28 @@ def preprocess(process_src, process_dst, process_width, process_height, process_
|
|||||||
shared.state.textinfo = "Preprocessing..."
|
shared.state.textinfo = "Preprocessing..."
|
||||||
shared.state.job_count = len(files)
|
shared.state.job_count = len(files)
|
||||||
|
|
||||||
|
def save_pic_with_caption(image, index):
|
||||||
|
caption = ""
|
||||||
|
|
||||||
if process_caption:
|
if process_caption:
|
||||||
shared.interrogator.load()
|
caption += shared.interrogator.generate_caption(image)
|
||||||
|
|
||||||
if process_caption_deepbooru:
|
if process_caption_deepbooru:
|
||||||
deepbooru.create_deepbooru_process(opts.interrogate_deepbooru_score_threshold, opts.deepbooru_sort_alpha)
|
if len(caption) > 0:
|
||||||
|
caption += ", "
|
||||||
|
caption += deepbooru.get_tags_from_process(image)
|
||||||
|
|
||||||
def save_pic_with_caption(image, index):
|
filename_part = filename
|
||||||
if process_caption:
|
filename_part = os.path.splitext(filename_part)[0]
|
||||||
caption = "-" + shared.interrogator.generate_caption(image)
|
filename_part = os.path.basename(filename_part)
|
||||||
caption = sanitize_caption(os.path.join(dst, f"{index:05}-{subindex[0]}"), caption, ".png")
|
|
||||||
elif process_caption_deepbooru:
|
basename = f"{index:05}-{subindex[0]}-{filename_part}"
|
||||||
shared.deepbooru_process_return["value"] = -1
|
image.save(os.path.join(dst, f"{basename}.png"))
|
||||||
shared.deepbooru_process_queue.put(image)
|
|
||||||
while shared.deepbooru_process_return["value"] == -1:
|
if len(caption) > 0:
|
||||||
time.sleep(0.2)
|
with open(os.path.join(dst, f"{basename}.txt"), "w", encoding="utf8") as file:
|
||||||
caption = "-" + shared.deepbooru_process_return["value"]
|
file.write(caption)
|
||||||
caption = sanitize_caption(os.path.join(dst, f"{index:05}-{subindex[0]}"), caption, ".png")
|
|
||||||
shared.deepbooru_process_return["value"] = -1
|
|
||||||
else:
|
|
||||||
caption = filename
|
|
||||||
caption = os.path.splitext(caption)[0]
|
|
||||||
caption = os.path.basename(caption)
|
|
||||||
|
|
||||||
image.save(os.path.join(dst, f"{index:05}-{subindex[0]}{caption}.png"))
|
|
||||||
subindex[0] += 1
|
subindex[0] += 1
|
||||||
|
|
||||||
def save_pic(image, index):
|
def save_pic(image, index):
|
||||||
@ -93,34 +112,3 @@ def preprocess(process_src, process_dst, process_width, process_height, process_
|
|||||||
save_pic(img, index)
|
save_pic(img, index)
|
||||||
|
|
||||||
shared.state.nextjob()
|
shared.state.nextjob()
|
||||||
|
|
||||||
if process_caption:
|
|
||||||
shared.interrogator.send_blip_to_ram()
|
|
||||||
|
|
||||||
if process_caption_deepbooru:
|
|
||||||
deepbooru.release_process()
|
|
||||||
|
|
||||||
|
|
||||||
def sanitize_caption(base_path, original_caption, suffix):
|
|
||||||
operating_system = platform.system().lower()
|
|
||||||
if (operating_system == "windows"):
|
|
||||||
invalid_path_characters = "\\/:*?\"<>|"
|
|
||||||
max_path_length = 259
|
|
||||||
else:
|
|
||||||
invalid_path_characters = "/" #linux/macos
|
|
||||||
max_path_length = 1023
|
|
||||||
caption = original_caption
|
|
||||||
for invalid_character in invalid_path_characters:
|
|
||||||
caption = caption.replace(invalid_character, "")
|
|
||||||
fixed_path_length = len(base_path) + len(suffix)
|
|
||||||
if fixed_path_length + len(caption) <= max_path_length:
|
|
||||||
return caption
|
|
||||||
caption_tokens = caption.split()
|
|
||||||
new_caption = ""
|
|
||||||
for token in caption_tokens:
|
|
||||||
last_caption = new_caption
|
|
||||||
new_caption = new_caption + token + " "
|
|
||||||
if (len(new_caption) + fixed_path_length - 1 > max_path_length):
|
|
||||||
break
|
|
||||||
print(f"\nPath will be too long. Truncated caption: {original_caption}\nto: {last_caption}", file=sys.stderr)
|
|
||||||
return last_caption.strip()
|
|
||||||
|
@ -1074,11 +1074,8 @@ def create_ui(wrap_gradio_gpu_call):
|
|||||||
with gr.Row():
|
with gr.Row():
|
||||||
process_flip = gr.Checkbox(label='Create flipped copies')
|
process_flip = gr.Checkbox(label='Create flipped copies')
|
||||||
process_split = gr.Checkbox(label='Split oversized images into two')
|
process_split = gr.Checkbox(label='Split oversized images into two')
|
||||||
process_caption = gr.Checkbox(label='Use BLIP caption as filename')
|
process_caption = gr.Checkbox(label='Use BLIP for caption')
|
||||||
if cmd_opts.deepdanbooru:
|
process_caption_deepbooru = gr.Checkbox(label='Use deepbooru for caption', visible=True if cmd_opts.deepdanbooru else False)
|
||||||
process_caption_deepbooru = gr.Checkbox(label='Use deepbooru caption as filename')
|
|
||||||
else:
|
|
||||||
process_caption_deepbooru = gr.Checkbox(label='Use deepbooru caption as filename', visible=False)
|
|
||||||
|
|
||||||
with gr.Row():
|
with gr.Row():
|
||||||
with gr.Column(scale=3):
|
with gr.Column(scale=3):
|
||||||
|
Loading…
Reference in New Issue
Block a user