Fix prompt usage

pull/5/head^2
Alvi 2024-05-07 05:26:55 -03:00
parent cb9ab96ad3
commit 4e2dcb9d9e
1 changed files with 21 additions and 25 deletions

View File

@ -1,19 +1,18 @@
import gradio as gr import gradio as gr
import unicodedata
from modules import scripts, deepbooru from modules import scripts, deepbooru
from modules.processing import process_images from modules.processing import process_images
import modules.shared as shared import modules.shared as shared
def remove_accents(input_str): """
nfkd_form = unicodedata.normalize('NFKD', input_str)
return u"".join(c for c in nfkd_form if not unicodedata.combining(c)) Thanks to Mathias Russ.
Thanks to RookHyena.
"""
class Script(scripts.Script): class Script(scripts.Script):
# :|
original_prompt = "Some"
interrogator = "Random text"
def title(self): def title(self):
return "Img2img batch interrogator" return "Img2img batch interrogator"
@ -31,31 +30,28 @@ class Script(scripts.Script):
def run(self, p, in_front, prompt_weight, use_deepbooru): def run(self, p, in_front, prompt_weight, use_deepbooru):
# :| raw_prompt = p.prompt
if in_front:
_check = f"{Script.original_prompt}, ({Script.interrogator}:{prompt_weight})"
else:
_check = f"({Script.interrogator}:{prompt_weight}), {Script.original_prompt}"
if p.prompt not in [_check, Script.interrogator]:
Script.original_prompt = remove_accents(p.prompt)
# fix alpha channel # fix alpha channel
p.init_images[0] = p.init_images[0].convert('RGB') p.init_images[0] = p.init_images[0].convert("RGB")
if use_deepbooru: if use_deepbooru:
prompt = deepbooru.model.tag(p.init_images[0]) interrogator = deepbooru.model.tag(p.init_images[0])
else: else:
prompt = shared.interrogator.interrogate(p.init_images[0]) interrogator = shared.interrogator.interrogate(p.init_images[0])
Script.interrogator = prompt
p.prompt = "" if p.prompt == "":
if Script.original_prompt in ["Some", ""]: p.prompt = interrogator
p.prompt = Script.interrogator
elif in_front: elif in_front:
p.prompt = f"{Script.original_prompt}, ({Script.interrogator}:{prompt_weight})" p.prompt = f"{p.prompt}, ({interrogator}:{prompt_weight})"
else: else:
p.prompt = f"({Script.interrogator}:{prompt_weight}), {Script.original_prompt}" p.prompt = f"({interrogator}:{prompt_weight}), {p.prompt}"
print(f"Prompt: {p.prompt}") print(f"Prompt: {p.prompt}")
return process_images(p)
processed = process_images(p)
# Restore the UI elements we modified
p.prompt = raw_prompt
return processed