File size: 7,722 Bytes
c336648 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
import os
import random
import datetime
import gradio as gr
from modules import script_callbacks, generation_parameters_copypaste, ui
# Directory of the current script
base_dir = os.path.dirname(__file__)
# Directory where data files are expected to be stored
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):
# First, check in the data directory
filepath = os.path.join(data_dir, filename)
if not os.path.exists(filepath):
# If not found, search in the entire PonyHelper directory
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]
# Read tags from the text files
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():
# Use the provided enabled_categories argument
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")]
# Use the modified add_tag_tab function as the callback for ui_tabs
script_callbacks.on_ui_tabs(lambda: add_tag_tab(tags)) |