|
import os
|
|
import random
|
|
import datetime
|
|
import gradio as gr
|
|
from modules import script_callbacks, generation_parameters_copypaste, ui
|
|
|
|
|
|
base_dir = os.path.dirname(__file__)
|
|
|
|
data_dir = os.path.join(base_dir, 'Data')
|
|
|
|
def find_file_in_directory(filename, search_directory):
|
|
"""Recursively search for a file in a directory."""
|
|
for root, _, files in os.walk(search_directory):
|
|
if filename in files:
|
|
return os.path.join(root, filename)
|
|
return None
|
|
|
|
def read_tags_from_file(filename):
|
|
|
|
filepath = os.path.join(data_dir, filename)
|
|
if not os.path.exists(filepath):
|
|
|
|
pony_helper_dir = os.path.dirname(base_dir)
|
|
filepath = find_file_in_directory(filename, pony_helper_dir)
|
|
if not filepath:
|
|
print(f"Warning: No such file: '{filename}' in the PonyHelper directory")
|
|
return []
|
|
with open(filepath, 'r') as file:
|
|
return [tag.strip() for tag in file]
|
|
|
|
|
|
tags = {
|
|
'general': read_tags_from_file('General'),
|
|
'civitAI-Prompt': read_tags_from_file('civitAI-Prompt'),
|
|
'characters': read_tags_from_file('Character'),
|
|
'clothes': read_tags_from_file('clothes'),
|
|
'accessory': read_tags_from_file('Accessory'),
|
|
'copyright': read_tags_from_file('Copyright'),
|
|
'artist': read_tags_from_file('Artist'),
|
|
'nsfw': read_tags_from_file('NSFW'),
|
|
'meta': read_tags_from_file('meta'),
|
|
'hair': read_tags_from_file('hair'),
|
|
'Age': read_tags_from_file('Age'),
|
|
'Poses': read_tags_from_file('Poses'),
|
|
'Backgrounds': read_tags_from_file('Backgrounds'),
|
|
'Lighting': read_tags_from_file('Lighting'),
|
|
'Quality': read_tags_from_file('Quality'),
|
|
'Tattoos': read_tags_from_file('Tattoos'),
|
|
'Emotions_Expressions': read_tags_from_file('Emotions_Expressions'),
|
|
'Environment_Setting': read_tags_from_file('Environment_Setting'),
|
|
'Season': read_tags_from_file('Season'),
|
|
'Weather': read_tags_from_file('Weather'),
|
|
'Themes': read_tags_from_file('Themes'),
|
|
'Body_Types': read_tags_from_file('Body_Types'),
|
|
'clothing_materials': read_tags_from_file('clothing_materials'),
|
|
'Ethnicity_Nationality': read_tags_from_file('Ethnicity_Nationality'),
|
|
'Perspective': read_tags_from_file('Perspective'),
|
|
'Animals_Creatures': read_tags_from_file('Animals_Creatures'),
|
|
'Props_Objects': read_tags_from_file('Props_Objects'),
|
|
}
|
|
|
|
def generate_words(num_words, category):
|
|
return random.sample(tags[category], min(num_words, len(tags[category])))
|
|
|
|
def generate_mixed_tags(num_words, rating_safe, rating_explicit, rating_questionable, rating_realistic, undanbooru_tags, *enabled_categories):
|
|
pony_options = {
|
|
'rating_safe': rating_safe,
|
|
'rating_explicit': rating_explicit,
|
|
'rating_questionable': rating_questionable,
|
|
'rating_realistic': rating_realistic,
|
|
'undanbooru_tags': undanbooru_tags,
|
|
}
|
|
|
|
clean_tags = undanbooru_tags
|
|
categories = [category for category, enabled in zip(tags.keys(), enabled_categories) if enabled]
|
|
mixed_tags_result = []
|
|
|
|
for category in categories:
|
|
words = generate_words(num_words, category)
|
|
if clean_tags:
|
|
words = [word.replace('_', ' ').replace('(', ' ').replace(')', ' ').replace('{', ' ').replace('}', ' ').replace('-', ' ').replace('!?', ' ').replace('&', ' ').replace('+', ' ').replace('=', ' ').replace('@', ' ').replace('^', ' ').replace('<', ' ').replace('>', ' ') for word in words]
|
|
mixed_tags_result.extend(words)
|
|
|
|
if pony_options.get('Pony_score', False):
|
|
score_string1 = 'score_9, score_8_up, score_7_up, score_6_up, score_5_up, score_4_up'
|
|
score_string2 = 'score_10, score_9_up, score_8_up, score_7_up, score_9, score_8_up, score_7_up, score_6_up, score_5_up, score_4_up'
|
|
current_minute = datetime.datetime.now().minute
|
|
selected_score_string = score_string1 if current_minute % 2 == 0 else score_string2
|
|
mixed_tags_result.append(selected_score_string)
|
|
|
|
if pony_options.get('rating_safe', False):
|
|
mixed_tags_result.append('rating_safe')
|
|
if pony_options.get('rating_explicit', False):
|
|
mixed_tags_result.append('rating_explicit')
|
|
if pony_options.get('rating_questionable', False):
|
|
mixed_tags_result.append('rating_questionable')
|
|
if pony_options.get('rating_realistic', False):
|
|
mixed_tags_result.append('score_hyper-realistic, realistic, ((Full Body:1.2))')
|
|
|
|
for word, selected in pony_options.items():
|
|
if word.startswith('source_') and selected:
|
|
mixed_tags_result.append(word)
|
|
|
|
mixed_tags_str = ', '.join(mixed_tags_result)
|
|
return mixed_tags_str
|
|
|
|
def find_prompts(fields):
|
|
field_prompt = [x for x in fields if x[1] == "Prompt"][0]
|
|
field_negative_prompt = [x for x in fields if x[1] == "Negative prompt"][0]
|
|
return [field_prompt[0], field_negative_prompt[0]]
|
|
|
|
def send_prompts(text):
|
|
if not text:
|
|
return "", gr.update()
|
|
params = generation_parameters_copypaste.parse_generation_parameters(text)
|
|
negative_prompt = params.get("Negative prompt", "")
|
|
return params.get("Prompt", ""), negative_prompt or gr.update()
|
|
|
|
def add_tag_tab(enabled_categories):
|
|
with gr.Blocks(css=".slider { margin: 0.5em 0; } .checkbox { margin: 0.5em 0; } .btn { margin: 1em 0; padding: 0.5em 1em; background-color: #4CAF50; color: white; border: none; cursor: pointer; }") as interface:
|
|
with gr.Tab("Pony Helper"):
|
|
with gr.Row():
|
|
num_words = gr.Slider(label="Number of Tags __________________________Pony Scores- 'score_9, score_8_up, score_7_up, score_6_up, score_5_up, score_4_up' ", minimum=1, maximum=5, step=1, value=1)
|
|
with gr.Row():
|
|
rating_safe = gr.Checkbox(label="Safe", value=False)
|
|
rating_explicit = gr.Checkbox(label="Explicit", value=False)
|
|
rating_questionable = gr.Checkbox(label="Questionable", value=False)
|
|
rating_realistic = gr.Checkbox(label="Realistic", value=False)
|
|
undanbooru_tags = gr.Checkbox(label="Undanbooru Tags", value=False)
|
|
with gr.Row():
|
|
|
|
enabled_categories = [gr.Checkbox(label=category.replace('_', ' ').title(), value=enabled) for category, enabled in enabled_categories.items()]
|
|
with gr.Row():
|
|
btn_generate = gr.Button("Generate Mixed Tags", css="btn")
|
|
with gr.Row():
|
|
output = gr.Textbox(label="Generated Tags")
|
|
with gr.Row():
|
|
btn_send_txt2img = gr.Button("Send to txt2img", css="btn")
|
|
btn_send_img2img = gr.Button("Send to img2img", css="btn")
|
|
|
|
btn_generate.click(fn=generate_mixed_tags, inputs=[num_words, rating_safe, rating_explicit, rating_questionable, rating_realistic, undanbooru_tags] + enabled_categories, outputs=output)
|
|
btn_send_txt2img.click(fn=lambda tags: send_prompts(tags), inputs=output, outputs=find_prompts(ui.txt2img_paste_fields))
|
|
btn_send_img2img.click(fn=lambda tags: send_prompts(tags), inputs=output, outputs=find_prompts(ui.img2img_paste_fields))
|
|
|
|
return [(interface, "Pony Helper", "pony-helper")]
|
|
|
|
|
|
script_callbacks.on_ui_tabs(lambda: add_tag_tab(tags)) |