mirror of
https://github.com/AUTOMATIC1111/stable-diffusion-webui.git
synced 2025-08-08 13:19:54 +00:00
refactored the deepbooru module to improve speed on running multiple interogations in a row. Added the option to generate deepbooru tags for textual inversion preproccessing.
This commit is contained in:
@@ -3,11 +3,14 @@ from PIL import Image, ImageOps
|
||||
import platform
|
||||
import sys
|
||||
import tqdm
|
||||
import time
|
||||
|
||||
from modules import shared, images
|
||||
from modules.shared import opts, cmd_opts
|
||||
if cmd_opts.deepdanbooru:
|
||||
import modules.deepbooru as deepbooru
|
||||
|
||||
|
||||
def preprocess(process_src, process_dst, process_flip, process_split, process_caption):
|
||||
def preprocess(process_src, process_dst, process_flip, process_split, process_caption, process_caption_deepbooru=False):
|
||||
size = 512
|
||||
src = os.path.abspath(process_src)
|
||||
dst = os.path.abspath(process_dst)
|
||||
@@ -24,10 +27,21 @@ def preprocess(process_src, process_dst, process_flip, process_split, process_ca
|
||||
if process_caption:
|
||||
shared.interrogator.load()
|
||||
|
||||
if process_caption_deepbooru:
|
||||
deepbooru.create_deepbooru_process()
|
||||
|
||||
def save_pic_with_caption(image, index):
|
||||
if process_caption:
|
||||
caption = "-" + shared.interrogator.generate_caption(image)
|
||||
caption = sanitize_caption(os.path.join(dst, f"{index:05}-{subindex[0]}"), caption, ".png")
|
||||
elif process_caption_deepbooru:
|
||||
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"]
|
||||
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]
|
||||
@@ -79,6 +93,10 @@ def preprocess(process_src, process_dst, process_flip, process_split, process_ca
|
||||
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"):
|
||||
|
Reference in New Issue
Block a user