File size: 4,972 Bytes
29a5ed9 | 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 150 151 152 153 154 155 | import os
import os.path
import re
import gradio as gr
from modules import script_callbacks, shared
from modules.ui import create_refresh_button
def add_tab():
with gr.Blocks(analytics_enabled=False) as ui:
with gr.Row().style(equal_height=False):
with gr.Column(variant='panel'):
gr.HTML(
value="<p>Differentiate prompts according to different types.")
replace_underscore = gr.Checkbox(
value=lambda: True,
label="Replace underscores with spaces",
elem_id="replace_underscore_checkbox")
need_split_prompts = gr.Textbox(
lines=6,
label="Prompts",
elem_id="need_split_prompts")
button_split_prompts = gr.Button(
elem_id="button_split_prompts",
value="Split",
variant='primary')
with gr.Column(variant='panel'):
submit_result = gr.Textbox(
label="Result",
elem_id="split_result",
interactive=False,
show_copy_button=True)
button_split_prompts.click(
fn=do_split,
inputs=[
need_split_prompts,
replace_underscore,
],
outputs=[submit_result]
)
return [(ui, "Split prompts", "split_prompts")]
def load_txt_to_dictionary(filePath, main_folder, file_contents):
with open(os.path.join(filePath), "r", encoding="utf-8") as f:
# Add the contents of the txt file to the main folder's list
file_contents[main_folder].extend(
[word.lower() for word in f.read().splitlines()])
def load_classification_files():
base_path = os.path.join(os.path.dirname(
os.path.dirname(os.path.realpath(__file__))), "分類")
# A dictionary to hold the contents of the txt files classified by the main folders
file_contents = {}
# Iterate through the main folders in the "分類" directory
for main_folder in os.listdir(base_path):
main_folder_path = os.path.join(base_path, main_folder)
# Only process directories (main folders)
if os.path.isdir(main_folder_path):
file_contents[main_folder] = []
# Iterate through all subdirectories and txt files inside the main folder
for dirpath, dirnames, filenames in os.walk(main_folder_path):
for file in filenames:
if file.endswith(".txt"):
load_txt_to_dictionary(
os.path.join(dirpath, file), main_folder, file_contents)
elif os.path.isfile(main_folder_path) and main_folder.endswith(".txt"):
file_contents[main_folder.replace(".txt", "")] = []
load_txt_to_dictionary(
main_folder_path, main_folder.replace(".txt", ""), file_contents)
return file_contents
def simplify_word(word):
# 使用正则表达式提取词部分
match = re.match(r'^\(([\w\s]+):\d+\.\d+\)$|^\(+([\w\s]+)\)+$', word)
if match:
processed_word = match.group(1) or match.group(2)
return processed_word
else:
return word
# Modify do_split function to use the classification files
def do_split(need_split_prompts, replace_underscore):
classifications = load_classification_files()
results = {key: [] for key in classifications.keys()}
if "其他" not in results:
results["其他"] = []
if "lora" not in results:
results["lora"] = []
prompts = re.split(r'[,|\n]', need_split_prompts)
for prompt in prompts:
prompt = prompt.strip().lower()
if not prompt:
continue
if prompt.startswith("<lora:"):
results["lora"].append(prompt)
continue
if replace_underscore:
prompt = prompt.replace("_", " ")
check_prompt = prompt
if prompt.startswith("("):
check_prompt = simplify_word(prompt)
classified = False
for file_name, keywords in classifications.items():
if any(keyword == check_prompt for keyword in keywords):
results[file_name].append(prompt)
classified = True
break
if not classified:
results["其他"].append(prompt)
# Format the results for output
seg_word = ", "
splited_result = ""
for file_name, prompts in results.items():
if len(prompts) == 0:
continue
if file_name.startswith("lora"):
seg_word = "\n"
else:
splited_result += f"[:{os.path.splitext(file_name)[0]}:99] "
splited_result += seg_word.join(prompts) + "\n\n"
return splited_result.rstrip("\n")
script_callbacks.on_ui_tabs(add_tab)
|